Skip to content

Commit 9c5b418

Browse files
Update location_manager.py to use new async task pattern
- Remove @dojo_async_task decorator - Use dojo_dispatch_task for calling Celery tasks - Convert task methods to static-like functions (no self parameter) - Make get_or_create_location and clean_unsaved_locations static methods
1 parent 102ef06 commit 9c5b418

1 file changed

Lines changed: 15 additions & 20 deletions

File tree

dojo/importers/location_manager.py

Lines changed: 15 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
from django.utils import timezone
77

88
from dojo.celery import app
9-
from dojo.decorators import dojo_async_task
9+
from dojo.celery_dispatch import dojo_dispatch_task
1010
from dojo.location.models import AbstractLocation, LocationFindingReference
1111
from dojo.location.status import FindingLocationStatus
1212
from dojo.models import (
@@ -24,41 +24,38 @@
2424

2525
# test_notifications.py: Implement Locations
2626
class LocationManager:
27-
def get_or_create_location(self, unsaved_location: AbstractLocation) -> AbstractLocation | None:
27+
@staticmethod
28+
def get_or_create_location(unsaved_location: AbstractLocation) -> AbstractLocation | None:
2829
if isinstance(unsaved_location, URL):
2930
return URL.get_or_create_from_object(unsaved_location)
3031
logger.debug(f"IMPORT_SCAN: Unsupported location type: {type(unsaved_location)}")
3132
return None
3233

33-
@dojo_async_task
34-
@app.task()
34+
@app.task
3535
def add_locations_to_unsaved_finding(
36-
self,
37-
finding: Finding,
36+
finding: Finding, # noqa: N805
3837
locations: list[AbstractLocation],
3938
**kwargs: dict,
4039
) -> None:
4140
"""Creates Endpoint objects for a single finding and creates the link via the endpoint status"""
4241
locations = list(set(locations))
4342

4443
logger.debug(f"IMPORT_SCAN: Adding {len(locations)} locations to finding: {finding}")
45-
self.clean_unsaved_locations(locations)
44+
LocationManager.clean_unsaved_locations(locations)
4645

4746
# LOCATION LOCATION LOCATION
4847
# TODO: bulk create the finding/product refs...
4948
locations_saved = 0
5049
for unsaved_location in locations:
51-
if saved_location := self.get_or_create_location(unsaved_location):
50+
if saved_location := LocationManager.get_or_create_location(unsaved_location):
5251
locations_saved += 1
5352
saved_location.location.associate_with_finding(finding, status=FindingLocationStatus.Active)
5453

5554
logger.debug(f"IMPORT_SCAN: {locations_saved} locations imported")
5655

57-
@dojo_async_task
58-
@app.task()
56+
@app.task
5957
def mitigate_location_status(
60-
self,
61-
location_refs: QuerySet[LocationFindingReference],
58+
location_refs: QuerySet[LocationFindingReference], # noqa: N805
6259
user: Dojo_User,
6360
**kwargs: dict,
6461
) -> None:
@@ -69,11 +66,9 @@ def mitigate_location_status(
6966
status=FindingLocationStatus.Mitigated,
7067
)
7168

72-
@dojo_async_task
73-
@app.task()
69+
@app.task
7470
def reactivate_location_status(
75-
self,
76-
location_refs: QuerySet[LocationFindingReference],
71+
location_refs: QuerySet[LocationFindingReference], # noqa: N805
7772
**kwargs: dict,
7873
) -> None:
7974
"""Reactivate all given (mitigated) locations refs"""
@@ -89,10 +84,10 @@ def chunk_locations_and_disperse(
8984
locations: list[AbstractLocation],
9085
**kwargs: dict,
9186
) -> None:
92-
self.add_locations_to_unsaved_finding(finding, locations, sync=True)
87+
dojo_dispatch_task(self.add_locations_to_unsaved_finding, finding, locations, sync=True)
9388

89+
@staticmethod
9490
def clean_unsaved_locations(
95-
self,
9691
locations: list[AbstractLocation],
9792
) -> None:
9893
"""
@@ -110,15 +105,15 @@ def chunk_locations_and_reactivate(
110105
location_refs: QuerySet[LocationFindingReference],
111106
**kwargs: dict,
112107
) -> None:
113-
self.reactivate_location_status(location_refs, sync=True)
108+
dojo_dispatch_task(self.reactivate_location_status, location_refs, sync=True)
114109

115110
def chunk_locations_and_mitigate(
116111
self,
117112
location_refs: QuerySet[LocationFindingReference],
118113
user: Dojo_User,
119114
**kwargs: dict,
120115
) -> None:
121-
self.mitigate_location_status(location_refs, user, sync=True)
116+
dojo_dispatch_task(self.mitigate_location_status, location_refs, user, sync=True)
122117

123118
def update_location_status(
124119
self,

0 commit comments

Comments
 (0)