Skip to content

Commit 1f907ce

Browse files
authored
Migration endpoints to locations fix (#14625)
* update migrate_endpoints_to_locations mgmt command to work when v3/locations enabled * include progress report on endpoint to location migration * linter fixes
1 parent 8ead201 commit 1f907ce

1 file changed

Lines changed: 31 additions & 7 deletions

File tree

dojo/management/commands/migrate_endpoints_to_locations.py

Lines changed: 31 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,9 @@
1111

1212
logger = logging.getLogger(__name__)
1313

14+
# Chunk size for DB cursor and progress report
15+
CHUNK_SIZE = 1000
16+
1417

1518
class Command(BaseCommand):
1619

@@ -88,10 +91,31 @@ def _associate_location_with_findings(self, endpoint: Endpoint, location: Locati
8891
location.associate_with_product(product)
8992

9093
def handle(self, *args, **options):
91-
# Start off with the endpoint objects - it should everything we need
92-
for endpoint in Endpoint.objects.all().iterator():
93-
# Get the URL object first
94-
location = self._endpoint_to_url(endpoint)
95-
# Associate the URL with the findings associated with the Findings
96-
# the association to a finding will also apply to a product automatically
97-
self._associate_location_with_findings(endpoint, location)
94+
# Allow endpoints to work with V3/Locations enabled
95+
with Endpoint.allow_endpoint_init():
96+
# Progress counter
97+
i = 0
98+
# Start off with the endpoint objects - it should contain everything we need
99+
queryset = Endpoint.objects.all()
100+
# Grab the total count so we can communicate progress
101+
endpoint_count = queryset.count()
102+
103+
# Process each endpoint
104+
for i, endpoint in enumerate(queryset.iterator(chunk_size=CHUNK_SIZE), 1):
105+
# Progress report every chunk
106+
if not i % CHUNK_SIZE:
107+
self.stdout.write(
108+
self.style.SUCCESS(
109+
f"Migrated {i}/{endpoint_count} endpoints...",
110+
),
111+
)
112+
# Get the URL object first
113+
location = self._endpoint_to_url(endpoint)
114+
# Associate the URL with the findings associated with the Findings
115+
# the association to a finding will also apply to a product automatically
116+
self._associate_location_with_findings(endpoint, location)
117+
self.stdout.write(
118+
self.style.SUCCESS(
119+
f"Migrated {i} total endpoints.",
120+
),
121+
)

0 commit comments

Comments
 (0)