From e6e7030ed787e6cc745d0d55588d587da064fc0c Mon Sep 17 00:00:00 2001 From: Cody Maffucci <46459665+Maffooch@users.noreply.github.com> Date: Mon, 17 Aug 2026 22:13:36 -0600 Subject: [PATCH] feat(locations): add a cancel hook to the endpoints-to-locations backfill The endpoints backfill can run for a long time. Add a cooperative cancel_callback stealth option, polled at the chunk boundary next to the existing progress hook, so an in-process caller (the Pro migration suite) can stop the run cleanly between chunks. Chunks are separately committed and the class is idempotent, so a cancelled run leaves a consistent, resumable state. - cancel_callback() -> bool added to stealth_options and read in handle(). - _should_cancel() polled after _emit_progress at the chunk boundary; a raising probe is treated as "keep going" so a transient error can't discard hours of work. - On cancel the run still applies tag inheritance for what it migrated and reports {"cancelled": True} in the summary; CLI runs are unaffected. Docs and the Pro changelog describe the new Cancel behaviour. Co-Authored-By: Claude Opus 5 --- .../admin/feature_flags/PRO__feature_flags.md | 2 +- .../PRO__migrating_from_endpoints.md | 2 + docs/content/releases/pro/changelog.md | 5 ++ .../migrate_endpoints_to_locations.py | 60 ++++++++++++-- .../test_migrate_endpoints_to_locations.py | 78 +++++++++++++++++++ 5 files changed, 138 insertions(+), 9 deletions(-) diff --git a/docs/content/admin/feature_flags/PRO__feature_flags.md b/docs/content/admin/feature_flags/PRO__feature_flags.md index 677c08e746..2158248147 100644 --- a/docs/content/admin/feature_flags/PRO__feature_flags.md +++ b/docs/content/admin/feature_flags/PRO__feature_flags.md @@ -68,7 +68,7 @@ The feature carries a **Restart Recommended** tag on the Feature Flags page for * The **Pro UI** and the **import pipeline** follow this toggle. After you enable Locations, new imports create Locations and the Locations pages appear on your next page load, without a restart. * The **Classic UI** pages and the `/api/v2` endpoint/location route wiring are decided from the `DD_V3_FEATURE_LOCATIONS` deployment setting when DefectDojo starts. This toggle does not change them, and restarting does not make it change them. If you use the Classic UI or depend on the `/api/v2` endpoint routes, set `DD_V3_FEATURE_LOCATIONS` to match and restart so every surface agrees. The stored toggle is seeded from that deployment setting on upgrade, so an instance that already ran with `DD_V3_FEATURE_LOCATIONS=True` comes up with the toggle already on (and locked), and the database owns the value from then on. -* Enabling existing history is not automatic. Your existing data stays as it is until you run the **data-migration suite** that appears under this row once Locations is on: three backfills (endpoints, dependencies, and source-code locations) followed by an identity rehash that unlocks once all three finish. Each is superuser-run, shows progress, and is safe to re-run. See [Migrating from Endpoints](/asset_modelling/locations/pro__migrating_from_endpoints/). +* Enabling existing history is not automatic. Your existing data stays as it is until you run the **data-migration suite** that appears under this row once Locations is on: three backfills (endpoints, dependencies, and source-code locations) followed by an identity rehash that unlocks once all three finish. Each is superuser-run, shows progress, is safe to re-run, and can be cancelled while running (it stops at the next batch boundary and can be resumed). See [Migrating from Endpoints](/asset_modelling/locations/pro__migrating_from_endpoints/). Enabling Locations is **self-service and one-way**: once it is on, the toggle locks (shown as **Cannot Be Disabled**), because turning it back off would require reversing the endpoint-to-location data migration, which is not yet supported. The feature carries a **Restart Recommended** tag for the Classic UI / API reason above. diff --git a/docs/content/asset_modelling/locations/PRO__migrating_from_endpoints.md b/docs/content/asset_modelling/locations/PRO__migrating_from_endpoints.md index bfaf2ca2d0..384c521286 100644 --- a/docs/content/asset_modelling/locations/PRO__migrating_from_endpoints.md +++ b/docs/content/asset_modelling/locations/PRO__migrating_from_endpoints.md @@ -15,6 +15,8 @@ Enabling Locations only changes behaviour for *new* imports; your existing histo When a backfill finishes it reports how many source objects it processed and how many distinct **Locations** those objects resolved to. The two numbers differ by design: several source objects can share one Location (many Endpoints normalising to the same URL, or many Findings sharing one component), so the Location count is normally lower than the object count. If any individual object could not be migrated it is skipped rather than aborting the run, and the number skipped is shown alongside the result. +A running item shows a **Cancel** button. Cancelling stops the run at the next batch boundary, so it is not instant: the current batch finishes and commits first. A cancelled run keeps everything it had already migrated, is reported as **Cancelled** with its partial counts, and because every step is idempotent, running the same item again resumes from where it stopped and converges on the same result as an uninterrupted run. Cancel is also the recovery path when a run's worker is lost: a run that stops reporting progress is marked failed on its own so the item becomes runnable again, and forcing a cancel releases a run that is otherwise wedged. + The suite has four items, because a Finding can carry three independent kinds of location: 1. **Endpoints to Locations backfill** — turns existing `Endpoint` rows into **URL Locations** (detailed below). diff --git a/docs/content/releases/pro/changelog.md b/docs/content/releases/pro/changelog.md index 206ca4785d..4e37ed6516 100644 --- a/docs/content/releases/pro/changelog.md +++ b/docs/content/releases/pro/changelog.md @@ -13,6 +13,11 @@ For Open Source release notes, please see the [Releases page on GitHub](https:// ## August 2026: v3.2 +### August 24, 2026: v3.2.300 + +Enhancements: +* **(Locations)** The data-migration suite on the Feature Flags page can now be cancelled while a backfill is running. Cancelling stops the run at the next batch boundary and keeps everything migrated so far, so re-running the item resumes and converges on the same result. A run whose worker is lost is now detected and marked failed on its own, so a stuck suite becomes runnable again instead of blocking every item. + ### August 17, 2026: v3.2.200 Enhancements: diff --git a/dojo/management/commands/migrate_endpoints_to_locations.py b/dojo/management/commands/migrate_endpoints_to_locations.py index c54040baf0..ea89b4e346 100644 --- a/dojo/management/commands/migrate_endpoints_to_locations.py +++ b/dojo/management/commands/migrate_endpoints_to_locations.py @@ -177,12 +177,14 @@ class Command(BaseCommand): help = "Usage: manage.py migrate_endpoints_to_locations" - # `progress_callback` and `summary_callback` are accepted by handle() but not - # exposed on the parser: they are programmatic hooks for in-process callers (Pro's - # migration suite passes them via call_command) and cannot be supplied on the - # command line. progress_callback(processed, total) drives a live bar; - # summary_callback(dict) reports the final {processed, total, locations, failures}. - stealth_options = ("progress_callback", "summary_callback") + # `progress_callback`, `summary_callback` and `cancel_callback` are accepted by + # handle() but not exposed on the parser: they are programmatic hooks for in-process + # callers (Pro's migration suite passes them via call_command) and cannot be supplied + # on the command line. progress_callback(processed, total) drives a live bar; + # summary_callback(dict) reports the final {processed, total, locations, failures, + # cancelled}; cancel_callback() -> bool is polled at each chunk boundary and, when it + # returns True, stops the run cleanly between chunks (see _should_cancel). + stealth_options = ("progress_callback", "summary_callback", "cancel_callback") def add_arguments(self, parser): parser.add_argument( @@ -741,6 +743,28 @@ def _emit_summary(self, summary: dict) -> None: except Exception: logger.warning("summary_callback raised; continuing migration", exc_info=True) + def _should_cancel(self) -> bool: + """ + Ask an optional cancel probe whether the run should stop. + + Pro's migration suite passes ``cancel_callback`` (a stealth option) that returns + True once a superuser has requested cancellation. Polled at each chunk boundary, + where the DB is already in a consistent, resumable state. + + Unlike progress/summary, a raising probe is treated as "do not cancel" rather than + propagated: a transient DB error while checking the flag must not throw away hours + of migration work. A probe that can never succeed is handled out of band by the + suite's stale-run recovery (force cancel). CLI runs pass no callback and never stop. + """ + callback = self.cancel_callback + if callback is None: + return False + try: + return bool(callback()) + except Exception: + logger.warning("cancel_callback raised; continuing migration", exc_info=True) + return False + def _log_progress( self, i: int, @@ -840,6 +864,9 @@ def handle(self, *args, **options): # Optional programmatic hooks (see stealth_options). None for CLI runs. self.progress_callback = options.get("progress_callback") self.summary_callback = options.get("summary_callback") + self.cancel_callback = options.get("cancel_callback") + # Set when the cancel probe asks us to stop, so the summary can report it. + self.cancelled = False # Per-phase wall-clock accumulators. self.timings = dict.fromkeys(PHASES, 0.0) @@ -955,6 +982,19 @@ def handle(self, *args, **options): # of the throttled stdout line below), for the migration suite. self._emit_progress(i, endpoint_count) + # Cancellation checkpoint. Chunks are separately committed and the + # tag queue was just flushed, so stopping here leaves a consistent, + # resumable state that a later run converges from (the class is + # idempotent by design). Break rather than return: the tag + # inheritance pass below still needs to run for what we migrated. + if self._should_cancel(): + self.cancelled = True + self.stdout.write(self.style.WARNING( + f"Cancellation requested; stopping after {i:,}/{endpoint_count:,} " + f"endpoints at a chunk boundary. Re-run to finish.", + )) + break + # 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: @@ -972,8 +1012,9 @@ def handle(self, *args, **options): elapsed = time.time() - run_t0 successful = i - len(self.failed_endpoints) + lead = "Stopped early." if self.cancelled else "Done." self.stdout.write(self.style.SUCCESS( - f"Done. Migrated {successful:,}/{i:,} endpoints in {self._fmt_duration(elapsed)} " + f"{lead} Migrated {successful:,}/{i:,} endpoints in {self._fmt_duration(elapsed)} " f"({(i / elapsed if elapsed else 0):.2f} endpoints/sec).", )) if self.failed_endpoints: @@ -1004,12 +1045,15 @@ def handle(self, *args, **options): # Summary for programmatic callers (the Pro Locations migration suite). CLI runs # pass no summary_callback and are unaffected. ``locations`` is the distinct- # Location count the endpoints collapsed onto (<= processed); ``failures`` are the - # per-endpoint errors isolated during the run so the suite can surface "N skipped". + # per-endpoint errors isolated during the run so the suite can surface "N skipped"; + # ``cancelled`` tells the suite the run stopped early (processed < total) rather + # than finishing, so it can land the run on "cancelled" instead of "completed". self._emit_summary( { "processed": i, "total": endpoint_count, "locations": len(self.migrated_location_ids), "failures": [{"id": endpoint_id, "error": error} for endpoint_id, error in self.failed_endpoints], + "cancelled": self.cancelled, }, ) diff --git a/unittests/test_migrate_endpoints_to_locations.py b/unittests/test_migrate_endpoints_to_locations.py index b1ad679746..d4ccf8a555 100644 --- a/unittests/test_migrate_endpoints_to_locations.py +++ b/unittests/test_migrate_endpoints_to_locations.py @@ -548,3 +548,81 @@ def test_inheritance_signal_is_suppressed_during_the_main_loop(self): [tag.name for tag in location.inherited_tags.all()], ["product-inherited"], ) + + def test_cancel_callback_stops_between_chunks(self): + # Three endpoints, one per chunk. The probe asks to cancel at the first + # chunk boundary, so the run stops with processed < total and the summary + # reports it as cancelled rather than completed. What committed before the + # cancel is persisted. + self._make_endpoint_with_status("first.example.com", active=True) + self._make_endpoint_with_status("second.example.com", active=True) + self._make_endpoint_with_status("third.example.com", active=True) + + calls = [] + + def cancel(): + calls.append(1) + return True # cancel at the first boundary + + summaries = [] + self._run(batch_size=1, cancel_callback=cancel, summary_callback=summaries.append) + + self.assertEqual(len(summaries), 1) + summary = summaries[0] + self.assertTrue(summary["cancelled"]) + self.assertEqual(summary["total"], 3) + self.assertEqual(summary["processed"], 1) + # Only the chunk that committed before the cancel is persisted. + self.assertEqual(URL.objects.count(), 1) + self.assertEqual(LocationFindingReference.objects.count(), 1) + + def test_rerun_after_cancel_converges(self): + for host in ("a.example.com", "b.example.com", "c.example.com"): + self._make_endpoint_with_status(host, active=True) + + first = [] + self._run(batch_size=1, cancel_callback=lambda: True, summary_callback=first.append) + self.assertTrue(first[0]["cancelled"]) + self.assertEqual(first[0]["processed"], 1) + self.assertEqual(URL.objects.count(), 1) + + # Re-running with no cancel finishes the migration and creates no + # duplicate rows for the chunk that already migrated. + second = [] + self._run(summary_callback=second.append) + self.assertFalse(second[0]["cancelled"]) + self.assertEqual(second[0]["processed"], 3) + self.assertEqual(URL.objects.count(), 3) + self.assertEqual(LocationFindingReference.objects.count(), 3) + + def test_cancel_callback_that_raises_does_not_abort(self): + # A probe that raises (e.g. a transient DB error) must not stop the run: + # discarding a multi-hour migration over a failed flag check is worse than + # finishing it. The exception is logged and the run completes normally. + self._make_endpoint_with_status("resilient.example.com", active=True) + + def boom(): + msg = "db blip while checking cancel flag" + raise RuntimeError(msg) + + summaries = [] + with self.assertLogs( + "dojo.management.commands.migrate_endpoints_to_locations", + level="WARNING", + ): + self._run(batch_size=1, cancel_callback=boom, summary_callback=summaries.append) + + self.assertFalse(summaries[0]["cancelled"]) + self.assertEqual(summaries[0]["processed"], 1) + self.assertEqual(URL.objects.count(), 1) + + def test_no_cancel_callback_reports_not_cancelled(self): + # The CLI path passes no cancel_callback and behaves exactly as before, + # reporting cancelled=False so the suite lands the run on "completed". + self._make_endpoint_with_status("plain-run.example.com", active=True) + + summaries = [] + self._run(summary_callback=summaries.append) + + self.assertFalse(summaries[0]["cancelled"]) + self.assertEqual(summaries[0]["processed"], 1)