From 5f957a115aecc875ea3cf1c5d76633d46c336dbc Mon Sep 17 00:00:00 2001 From: hjlarry Date: Mon, 15 Dec 2025 16:17:05 +0800 Subject: [PATCH] add dry run param for clear command --- api/commands.py | 7 +++ ...ear_free_plan_expired_workflow_run_logs.py | 53 +++++++++++++++---- ...ear_free_plan_expired_workflow_run_logs.py | 15 ++++++ 3 files changed, 65 insertions(+), 10 deletions(-) diff --git a/api/commands.py b/api/commands.py index 9a990459c0..16554dbedc 100644 --- a/api/commands.py +++ b/api/commands.py @@ -869,11 +869,17 @@ def clear_free_plan_tenant_expired_logs(days: int, batch: int, tenant_ids: list[ default=None, help="Optional upper bound (exclusive) for created_at; must be paired with --start-after.", ) +@click.option( + "--dry-run", + is_flag=True, + help="Preview cleanup results without deleting any workflow run data.", +) def clean_workflow_runs( days: int, batch_size: int, start_after: datetime.datetime | None, end_before: datetime.datetime | None, + dry_run: bool, ): """ Clean workflow runs and related workflow data for free tenants. @@ -888,6 +894,7 @@ def clean_workflow_runs( batch_size=batch_size, start_after=start_after, end_before=end_before, + dry_run=dry_run, ).run() click.echo(click.style("Workflow run cleanup completed.", fg="green")) diff --git a/api/services/clear_free_plan_expired_workflow_run_logs.py b/api/services/clear_free_plan_expired_workflow_run_logs.py index 1fe2bad2d0..b5e5f40823 100644 --- a/api/services/clear_free_plan_expired_workflow_run_logs.py +++ b/api/services/clear_free_plan_expired_workflow_run_logs.py @@ -27,6 +27,7 @@ class WorkflowRunCleanup: start_after: datetime.datetime | None = None, end_before: datetime.datetime | None = None, workflow_run_repo: APIWorkflowRunRepository | None = None, + dry_run: bool = False, ): if (start_after is None) ^ (end_before is None): raise ValueError("start_after and end_before must be both set or both omitted.") @@ -40,6 +41,7 @@ class WorkflowRunCleanup: self.batch_size = batch_size self.billing_cache: dict[str, CloudPlan | None] = {} + self.dry_run = dry_run self.workflow_run_repo: APIWorkflowRunRepository if workflow_run_repo: self.workflow_run_repo = workflow_run_repo @@ -53,14 +55,17 @@ class WorkflowRunCleanup: def run(self) -> None: click.echo( click.style( - f"Cleaning workflow runs " + f"{'Inspecting' if self.dry_run else 'Cleaning'} workflow runs " f"{'between ' + self.window_start.isoformat() + ' and ' if self.window_start else 'before '}" f"{self.window_end.isoformat()} (batch={self.batch_size})", fg="white", ) ) + if self.dry_run: + click.echo(click.style("Dry run mode enabled. No data will be deleted.", fg="yellow")) total_runs_deleted = 0 + total_runs_targeted = 0 batch_index = 0 last_seen: tuple[datetime.datetime, str] | None = None @@ -90,6 +95,19 @@ class WorkflowRunCleanup: ) continue + total_runs_targeted += len(free_runs) + + if self.dry_run: + sample_ids = ", ".join(run.id for run in free_runs[:5]) + click.echo( + click.style( + f"[batch #{batch_index}] would delete {len(free_runs)} runs " + f"(sample ids: {sample_ids}) and skip {paid_or_skipped} paid/unknown", + fg="yellow", + ) + ) + continue + try: counts = self.workflow_run_repo.delete_runs_with_related( free_runs, @@ -112,17 +130,32 @@ class WorkflowRunCleanup: ) ) - if self.window_start: - summary_message = ( - f"Cleanup complete. Deleted {total_runs_deleted} workflow runs " - f"between {self.window_start.isoformat()} and {self.window_end.isoformat()}" - ) + if self.dry_run: + if self.window_start: + summary_message = ( + f"Dry run complete. Would delete {total_runs_targeted} workflow runs " + f"between {self.window_start.isoformat()} and {self.window_end.isoformat()}" + ) + else: + summary_message = ( + f"Dry run complete. Would delete {total_runs_targeted} workflow runs " + f"before {self.window_end.isoformat()}" + ) + summary_color = "yellow" else: - summary_message = ( - f"Cleanup complete. Deleted {total_runs_deleted} workflow runs before {self.window_end.isoformat()}" - ) + if self.window_start: + summary_message = ( + f"Cleanup complete. Deleted {total_runs_deleted} workflow runs " + f"between {self.window_start.isoformat()} and {self.window_end.isoformat()}" + ) + else: + summary_message = ( + f"Cleanup complete. Deleted {total_runs_deleted} workflow runs " + f"before {self.window_end.isoformat()}" + ) + summary_color = "white" - click.echo(click.style(summary_message, fg="white")) + click.echo(click.style(summary_message, fg=summary_color)) def _filter_free_tenants(self, tenant_ids: Iterable[str]) -> set[str]: if not dify_config.BILLING_ENABLED: diff --git a/api/tests/unit_tests/services/test_clear_free_plan_expired_workflow_run_logs.py b/api/tests/unit_tests/services/test_clear_free_plan_expired_workflow_run_logs.py index a1685fcfb0..24a07e7937 100644 --- a/api/tests/unit_tests/services/test_clear_free_plan_expired_workflow_run_logs.py +++ b/api/tests/unit_tests/services/test_clear_free_plan_expired_workflow_run_logs.py @@ -163,6 +163,21 @@ def test_run_exits_on_empty_batch(monkeypatch: pytest.MonkeyPatch) -> None: cleanup.run() +def test_run_dry_run_skips_deletions(monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]) -> None: + cutoff = datetime.datetime.now() + repo = FakeRepo(batches=[[FakeRun("run-free", "t_free", cutoff)]]) + cleanup = create_cleanup(monkeypatch, repo=repo, days=30, batch_size=10, dry_run=True) + + monkeypatch.setattr(cleanup_module.dify_config, "BILLING_ENABLED", False) + + cleanup.run() + + assert repo.deleted == [] + captured = capsys.readouterr().out + assert "Dry run mode enabled" in captured + assert "would delete 1 runs" in captured + + def test_between_sets_window_bounds(monkeypatch: pytest.MonkeyPatch) -> None: start_after = datetime.datetime(2024, 5, 1, 0, 0, 0) end_before = datetime.datetime(2024, 6, 1, 0, 0, 0)