diff --git a/dojo/management/commands/migrate_endpoints_to_locations.py b/dojo/management/commands/migrate_endpoints_to_locations.py index 6ae1996de5..fd62699862 100644 --- a/dojo/management/commands/migrate_endpoints_to_locations.py +++ b/dojo/management/commands/migrate_endpoints_to_locations.py @@ -3,10 +3,12 @@ import logging import time from collections import defaultdict +from itertools import islice +from django.conf import settings from django.core.management.base import BaseCommand from django.db import connection, transaction -from django.db.models import Prefetch +from django.db.models import Case, Exists, OuterRef, Prefetch, Value, When from django.utils import timezone from dojo.location.models import Location, LocationFindingReference, LocationProductReference @@ -17,9 +19,10 @@ logger = logging.getLogger(__name__) -# Chunk size for the DB iterator. Tunable via --batch-size. +# Endpoints read per DB iterator chunk, and rows per bulk write. Tunable via +# --batch-size. DEFAULT_BATCH_SIZE = 1000 -# How often to emit per-chunk progress lines. Tunable via --progress-every. +# Minimum endpoints between progress lines. Tunable via --progress-every. DEFAULT_PROGRESS_EVERY = 50 @@ -42,14 +45,96 @@ def _suspend_auto_now_add(model, field_name: str): field.auto_now_add = saved +class _ChunkRows: + + """ + Rows accumulated for one chunk, deduplicated on their unique constraint. + + The migration builds every row for a chunk in memory and then issues one + write per kind. Deduplicating here is what makes those writes safe: several + legacy Endpoints can normalize onto a single Location, and Postgres rejects + an ``ON CONFLICT DO UPDATE`` statement that would touch the same row twice. + + Each row is paired with the id of the endpoint that produced it so a failed + batch write can be retried per row and attributed to the right endpoint. + """ + + def __init__(self) -> None: + self.meta: list[DojoMeta] = [] + self.meta_endpoint_ids: list[int | None] = [] + self._seen_meta: set[tuple[int, str]] = set() + + self.finding_refs: list[LocationFindingReference] = [] + self.finding_ref_endpoint_ids: list[int | None] = [] + self._seen_finding_refs: set[tuple[int, int]] = set() + + # (location_id, product_id) -> (status, product, location, endpoint_id) + self._product_refs: dict[tuple[int, int], tuple[str, Product, Location, int | None]] = {} + + def add_meta(self, meta: DojoMeta, location: Location, endpoint_id: int | None) -> None: + """Queue a DojoMeta copy, keyed on its unique_together (location, name).""" + key = (location.id, meta.name) + if key in self._seen_meta: + return + self._seen_meta.add(key) + self.meta.append(DojoMeta(name=meta.name, value=meta.value, location=location)) + self.meta_endpoint_ids.append(endpoint_id) + + def add_finding_ref(self, reference: LocationFindingReference, endpoint_id: int | None) -> None: + """Queue a finding reference, keyed on its unique (location, finding).""" + key = (reference.location.id, reference.finding_id) + if key in self._seen_finding_refs: + return + self._seen_finding_refs.add(key) + self.finding_refs.append(reference) + self.finding_ref_endpoint_ids.append(endpoint_id) + + def record_product( + self, + product: Product, + location: Location, + status: str, + endpoint_id: int | None, + ) -> None: + """ + Queue a product reference, letting Active win within the chunk. + + The status recorded here only reflects this chunk; + ``_reconcile_product_statuses`` has the final say once every finding + reference is in place. + """ + key = (location.id, product.id) + existing = self._product_refs.get(key) + if existing is None or ( + status == ProductLocationStatus.Active + and existing[0] != ProductLocationStatus.Active + ): + self._product_refs[key] = (status, product, location, endpoint_id) + + def product_ref_rows(self) -> tuple[list[LocationProductReference], list[int | None]]: + rows = [ + LocationProductReference( + location=location, + product=product, + status=status, + relationship="", + relationship_data={}, + ) + for status, product, location, _ in self._product_refs.values() + ] + endpoint_ids = [endpoint_id for _, _, _, endpoint_id in self._product_refs.values()] + return rows, endpoint_ids + + # Phases tracked by --benchmark. Order is preserved in the summary table. PHASES = ( - "fetch_endpoint", # iterator yields the next endpoint - "url_create", # URL.get_or_create_from_values + Location side-effect + "fetch_chunk", # iterator + prefetch queries for the next chunk + "url_create", # URL.bulk_get_or_create + Location rows for the chunk "tags", # batched endpoint tag copy onto locations "meta", # DojoMeta copy onto the location "finding_refs", # LocationFindingReference creation per Endpoint_Status "product_refs", # LocationProductReference creation + "reconcile", # recompute product-reference status from finding refs ) @@ -61,6 +146,24 @@ class Command(BaseCommand): - Endpoints -> URL (which will create a Location) - Products on Endpoint -> LocationProductReference - Findings on Endpoints -> LocationProductReference + + The command is safe to re-run and converges on the state the source + Endpoints/Endpoint_Statuses describe: + + - ``URL``/``Location`` rows are matched on ``identity_hash``, so an endpoint + that already migrated reuses its location instead of duplicating it. + - ``LocationFindingReference`` is upserted on (location, finding): a status, + auditor or audit time that has since changed on the source + ``Endpoint_Status`` is written over the stale value, while the original + ``created`` timestamp is preserved. + - ``LocationProductReference`` status is recomputed from the finding + references after each chunk, using the same rule as + ``Location.status_from_product`` (Active when the location has at least + one Active finding for that product). This repairs statuses left behind by + an earlier run and makes the outcome independent of the order endpoints + are visited in. + - ``DojoMeta`` rows and tag copies are insert-only, so re-runs neither + duplicate them nor overwrite edits made after the first migration. """ help = "Usage: manage.py migrate_endpoints_to_locations" @@ -70,13 +173,16 @@ def add_arguments(self, parser): "--batch-size", type=int, default=DEFAULT_BATCH_SIZE, - help=f"Endpoint.objects.iterator() chunk size (default: {DEFAULT_BATCH_SIZE}).", + help="Endpoints read, and rows written, per batch " + f"(default: {DEFAULT_BATCH_SIZE}).", ) parser.add_argument( "--progress-every", type=int, default=DEFAULT_PROGRESS_EVERY, - help=f"Emit a progress line every N endpoints (default: {DEFAULT_PROGRESS_EVERY}).", + help="Emit a progress line once at least N endpoints have been migrated " + "since the last one; lines land on batch boundaries, so a value " + f"below --batch-size reports once per batch (default: {DEFAULT_PROGRESS_EVERY}).", ) parser.add_argument( "--benchmark", @@ -196,64 +302,307 @@ def _flush_location_tags(self) -> None: self.pending_tag_locations.clear() self.pending_endpoint_tags.clear() - # -- Migration logic -------------------------------------------------- + # -- Chunked reads --------------------------------------------------------- - def _endpoint_to_url(self, endpoint: Endpoint) -> Location: - # Create the raw URL object first - # This should create the location object as well + def _iter_chunks(self, queryset): + """ + Yield lists of at most ``batch_size`` endpoints. + + ``queryset.iterator(chunk_size=...)`` already fetches and prefetches a + chunk at a time, so re-buffering into a list of the same size costs no + extra memory over what Django is already holding, and it gives the + migration a batch to write in one shot per phase. + """ + iterator = queryset.iterator(chunk_size=self.batch_size) + while True: + t = self._bench_start() + chunk = list(islice(iterator, self.batch_size)) + self._bench_end("fetch_chunk", t) + if not chunk: + return + yield chunk + + # -- Bulk writes with per-row fallback ------------------------------------ + + def _bulk_create_rows( + self, + phase: str, + model, + rows: list, + endpoint_ids: list[int | None], + cm_factory=None, + conflict_kwargs: dict | None = None, + ) -> None: + """ + Write ``rows`` in one ``bulk_create``, retrying per row if that fails. + + The whole chunk goes out as a single statement on the happy path. If it + raises we fall back to one row at a time so a single bad row can't + discard the valid rows it was batched with — the same + batch-then-isolate shape ``_flush_location_tags`` uses. Failures are + attributed to the endpoint that produced the row via + ``endpoint_ids[i]``. + + ``conflict_kwargs`` selects the upsert behaviour (``ignore_conflicts`` + vs ``update_conflicts`` + ``update_fields``/``unique_fields``); callers + must deduplicate rows on the conflict target first, because Postgres + rejects a statement whose ``ON CONFLICT DO UPDATE`` would touch the + same row twice. + + ``cm_factory`` is a callable returning a context manager to wrap each + write (used for ``_suspend_auto_now_add``); it must be a factory rather + than an instance because the retry path enters it once per row. + """ + if not rows: + return + cm_factory = cm_factory or contextlib.nullcontext + conflict_kwargs = conflict_kwargs or {"ignore_conflicts": True} t = self._bench_start() - url = URL.get_or_create_from_values( - protocol=endpoint.protocol, - user_info=endpoint.userinfo, - host=endpoint.host, - port=endpoint.port, - path=endpoint.path, - query=endpoint.query, - fragment=endpoint.fragment, - ) - self._bench_end("url_create", t) + try: + try: + with cm_factory(): + model.objects.bulk_create( + rows, batch_size=self.batch_size, **conflict_kwargs, + ) + except Exception: + logger.exception( + "Batched %s write failed for %d row(s); " + "retrying one row at a time", + phase, len(rows), + ) + for row, endpoint_id in zip(rows, endpoint_ids, strict=True): + try: + with transaction.atomic(), cm_factory(): + model.objects.bulk_create([row], **conflict_kwargs) + except Exception as exc: + logger.exception( + "Failed to write %s row for endpoint id=%s; continuing", + phase, endpoint_id, + ) + self._record_endpoint_failure(endpoint_id, exc) + finally: + self._bench_end(phase, t) - # Queue endpoint tags for one bulk write per migration batch instead - # of making Tagulous look up and attach tags once per endpoint. + def _reconcile_product_statuses(self, location_ids: list[int]) -> None: + """ + Recompute ``LocationProductReference.status`` for the chunk's locations. + + A product reference is Active exactly when the location has at least + one Active finding reference for that product — the same rule + ``Location.status_from_product`` applies. Recomputing it from the + finding refs in one set-based UPDATE per chunk is what makes re-runs + converge: + + - it repairs stale statuses left by an earlier run (or by the previous + first-write-wins behaviour of ``associate_with_product``); + - it is order independent, so a location shared by endpoints in + different chunks ends up Active if *any* of its findings is Active, + without the migration having to remember cross-chunk state. + + Only locations touched by this chunk are considered, so the statement + stays bounded no matter how large the install is. + """ + if not location_ids: + return t = self._bench_start() - tag_names = {tag.name for tag in endpoint.tags.all()} - if tag_names: - self._queue_location_tags(endpoint, url.location, tag_names) - self._bench_end("tags", t) - - # Add any metadata from the endpoint to the location. - # bulk_create with ignore_conflicts mirrors the previous get_or_create - # semantics — DojoMeta.unique_together = (location, name) so any - # conflict is by definition the same row we'd otherwise have fetched. - # One INSERT per endpoint instead of SELECT+INSERT per meta row. + try: + has_active_finding = LocationFindingReference.objects.filter( + location_id=OuterRef("location_id"), + finding__test__engagement__product_id=OuterRef("product_id"), + status=FindingLocationStatus.Active, + ) + try: + LocationProductReference.objects.filter( + location_id__in=location_ids, + ).update( + status=Case( + When(Exists(has_active_finding), then=Value(ProductLocationStatus.Active)), + default=Value(ProductLocationStatus.Mitigated), + ), + updated=timezone.now(), + ) + except Exception: + # Reconciliation is a repair pass over rows that are already + # committed, so a failure here costs accuracy on this chunk's + # statuses, not data. Keep going rather than abort the run. + logger.exception( + "Product-reference status reconciliation failed for %d location(s); " + "continuing", len(location_ids), + ) + finally: + self._bench_end("reconcile", t) + + # -- Migration logic -------------------------------------------------- + + def _resolve_locations(self, endpoints: list[Endpoint]) -> list[tuple[Endpoint, Location]]: + """ + Resolve one Location per endpoint for the whole chunk. + + ``URL.bulk_get_or_create`` does the work in ~3 queries per chunk (one + lookup by ``identity_hash`` plus a ``bulk_create`` for the parent + Location rows and the URL rows) in place of the per-endpoint + ``get_or_create``, which cost a savepoint pair, a SELECT, a + ``validate_unique`` SELECT and two INSERTs each. + + Each URL is validated individually first. ``bulk_get_or_create`` only + calls ``clean()``, whereas the per-endpoint ``save()`` this replaces ran + ``full_clean()`` — so validating here keeps an endpoint whose values + can't pass field validation (empty host, out-of-range port, …) from + being written as an invalid Location, and reports it against its own id + instead of taking the chunk down. ``validate_unique`` is left off: the + identity_hash lookup ``bulk_get_or_create`` already does resolves + duplicates without a SELECT per row. + + If the bulk write itself still fails we fall back to the per-endpoint + ``get_or_create`` path for the chunk, which migrates every good endpoint + and attributes the failure to the one that caused it. + """ t = self._bench_start() - meta_rows = [ - DojoMeta(name=m.name, value=m.value, location=url.location) - for m in endpoint.endpoint_meta.all() - ] - if meta_rows: - DojoMeta.objects.bulk_create(meta_rows, ignore_conflicts=True) - self._bench_end("meta", t) + try: + pairs: list[tuple[Endpoint, URL]] = [] + for endpoint in endpoints: + url = URL.from_parts( + protocol=endpoint.protocol, + user_info=endpoint.userinfo, + host=endpoint.host, + port=endpoint.port, + path=endpoint.path, + query=endpoint.query, + fragment=endpoint.fragment, + ) + try: + # Mirrors what BaseModelWithoutTimeMeta.save() would have + # done for this URL, including the flag it keys off. + if settings.V3_FEATURE_LOCATIONS: + url.full_clean(validate_unique=False, validate_constraints=False) + else: + url.clean() + except Exception as exc: + endpoint_id = getattr(endpoint, "id", None) + logger.exception( + "Failed to migrate endpoint id=%s; continuing", endpoint_id, + ) + self._record_endpoint_failure(endpoint_id, exc) + continue + pairs.append((endpoint, url)) - return url.location + if not pairs: + return [] - def _convert_endpoint_status_to_string_status(self, endpoint_status: Endpoint_Status) -> str: + try: + saved = URL.bulk_get_or_create([url for _, url in pairs]) + except Exception: + logger.exception( + "Bulk location resolution failed for %d endpoint(s); " + "falling back to one get_or_create per endpoint", + len(pairs), + ) + resolved = [] + for endpoint, url in pairs: + try: + resolved.append((endpoint, URL.get_or_create_from_object(url).location)) + except Exception as exc: + endpoint_id = getattr(endpoint, "id", None) + logger.exception( + "Failed to migrate endpoint id=%s; continuing", endpoint_id, + ) + self._record_endpoint_failure(endpoint_id, exc) + return resolved + + return [ + (endpoint, saved_url.location) + for (endpoint, _), saved_url in zip(pairs, saved, strict=True) + ] + finally: + self._bench_end("url_create", t) + + def _process_chunk(self, endpoints: list[Endpoint]) -> None: """ - Start the conversion with the "special" statuses first since we are moving to a model - of having a single status possible rather than a combo of many + Migrate a chunk of endpoints with one bulk write per phase. + + Everything between the location resolve and the writes is pure Python + over already-prefetched data, so the DB sees a fixed handful of + statements per chunk rather than a dozen per endpoint. """ - if endpoint_status.risk_accepted: - return FindingLocationStatus.RiskAccepted - if endpoint_status.false_positive: - return FindingLocationStatus.FalsePositive - if endpoint_status.out_of_scope: - return FindingLocationStatus.OutOfScope - if endpoint_status.mitigated: - return FindingLocationStatus.Mitigated - # Default to Active - return FindingLocationStatus.Active + resolved = self._resolve_locations(endpoints) + if not resolved: + return + + rows = _ChunkRows() + for endpoint, location in resolved: + endpoint_id = getattr(endpoint, "id", None) + try: + # Queue endpoint tags for one bulk write per chunk instead of + # making Tagulous look up and attach tags once per endpoint. + t = self._bench_start() + tag_names = {tag.name for tag in endpoint.tags.all()} + if tag_names: + self._queue_location_tags(endpoint, location, tag_names) + self._bench_end("tags", t) + + for meta in endpoint.endpoint_meta.all(): + rows.add_meta(meta, location, endpoint_id) + + # Track the endpoint's own product as a contributor for the + # post-migration tag inheritance pass (the no-findings branch + # of `_collect_references` also depends on this product, and it + # won't be tracked otherwise). + if endpoint.product_id: + self._track_product_location(endpoint.product, location) + + self._collect_references(endpoint, location, rows) + except Exception as exc: + logger.exception("Failed to migrate endpoint id=%s; continuing", endpoint_id) + self._record_endpoint_failure(endpoint_id, exc) + + # DojoMeta: `ignore_conflicts` on unique_together (location, name). A + # conflict is by definition the row we would otherwise have fetched, so + # skipping it keeps re-runs no-ops and leaves any post-migration edit to + # a metadata value alone. + self._bulk_create_rows("meta", DojoMeta, rows.meta, rows.meta_endpoint_ids) + + # LocationFindingReference: upsert on (location, finding) so a re-run + # syncs the status/auditor/audit_time when the source Endpoint_Status + # has moved on. `created` is deliberately absent from update_fields — + # the original creation timestamp is preserved on repeat runs. + self._bulk_create_rows( + "finding_refs", LocationFindingReference, + rows.finding_refs, rows.finding_ref_endpoint_ids, + cm_factory=lambda: _suspend_auto_now_add(LocationFindingReference, "created"), + conflict_kwargs={ + "update_conflicts": True, + "update_fields": ["status", "auditor", "audit_time", "updated"], + "unique_fields": ["location", "finding"], + }, + ) + + # LocationProductReference: insert-only. The status carried here is + # derived from this chunk alone, so it is not authoritative — the + # reconciliation pass below recomputes it from the finding refs. + product_ref_rows, product_ref_endpoint_ids = rows.product_ref_rows() + self._bulk_create_rows( + "product_refs", LocationProductReference, + product_ref_rows, product_ref_endpoint_ids, + ) + + self._reconcile_product_statuses(list({location.id for _, location in resolved})) + + def _collect_references( + self, + endpoint: Endpoint, + location: Location, + rows: _ChunkRows, + ) -> None: + """ + Accumulate this endpoint's finding/product reference rows. + + Pure Python over the prefetched ``status_endpoint`` list — no queries. + Bypasses ``Location.associate_with_finding`` (which would trigger + full_clean validation plus the post_save inherit_tags signal per row) + and is semantically equivalent to it for this migration. + """ + endpoint_id = getattr(endpoint, "id", None) - def _associate_location_with_findings(self, endpoint: Endpoint, location: Location) -> None: # Pull the prefetched list once. Avoids the redundant `.exists()` round- # trip the prior code did and lets the loop iterate prefetched data. statuses = list(endpoint.status_endpoint.all()) @@ -261,34 +610,11 @@ def _associate_location_with_findings(self, endpoint: Endpoint, location: Locati # No findings — associate with the endpoint's product if one exists. if not statuses: if endpoint.product_id: - t_p = self._bench_start() - LocationProductReference.objects.bulk_create( - [LocationProductReference( - location=location, - product=endpoint.product, - status=ProductLocationStatus.Mitigated, - relationship="", - relationship_data={}, - )], - ignore_conflicts=True, + rows.record_product( + endpoint.product, location, ProductLocationStatus.Mitigated, endpoint_id, ) - self._bench_end("product_refs", t_p) return - # Build LFR rows for every status, and build LPR rows deduplicated by - # product, deriving the product status as Active iff any of THIS - # endpoint's findings on that product are Active. This bypasses - # `Location.associate_with_finding` (which would trigger full_clean - # validation + the post_save inherit_tags signal per row) and is - # semantically equivalent to the prior behavior in the common case - # where each endpoint maps to a distinct location. As a side-effect - # it also fixes the existing `associate_with_product` first-write- - # wins bug (where a Mitigated status would stick even when later - # Active findings come in for the same product). - finding_refs: list[LocationFindingReference] = [] - product_status_by_id: dict[int, str] = {} - product_obj_by_id: dict[int, object] = {} - for endpoint_status in statuses: finding = endpoint_status.finding if finding is None: @@ -299,6 +625,15 @@ def _associate_location_with_findings(self, endpoint: Endpoint, location: Locati # differs from endpoint.product). self._track_product_location(product, location) status = self._convert_endpoint_status_to_string_status(endpoint_status) + rows.record_product( + product, + location, + ProductLocationStatus.Active + if status == FindingLocationStatus.Active + else ProductLocationStatus.Mitigated, + endpoint_id, + ) + # Endpoint_Status.date is a Date; the original code persisted # the same midnight-aware datetime in a post-save UPDATE. We # set it directly here — bulk_create skips auto_now_add so the @@ -308,7 +643,7 @@ def _associate_location_with_findings(self, endpoint: Endpoint, location: Locati endpoint_status.date.month, endpoint_status.date.day, )) - finding_refs.append(LocationFindingReference( + rows.add_finding_ref(LocationFindingReference( location=location, finding=finding, status=status, @@ -317,42 +652,23 @@ def _associate_location_with_findings(self, endpoint: Endpoint, location: Locati relationship="", relationship_data={}, created=created_dt, - )) - if product.id not in product_obj_by_id: - product_obj_by_id[product.id] = product - product_status_by_id[product.id] = ( - ProductLocationStatus.Active - if status == FindingLocationStatus.Active - else ProductLocationStatus.Mitigated - ) - elif (status == FindingLocationStatus.Active - and product_status_by_id[product.id] != ProductLocationStatus.Active): - product_status_by_id[product.id] = ProductLocationStatus.Active - - t_f = self._bench_start() - if finding_refs: - with _suspend_auto_now_add(LocationFindingReference, "created"): - LocationFindingReference.objects.bulk_create( - finding_refs, ignore_conflicts=True, batch_size=500, - ) - self._bench_end("finding_refs", t_f) - - t_p = self._bench_start() - if product_obj_by_id: - product_refs = [ - LocationProductReference( - location=location, - product=product_obj_by_id[pid], - status=product_status_by_id[pid], - relationship="", - relationship_data={}, - ) - for pid in product_obj_by_id - ] - LocationProductReference.objects.bulk_create( - product_refs, ignore_conflicts=True, batch_size=500, - ) - self._bench_end("product_refs", t_p) + ), endpoint_id) + + def _convert_endpoint_status_to_string_status(self, endpoint_status: Endpoint_Status) -> str: + """ + Start the conversion with the "special" statuses first since we are moving to a model + of having a single status possible rather than a combo of many + """ + if endpoint_status.risk_accepted: + return FindingLocationStatus.RiskAccepted + if endpoint_status.false_positive: + return FindingLocationStatus.FalsePositive + if endpoint_status.out_of_scope: + return FindingLocationStatus.OutOfScope + if endpoint_status.mitigated: + return FindingLocationStatus.Mitigated + # Default to Active + return FindingLocationStatus.Active # -- Progress + summary reporting ---------------------------------------- @@ -367,17 +683,23 @@ def _fmt_duration(seconds: float) -> str: return f"{m}m {s}s" return f"{s}s" - def _log_progress(self, i: int, total: int, run_t0: float, queries_this_chunk: int | None) -> None: + def _log_progress( + self, + i: int, + total: int, + run_t0: float, + queries_this_window: int | None, + endpoints_this_window: int, + ) -> None: elapsed = time.time() - run_t0 rate = i / elapsed if elapsed > 0 else 0.0 remaining = (total - i) / rate if rate > 0 else 0.0 pct = (i / total * 100.0) if total else 100.0 line = (f"Migrated {i:,}/{total:,} endpoints ({pct:.1f}%) — " f"{rate:.1f} endpoints/sec — ETA {self._fmt_duration(remaining)}") - if queries_this_chunk is not None: - # Per-endpoint query count for this chunk window only. - chunk_size = self.progress_every - line += f" — {queries_this_chunk / chunk_size:.1f} queries/endpoint" + if queries_this_window is not None and endpoints_this_window: + # Per-endpoint query count for this reporting window only. + line += f" — {queries_this_window / endpoints_this_window:.1f} queries/endpoint" self.stdout.write(self.style.SUCCESS(line)) if self.benchmark: @@ -485,11 +807,30 @@ def handle(self, *args, **options): if self.query_count: connection.force_debug_cursor = True - queries_at_chunk_start = len(connection.queries) - else: - queries_at_chunk_start = 0 # unused + queries_at_window_start = len(connection.queries) if self.query_count else 0 - # Allow endpoints to work with V3/Locations enabled + # Lazy import: the inheritance module imports the full model layer, so + # keep it out of management-command discovery. + from dojo.tags import inheritance as tag_inheritance # noqa: PLC0415 + + # Allow endpoints to work with V3/Locations enabled, and keep the + # per-instance tag-inheritance signals out of the hot loop. + # + # Every `Location` row created here fires `inherit_tags_on_instance`, + # which calls `Location.all_related_products()` — an OR across two + # multi-join paths (`LocationProductReference` and + # `Finding -> Test -> Engagement -> Product` via + # `LocationFindingReference`) that Postgres can only answer by + # materialising the whole join and filtering it. Because this migration + # is itself filling `LocationFindingReference`, that query gets more + # expensive with every endpoint migrated, which is what turned the run + # into an O(n^2) crawl (~105ms of the ~132ms each `Location` insert + # cost, on a small dataset, growing from there). + # + # The work is redundant regardless: a freshly created Location has no + # product references yet, so the signal has nothing to inherit. Correct + # inheritance is applied in bulk by `_run_tag_inheritance()` once the + # references exist, which is why that pass already exists. with Endpoint.allow_endpoint_init(): # Prefetch everything the per-endpoint loop will touch so the # iterator doesn't trigger N+1 joins: @@ -498,7 +839,7 @@ def handle(self, *args, **options): # - `tags` and `endpoint_meta` are prefetched managers # - `status_endpoint` is prefetched together with the FK chain # `finding -> test -> engagement -> product` and `mitigated_by` - # so `associate_with_finding` can read them without queries. + # so the reference rows can be built without queries. queryset = ( Endpoint.objects.all() .select_related("product") @@ -525,55 +866,37 @@ def handle(self, *args, **options): run_t0 = time.time() i = 0 - # Process each endpoint - for i, endpoint in enumerate(queryset.iterator(chunk_size=self.batch_size), 1): - t_fetch = self._bench_start() - # iterator already produced `endpoint`; bill nothing meaningful - # to fetch_endpoint here — kept as a placeholder that B1's - # prefetch will start incrementing. - self._bench_end("fetch_endpoint", t_fetch) - - # Wrap the per-endpoint work so one failure doesn't abort a - # multi-hour migration. We log the full traceback and record - # the endpoint id, then continue. The bulk_create-based hot - # path makes partial-state on failure unlikely (each phase - # is its own bulk insert), and any rows that DID land remain - # valid and idempotent on re-run. - try: - # Get the URL object first - location = self._endpoint_to_url(endpoint) - # Track the endpoint's own product as a contributor for the - # post-migration tag inheritance pass (the no-findings - # branch of _associate_location_with_findings also depends - # on this product, and it won't be tracked otherwise). - if endpoint.product_id: - self._track_product_location(endpoint.product, location) - # Associate the URL with the findings associated with the Findings - # the association to a finding will also apply to a product automatically - self._associate_location_with_findings(endpoint, location) - except Exception as exc: - endpoint_id = getattr(endpoint, "id", None) - logger.exception("Failed to migrate endpoint id=%s; continuing", endpoint_id) - self._record_endpoint_failure(endpoint_id, exc) - - # Flush independently of per-endpoint success so a failing - # endpoint at a batch boundary cannot leave the queue growing. - if i % self.batch_size == 0: + last_reported = 0 + # Process endpoints a chunk at a time. Each chunk issues a fixed + # handful of statements — one lookup + two inserts for the + # locations, then one write per reference/meta/tag phase — instead + # of a dozen round trips per endpoint. `_process_chunk` isolates + # failures itself: per-endpoint for anything that can raise while + # building rows, and per-row on a failed batch write, so one bad + # endpoint still can't abort a multi-hour migration. + with tag_inheritance.suppress_tag_inheritance(): + for chunk in self._iter_chunks(queryset): + self._process_chunk(chunk) + i += len(chunk) + + # Flush independently of per-endpoint success so a failing + # endpoint at a chunk boundary cannot leave the queue growing. self._flush_location_tags() - # Progress report every --progress-every endpoints - if i % self.progress_every == 0: - queries_in_chunk = None - if self.query_count: - queries_in_chunk = len(connection.queries) - queries_at_chunk_start - # Trim the query log so memory doesn't balloon on long runs; - # after clear() the next chunk's baseline is 0. - connection.queries_log.clear() - queries_at_chunk_start = 0 - self._log_progress(i, endpoint_count, run_t0, queries_in_chunk) - - # Persist the final partial batch before reporting completion. - self._flush_location_tags() + # Progress report once at least --progress-every endpoints + # have been migrated since the last line. + if i - last_reported >= self.progress_every or i >= endpoint_count: + queries_in_window = None + if self.query_count: + queries_in_window = len(connection.queries) - queries_at_window_start + # Trim the query log so memory doesn't balloon on long runs; + # after clear() the next window's baseline is 0. + connection.queries_log.clear() + queries_at_window_start = 0 + self._log_progress( + i, endpoint_count, run_t0, queries_in_window, i - last_reported, + ) + last_reported = i elapsed = time.time() - run_t0 successful = i - len(self.failed_endpoints) diff --git a/unittests/test_migrate_endpoints_to_locations.py b/unittests/test_migrate_endpoints_to_locations.py index 8079eb707c..aacda60b98 100644 --- a/unittests/test_migrate_endpoints_to_locations.py +++ b/unittests/test_migrate_endpoints_to_locations.py @@ -1,11 +1,29 @@ +import datetime from io import StringIO from unittest.mock import patch from django.core.management import call_command from django.test import TestCase, override_settings +from django.utils import timezone -from dojo.location.models import Location -from dojo.models import Endpoint, Product, Product_Type +from dojo.location.models import ( + Location, + LocationFindingReference, + LocationProductReference, +) +from dojo.location.status import FindingLocationStatus, ProductLocationStatus +from dojo.models import ( + Dojo_User, + DojoMeta, + Endpoint, + Endpoint_Status, + Engagement, + Finding, + Product, + Product_Type, + Test, + Test_Type, +) from dojo.tags.utils import bulk_add_tag_mapping from dojo.url.models import URL @@ -19,6 +37,7 @@ def setUp(self): description="Test product", prod_type=product_type, ) + self.reporter = Dojo_User.objects.create(username="endpoint-migration-reporter") def _make_endpoint(self, host, tags): with Endpoint.allow_endpoint_init(): @@ -30,6 +49,62 @@ def _make_endpoint(self, host, tags): endpoint.tags.add(*tags) return endpoint + def _make_test(self): + engagement = Engagement.objects.create( + name="Endpoint migration engagement", + product=self.product, + target_start=timezone.now().date(), + target_end=timezone.now().date(), + ) + test_type, _ = Test_Type.objects.get_or_create(name="Endpoint migration test type") + return Test.objects.create( + engagement=engagement, + test_type=test_type, + scan_type="Endpoint migration scan", + target_start=timezone.now(), + target_end=timezone.now(), + ) + + def _make_endpoint_with_status(self, host, *, active, test=None, title=None): + """Create an Endpoint carrying one Endpoint_Status for a new Finding.""" + test = test or self._make_test() + finding = Finding.objects.create( + title=title or f"Finding for {host}", + test=test, + severity="High", + numerical_severity="S1", + description="Test finding", + active=active, + verified=False, + reporter=self.reporter, + ) + with Endpoint.allow_endpoint_init(): + endpoint = Endpoint.objects.create( + protocol="https", + host=host, + product=self.product, + ) + status = Endpoint_Status.objects.create( + endpoint=endpoint, + finding=finding, + date=datetime.date(2024, 5, 17), + mitigated=not active, + ) + return endpoint, finding, status + + def _run(self, **kwargs): + stdout = StringIO() + call_command( + "migrate_endpoints_to_locations", + progress_every=kwargs.pop("progress_every", 100), + stdout=stdout, + **kwargs, + ) + return stdout.getvalue() + + def _location_for(self, host): + return URL.objects.get(host=host).location + def test_endpoint_tags_are_copied_in_deduplicated_batches(self): # Four legacy Endpoints resolve to one Location. With batches of three, # at least one batch contains that Location more than once regardless of @@ -203,3 +278,241 @@ def fail_one_tag(tag_to_locations, **kwargs): tag_model = Location.tags.tag_model self.assertEqual(tag_model.objects.get(name="failing-tag").count, 1) self.assertEqual(tag_model.objects.get(name="healthy-tag").count, 1) + + def test_rerun_creates_no_duplicate_rows(self): + self._make_endpoint_with_status("idempotent.example.com", active=True) + endpoint = self._make_endpoint("plain.example.com", ["plain-tag"]) + DojoMeta.objects.create(name="owner", value="team-a", endpoint=endpoint) + + self._run() + location = self._location_for("idempotent.example.com") + first_created = LocationFindingReference.objects.get(location=location).created + counts = ( + Location.objects.count(), + URL.objects.count(), + LocationFindingReference.objects.count(), + LocationProductReference.objects.count(), + DojoMeta.objects.filter(location__isnull=False).count(), + ) + self.assertEqual(counts, (2, 2, 1, 2, 1)) + + self._run() + + self.assertEqual( + ( + Location.objects.count(), + URL.objects.count(), + LocationFindingReference.objects.count(), + LocationProductReference.objects.count(), + DojoMeta.objects.filter(location__isnull=False).count(), + ), + counts, + ) + # A rerun must not restamp `created` on a reference it already wrote. + self.assertEqual( + LocationFindingReference.objects.get(location=location).created, + first_created, + ) + + def test_rerun_syncs_changed_finding_reference_status(self): + _, _, status = self._make_endpoint_with_status("changed.example.com", active=True) + self._run() + + location = self._location_for("changed.example.com") + reference = LocationFindingReference.objects.get(location=location) + self.assertEqual(reference.status, FindingLocationStatus.Active) + created_before = reference.created + + # The source Endpoint_Status moves on after the first migration. + Endpoint_Status.objects.filter(pk=status.pk).update(risk_accepted=True) + + self._run() + + reference.refresh_from_db() + self.assertEqual(reference.status, FindingLocationStatus.RiskAccepted) + self.assertEqual(reference.created, created_before) + self.assertEqual(LocationFindingReference.objects.count(), 1) + + def test_product_reference_status_is_reconciled_from_finding_references(self): + # Two endpoints sharing one product; only the first has an active + # finding, so the product reference for its location must be Active and + # the other's Mitigated. + test = self._make_test() + _, _, active_status = self._make_endpoint_with_status( + "active.example.com", active=True, test=test, title="Active finding", + ) + self._make_endpoint_with_status( + "mitigated.example.com", active=False, test=test, title="Mitigated finding", + ) + + self._run() + + active_location = self._location_for("active.example.com") + mitigated_location = self._location_for("mitigated.example.com") + self.assertEqual( + LocationProductReference.objects.get( + location=active_location, product=self.product, + ).status, + ProductLocationStatus.Active, + ) + self.assertEqual( + LocationProductReference.objects.get( + location=mitigated_location, product=self.product, + ).status, + ProductLocationStatus.Mitigated, + ) + + # A status that drifted out of sync with the finding references is + # repaired rather than left alone by `ignore_conflicts`. + LocationProductReference.objects.filter(location=active_location).update( + status=ProductLocationStatus.Mitigated, + ) + self._run() + self.assertEqual( + LocationProductReference.objects.get( + location=active_location, product=self.product, + ).status, + ProductLocationStatus.Active, + ) + + # And a genuine downgrade in the source data propagates: the previous + # per-endpoint code left the first-written Active status in place. + Endpoint_Status.objects.filter(pk=active_status.pk).update(mitigated=True) + self._run() + self.assertEqual( + LocationProductReference.objects.get( + location=active_location, product=self.product, + ).status, + ProductLocationStatus.Mitigated, + ) + + def test_shared_location_product_status_is_order_independent(self): + # Two endpoints normalising onto the same Location, split across + # separate chunks, with only the second carrying the active finding. + test = self._make_test() + for title, active in (("Mitigated first", False), ("Active second", True)): + finding = Finding.objects.create( + title=title, + test=test, + severity="High", + numerical_severity="S1", + description="Test finding", + active=active, + verified=False, + reporter=self.reporter, + ) + with Endpoint.allow_endpoint_init(): + endpoint = Endpoint.objects.create( + protocol="https", host="shared-status.example.com", product=self.product, + ) + Endpoint_Status.objects.create( + endpoint=endpoint, + finding=finding, + date=datetime.date(2024, 5, 17), + mitigated=not active, + ) + + self._run(batch_size=1) + + location = self._location_for("shared-status.example.com") + self.assertEqual(URL.objects.filter(host="shared-status.example.com").count(), 1) + self.assertEqual(LocationFindingReference.objects.filter(location=location).count(), 2) + self.assertEqual( + LocationProductReference.objects.get( + location=location, product=self.product, + ).status, + ProductLocationStatus.Active, + ) + + def test_invalid_endpoint_is_reported_and_its_chunk_still_migrates(self): + # `bulk_get_or_create` only calls clean(), so the command validates each + # URL itself to keep an endpoint that cannot pass field validation from + # being written as an invalid Location. + self._make_endpoint("good-one.example.com", []) + self._make_endpoint("good-two.example.com", []) + broken = self._make_endpoint("broken.example.com", []) + Endpoint.objects.filter(pk=broken.pk).update(host="") + + stdout = StringIO() + with self.assertLogs( + "dojo.management.commands.migrate_endpoints_to_locations", + level="ERROR", + ) as logs: + call_command( + "migrate_endpoints_to_locations", + batch_size=10, + progress_every=100, + stdout=stdout, + ) + + self.assertIn("Migrated 2/3 endpoints", stdout.getvalue()) + self.assertIn(str(broken.id), stdout.getvalue()) + self.assertTrue( + any(f"Failed to migrate endpoint id={broken.id}" in line for line in logs.output), + ) + self.assertEqual( + sorted(URL.objects.values_list("host", flat=True)), + ["good-one.example.com", "good-two.example.com"], + ) + + def test_failed_bulk_location_write_falls_back_per_endpoint(self): + self._make_endpoint("first-bulk.example.com", []) + self._make_endpoint("second-bulk.example.com", []) + + original = URL.bulk_get_or_create + calls = [] + + def fail_first_chunk(locations): + calls.append(len(locations)) + if len(calls) == 1: + msg = "simulated bulk location write failure" + raise RuntimeError(msg) + return original(locations) + + stdout = StringIO() + with ( + patch.object(URL, "bulk_get_or_create", side_effect=fail_first_chunk), + self.assertLogs( + "dojo.management.commands.migrate_endpoints_to_locations", + level="ERROR", + ) as logs, + ): + call_command( + "migrate_endpoints_to_locations", + batch_size=10, + progress_every=100, + stdout=stdout, + ) + + self.assertTrue( + any("falling back to one get_or_create per endpoint" in line for line in logs.output), + ) + # Every endpoint still migrates, via the per-endpoint path. + self.assertIn("Migrated 2/2 endpoints", stdout.getvalue()) + self.assertEqual( + sorted(URL.objects.values_list("host", flat=True)), + ["first-bulk.example.com", "second-bulk.example.com"], + ) + + def test_inheritance_signal_is_suppressed_during_the_main_loop(self): + # The per-Location post_save inheritance signal issues an OR-joined + # query whose cost grows with LocationFindingReference, so the hot loop + # must not fire it; `_run_tag_inheritance` applies inheritance in bulk + # once the reference rows exist. + self.product.tags.add("product-inherited") + self.product.enable_product_tag_inheritance = True + self.product.save(update_fields=["enable_product_tag_inheritance"]) + self._make_endpoint_with_status("inherit.example.com", active=True) + + with patch( + "dojo.location.models.Location.all_related_products", + ) as all_related_products: + self._run() + + all_related_products.assert_not_called() + + location = self._location_for("inherit.example.com") + self.assertEqual( + [tag.name for tag in location.inherited_tags.all()], + ["product-inherited"], + )