|
11 | 11 |
|
12 | 12 | logger = logging.getLogger(__name__) |
13 | 13 |
|
| 14 | +# Chunk size for DB cursor and progress report |
| 15 | +CHUNK_SIZE = 1000 |
| 16 | + |
14 | 17 |
|
15 | 18 | class Command(BaseCommand): |
16 | 19 |
|
@@ -88,10 +91,31 @@ def _associate_location_with_findings(self, endpoint: Endpoint, location: Locati |
88 | 91 | location.associate_with_product(product) |
89 | 92 |
|
90 | 93 | 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