add clean up whitelist

api/services/clear_free_plan_expired_workflow_run_logs.py

@@ -44,6 +44,7 @@ class WorkflowRunCleanup:
 
         self.batch_size = batch_size
         self.billing_cache: dict[str, SubscriptionPlan | None] = {}
+        self._cleanup_whitelist: set[str] | None = None
         self.dry_run = dry_run
         self.free_plan_grace_period_days = dify_config.SANDBOX_EXPIRED_RECORDS_CLEAN_GRACEFUL_PERIOD
         self.workflow_run_repo: APIWorkflowRunRepository
@@ -169,6 +170,8 @@ class WorkflowRunCleanup:
         if not tenant_id_list:
             return set()
 
+        cleanup_whitelist = self._get_cleanup_whitelist()
+
         uncached_tenants = [tenant_id for tenant_id in tenant_id_list if tenant_id not in self.billing_cache]
 
         if uncached_tenants:
@@ -186,6 +189,9 @@ class WorkflowRunCleanup:
 
         eligible_free_tenants: set[str] = set()
         for tenant_id in tenant_id_list:
+            if tenant_id in cleanup_whitelist:
+                continue
+
             info = self.billing_cache.get(tenant_id)
             if not info:
                 continue
@@ -222,6 +228,23 @@ class WorkflowRunCleanup:
         grace_deadline = expiration_at + datetime.timedelta(days=self.free_plan_grace_period_days)
         return datetime.datetime.now(datetime.UTC) < grace_deadline
 
+    def _get_cleanup_whitelist(self) -> set[str]:
+        if self._cleanup_whitelist is not None:
+            return self._cleanup_whitelist
+
+        if not dify_config.BILLING_ENABLED:
+            self._cleanup_whitelist = set()
+            return self._cleanup_whitelist
+
+        try:
+            whitelist_ids = BillingService.get_expired_subscription_cleanup_whitelist()
+        except Exception:
+            logger.exception("Failed to fetch cleanup whitelist from billing service")
+            whitelist_ids = []
+
+        self._cleanup_whitelist = set(whitelist_ids)
+        return self._cleanup_whitelist
+
     def _delete_trigger_logs(self, session: Session, run_ids: Sequence[str]) -> int:
         trigger_repo = SQLAlchemyWorkflowTriggerLogRepository(session)
         return trigger_repo.delete_by_run_ids(run_ids)

api/tests/unit_tests/services/test_clear_free_plan_expired_workflow_run_logs.py

@@ -72,6 +72,7 @@ def create_cleanup(
     repo: FakeRepo,
     *,
     grace_period_days: int = 0,
+    whitelist: set[str] | None = None,
     **kwargs: Any,
 ) -> WorkflowRunCleanup:
     monkeypatch.setattr(
@@ -79,6 +80,11 @@ def create_cleanup(
         "SANDBOX_EXPIRED_RECORDS_CLEAN_GRACEFUL_PERIOD",
         grace_period_days,
     )
+    monkeypatch.setattr(
+        cleanup_module.WorkflowRunCleanup,
+        "_get_cleanup_whitelist",
+        lambda self: whitelist or set(),
+    )
     return WorkflowRunCleanup(workflow_run_repo=repo, **kwargs)
 
 
@@ -136,6 +142,29 @@ def test_filter_free_tenants_respects_grace_period(monkeypatch: pytest.MonkeyPat
     assert free == {"long_sandbox"}
 
 
+def test_filter_free_tenants_skips_cleanup_whitelist(monkeypatch: pytest.MonkeyPatch) -> None:
+    cleanup = create_cleanup(
+        monkeypatch,
+        repo=FakeRepo([]),
+        days=30,
+        batch_size=10,
+        whitelist={"tenant_whitelist"},
+    )
+
+    monkeypatch.setattr(cleanup_module.dify_config, "BILLING_ENABLED", True)
+    cleanup.billing_cache["tenant_whitelist"] = plan_info("sandbox", -1)
+    monkeypatch.setattr(
+        cleanup_module.BillingService,
+        "get_plan_bulk",
+        staticmethod(lambda tenant_ids: {tenant_id: plan_info("sandbox", -1) for tenant_id in tenant_ids}),
+    )
+
+    tenants = {"tenant_whitelist", "tenant_regular"}
+    free = cleanup._filter_free_tenants(tenants)
+
+    assert free == {"tenant_regular"}
+
+
 def test_filter_free_tenants_bulk_failure(monkeypatch: pytest.MonkeyPatch) -> None:
     cleanup = create_cleanup(monkeypatch, repo=FakeRepo([]), days=30, batch_size=10)