From e14a5797d44df93a66646715b68794a73f90feee Mon Sep 17 00:00:00 2001 From: hjlarry Date: Thu, 11 Dec 2025 12:02:42 +0800 Subject: [PATCH 01/42] add workflow run clean up --- api/commands.py | 41 +++ api/configs/feature/__init__.py | 5 + api/extensions/ext_celery.py | 7 + api/extensions/ext_commands.py | 2 + ...d7c23e_add_workflow_runs_created_at_idx.py | 29 +++ api/schedule/clean_workflow_runs_task.py | 30 +++ api/services/billing_service.py | 22 ++ ...ear_free_plan_expired_workflow_run_logs.py | 235 ++++++++++++++++++ ...ear_free_plan_expired_workflow_run_logs.py | 231 +++++++++++++++++ 9 files changed, 602 insertions(+) create mode 100644 api/migrations/versions/2025_12_10_1504-8a7f2ad7c23e_add_workflow_runs_created_at_idx.py create mode 100644 api/schedule/clean_workflow_runs_task.py create mode 100644 api/services/clear_free_plan_expired_workflow_run_logs.py create mode 100644 api/tests/unit_tests/services/test_clear_free_plan_expired_workflow_run_logs.py diff --git a/api/commands.py b/api/commands.py index a8d89ac200..cdf3a09bb7 100644 --- a/api/commands.py +++ b/api/commands.py @@ -1,4 +1,5 @@ import base64 +import datetime import json import logging import secrets @@ -41,6 +42,7 @@ from models.provider_ids import DatasourceProviderID, ToolProviderID from models.source import DataSourceApiKeyAuthBinding, DataSourceOauthBinding from models.tools import ToolOAuthSystemClient from services.account_service import AccountService, RegisterService, TenantService +from services.clear_free_plan_expired_workflow_run_logs import WorkflowRunCleanup from services.clear_free_plan_tenant_expired_logs import ClearFreePlanTenantExpiredLogs from services.plugin.data_migration import PluginDataMigration from services.plugin.plugin_migration import PluginMigration @@ -852,6 +854,45 @@ def clear_free_plan_tenant_expired_logs(days: int, batch: int, tenant_ids: list[ click.echo(click.style("Clear free plan tenant expired logs completed.", fg="green")) +@click.command("clean-workflow-runs", help="Clean expired workflow runs and related data for free tenants.") +@click.option("--days", default=30, show_default=True, help="Delete workflow runs created before N days ago.") +@click.option("--batch-size", default=1000, show_default=True, help="Batch size for selecting workflow runs.") +@click.option( + "--start-after", + type=click.DateTime(formats=["%Y-%m-%d", "%Y-%m-%dT%H:%M:%S"]), + default=None, + help="Optional lower bound (inclusive) for created_at; must be paired with --end-before.", +) +@click.option( + "--end-before", + type=click.DateTime(formats=["%Y-%m-%d", "%Y-%m-%dT%H:%M:%S"]), + default=None, + help="Optional upper bound (exclusive) for created_at; must be paired with --start-after.", +) +def clean_workflow_runs( + days: int, + batch_size: int, + start_after: datetime.datetime | None, + end_before: datetime.datetime | None, +): + """ + Clean workflow runs and related workflow data for free tenants. + """ + if (start_after is None) ^ (end_before is None): + raise click.UsageError("--start-after and --end-before must be provided together.") + + click.echo(click.style("Starting workflow run cleanup.", fg="white")) + + WorkflowRunCleanup( + days=days, + batch_size=batch_size, + start_after=start_after, + end_before=end_before, + ).run() + + click.echo(click.style("Workflow run cleanup completed.", fg="green")) + + @click.option("-f", "--force", is_flag=True, help="Skip user confirmation and force the command to execute.") @click.command("clear-orphaned-file-records", help="Clear orphaned file records.") def clear_orphaned_file_records(force: bool): diff --git a/api/configs/feature/__init__.py b/api/configs/feature/__init__.py index a5916241df..e237470709 100644 --- a/api/configs/feature/__init__.py +++ b/api/configs/feature/__init__.py @@ -1116,6 +1116,11 @@ class CeleryScheduleTasksConfig(BaseSettings): default=60 * 60, ) + ENABLE_WORKFLOW_RUN_CLEANUP_TASK: bool = Field( + description="Enable scheduled workflow run cleanup task", + default=False, + ) + class PositionConfig(BaseSettings): POSITION_PROVIDER_PINS: str = Field( diff --git a/api/extensions/ext_celery.py b/api/extensions/ext_celery.py index 5cf4984709..c19763372d 100644 --- a/api/extensions/ext_celery.py +++ b/api/extensions/ext_celery.py @@ -160,6 +160,13 @@ def init_app(app: DifyApp) -> Celery: "task": "schedule.clean_workflow_runlogs_precise.clean_workflow_runlogs_precise", "schedule": crontab(minute="0", hour="2"), } + if dify_config.ENABLE_WORKFLOW_RUN_CLEANUP_TASK: + # for saas only + imports.append("schedule.clean_workflow_runs_task") + beat_schedule["clean_workflow_runs_task"] = { + "task": "schedule.clean_workflow_runs_task.clean_workflow_runs_task", + "schedule": crontab(minute="0", hour="0"), + } if dify_config.ENABLE_WORKFLOW_SCHEDULE_POLLER_TASK: imports.append("schedule.workflow_schedule_task") beat_schedule["workflow_schedule_task"] = { diff --git a/api/extensions/ext_commands.py b/api/extensions/ext_commands.py index 71a63168a5..6f6322827c 100644 --- a/api/extensions/ext_commands.py +++ b/api/extensions/ext_commands.py @@ -4,6 +4,7 @@ from dify_app import DifyApp def init_app(app: DifyApp): from commands import ( add_qdrant_index, + clean_workflow_runs, cleanup_orphaned_draft_variables, clear_free_plan_tenant_expired_logs, clear_orphaned_file_records, @@ -54,6 +55,7 @@ def init_app(app: DifyApp): setup_datasource_oauth_client, transform_datasource_credentials, install_rag_pipeline_plugins, + clean_workflow_runs, ] for cmd in cmds_to_register: app.cli.add_command(cmd) diff --git a/api/migrations/versions/2025_12_10_1504-8a7f2ad7c23e_add_workflow_runs_created_at_idx.py b/api/migrations/versions/2025_12_10_1504-8a7f2ad7c23e_add_workflow_runs_created_at_idx.py new file mode 100644 index 0000000000..7968429ca8 --- /dev/null +++ b/api/migrations/versions/2025_12_10_1504-8a7f2ad7c23e_add_workflow_runs_created_at_idx.py @@ -0,0 +1,29 @@ +"""Add index on workflow_runs.created_at + +Revision ID: 8a7f2ad7c23e +Revises: d57accd375ae +Create Date: 2025-12-10 15:04:00.000000 +""" + +from alembic import op +import sqlalchemy as sa + +# revision identifiers, used by Alembic. +revision = "8a7f2ad7c23e" +down_revision = "d57accd375ae" +branch_labels = None +depends_on = None + + +def upgrade(): + with op.batch_alter_table("workflow_runs", schema=None) as batch_op: + batch_op.create_index( + batch_op.f("workflow_runs_created_at_idx"), + ["created_at"], + unique=False, + ) + + +def downgrade(): + with op.batch_alter_table("workflow_runs", schema=None) as batch_op: + batch_op.drop_index(batch_op.f("workflow_runs_created_at_idx")) diff --git a/api/schedule/clean_workflow_runs_task.py b/api/schedule/clean_workflow_runs_task.py new file mode 100644 index 0000000000..b59dc7f823 --- /dev/null +++ b/api/schedule/clean_workflow_runs_task.py @@ -0,0 +1,30 @@ +import click + +import app +from configs import dify_config +from services.clear_free_plan_expired_workflow_run_logs import WorkflowRunCleanup + +CLEANUP_QUEUE = "retention" + + +@app.celery.task(queue=CLEANUP_QUEUE) +def clean_workflow_runs_task() -> None: + """ + Scheduled cleanup for workflow runs and related records (sandbox tenants only). + """ + click.echo( + click.style( + f"Scheduled workflow run cleanup starting: cutoff={dify_config.WORKFLOW_LOG_RETENTION_DAYS} days, " + f"batch={dify_config.WORKFLOW_LOG_CLEANUP_BATCH_SIZE}", + fg="green", + ) + ) + + WorkflowRunCleanup( + days=dify_config.WORKFLOW_LOG_RETENTION_DAYS, + batch_size=dify_config.WORKFLOW_LOG_CLEANUP_BATCH_SIZE, + start_after=None, + end_before=None, + ).run() + + click.echo(click.style("Scheduled workflow run cleanup finished.", fg="green")) diff --git a/api/services/billing_service.py b/api/services/billing_service.py index 54e1c9d285..4fc61c6c0f 100644 --- a/api/services/billing_service.py +++ b/api/services/billing_service.py @@ -1,3 +1,4 @@ +import logging import os from typing import Literal @@ -11,6 +12,8 @@ from extensions.ext_redis import redis_client from libs.helper import RateLimiter from models import Account, TenantAccountJoin, TenantAccountRole +logger = logging.getLogger(__name__) + class BillingService: base_url = os.environ.get("BILLING_API_URL", "BILLING_API_URL") @@ -25,6 +28,25 @@ class BillingService: billing_info = cls._send_request("GET", "/subscription/info", params=params) return billing_info + @classmethod + def get_info_bulk(cls, tenant_ids: list[str]) -> dict[str, dict]: + """ + Temporary bulk billing info fetch. Will be replaced by a real batch API. + + Args: + tenant_ids: list of tenant ids + + Returns: + Mapping of tenant_id -> billing info dict + """ + result: dict[str, dict] = {} + for tenant_id in tenant_ids: + try: + result[tenant_id] = cls.get_info(tenant_id) + except Exception: + logger.exception("Failed to fetch billing info for tenant %s in bulk mode", tenant_id) + return result + @classmethod def get_tenant_feature_plan_usage_info(cls, tenant_id: str): params = {"tenant_id": tenant_id} diff --git a/api/services/clear_free_plan_expired_workflow_run_logs.py b/api/services/clear_free_plan_expired_workflow_run_logs.py new file mode 100644 index 0000000000..98543c2928 --- /dev/null +++ b/api/services/clear_free_plan_expired_workflow_run_logs.py @@ -0,0 +1,235 @@ +import datetime +import logging +from collections.abc import Iterable, Sequence +from dataclasses import dataclass + +import click +import sqlalchemy as sa +from sqlalchemy import select +from sqlalchemy.orm import Session + +from configs import dify_config +from enums.cloud_plan import CloudPlan +from extensions.ext_database import db +from models import WorkflowAppLog, WorkflowNodeExecutionModel, WorkflowRun +from models.trigger import WorkflowTriggerLog +from models.workflow import WorkflowNodeExecutionOffload, WorkflowPause, WorkflowPauseReason +from services.billing_service import BillingService + +logger = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class WorkflowRunRow: + id: str + tenant_id: str + created_at: datetime.datetime + + +class WorkflowRunCleanup: + def __init__( + self, + days: int, + batch_size: int, + start_after: datetime.datetime | None = None, + end_before: datetime.datetime | None = None, + ): + if (start_after is None) ^ (end_before is None): + raise ValueError("start_after and end_before must be both set or both omitted.") + + computed_cutoff = datetime.datetime.now() - datetime.timedelta(days=days) + self.window_start = start_after + self.window_end = end_before or computed_cutoff + + if self.window_start and self.window_end <= self.window_start: + raise ValueError("end_before must be greater than start_after.") + + self.batch_size = batch_size + self.billing_cache: dict[str, CloudPlan | None] = {} + + def run(self) -> None: + click.echo( + click.style( + f"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", + ) + ) + + total_runs_deleted = 0 + batch_index = 0 + last_seen: tuple[datetime.datetime, str] | None = None + + while True: + with Session(db.engine) as session: + run_rows = self._load_batch(session, last_seen) + if not run_rows: + break + + batch_index += 1 + last_seen = (run_rows[-1].created_at, run_rows[-1].id) + tenant_ids = {row.tenant_id for row in run_rows} + free_tenants = self._filter_free_tenants(tenant_ids) + free_run_ids = [row.id for row in run_rows if row.tenant_id in free_tenants] + paid_or_skipped = len(run_rows) - len(free_run_ids) + + if not free_run_ids: + click.echo( + click.style( + f"[batch #{batch_index}] skipped (no sandbox runs in batch, {paid_or_skipped} paid)", + fg="yellow", + ) + ) + continue + + try: + counts = self._delete_runs(session, free_run_ids) + session.commit() + except Exception: + session.rollback() + logger.exception("Failed to delete workflow runs batch ending at %s", last_seen[0]) + raise + + total_runs_deleted += counts["runs"] + click.echo( + click.style( + f"[batch #{batch_index}] deleted runs: {counts['runs']} " + f"(nodes {counts['node_executions']}, offloads {counts['offloads']}, " + f"app_logs {counts['app_logs']}, trigger_logs {counts['trigger_logs']}, " + f"pauses {counts['pauses']}, pause_reasons {counts['pause_reasons']}); " + f"skipped {paid_or_skipped} paid/unknown", + fg="green", + ) + ) + + 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 before {self.window_end.isoformat()}" + ) + + click.echo(click.style(summary_message, fg="white")) + + def _load_batch( + self, session: Session, last_seen: tuple[datetime.datetime, str] | None + ) -> list[WorkflowRunRow]: + stmt = ( + select(WorkflowRun.id, WorkflowRun.tenant_id, WorkflowRun.created_at) + .where(WorkflowRun.created_at < self.window_end) + .order_by(WorkflowRun.created_at.asc(), WorkflowRun.id.asc()) + .limit(self.batch_size) + ) + + if self.window_start: + stmt = stmt.where(WorkflowRun.created_at >= self.window_start) + + if last_seen: + stmt = stmt.where( + sa.or_( + WorkflowRun.created_at > last_seen[0], + sa.and_(WorkflowRun.created_at == last_seen[0], WorkflowRun.id > last_seen[1]), + ) + ) + + rows = session.execute(stmt).all() + return [WorkflowRunRow(id=row.id, tenant_id=row.tenant_id, created_at=row.created_at) for row in rows] + + def _filter_free_tenants(self, tenant_ids: Iterable[str]) -> set[str]: + if not dify_config.BILLING_ENABLED: + return set(tenant_ids) + + tenant_id_list = list(tenant_ids) + uncached_tenants = [tenant_id for tenant_id in tenant_id_list if tenant_id not in self.billing_cache] + + if uncached_tenants: + try: + bulk_info = BillingService.get_info_bulk(uncached_tenants) + except Exception: + bulk_info = {} + logger.exception("Failed to fetch billing plans in bulk for tenants: %s", uncached_tenants) + + for tenant_id in uncached_tenants: + plan: CloudPlan | None = None + info = bulk_info.get(tenant_id) + if info: + try: + raw_plan = info.get("subscription", {}).get("plan") + plan = CloudPlan(raw_plan) + except Exception: + logger.exception("Failed to parse billing plan for tenant %s", tenant_id) + else: + logger.warning("Missing billing info for tenant %s in bulk resp; treating as non-free", tenant_id) + + self.billing_cache[tenant_id] = plan + + return {tenant_id for tenant_id in tenant_id_list if self.billing_cache.get(tenant_id) == CloudPlan.SANDBOX} + + def _delete_runs(self, session: Session, workflow_run_ids: Sequence[str]) -> dict[str, int]: + node_execution_ids = session.scalars( + select(WorkflowNodeExecutionModel.id).where(WorkflowNodeExecutionModel.workflow_run_id.in_(workflow_run_ids)) + ).all() + + offloads_deleted = 0 + if node_execution_ids: + offloads_deleted = ( + session.query(WorkflowNodeExecutionOffload) + .where(WorkflowNodeExecutionOffload.node_execution_id.in_(node_execution_ids)) + .delete(synchronize_session=False) + ) + + node_executions_deleted = 0 + if node_execution_ids: + node_executions_deleted = ( + session.query(WorkflowNodeExecutionModel) + .where(WorkflowNodeExecutionModel.id.in_(node_execution_ids)) + .delete(synchronize_session=False) + ) + + app_logs_deleted = ( + session.query(WorkflowAppLog) + .where(WorkflowAppLog.workflow_run_id.in_(workflow_run_ids)) + .delete(synchronize_session=False) + ) + + pause_ids = session.scalars( + select(WorkflowPause.id).where(WorkflowPause.workflow_run_id.in_(workflow_run_ids)) + ).all() + pause_reasons_deleted = 0 + pauses_deleted = 0 + + if pause_ids: + pause_reasons_deleted = ( + session.query(WorkflowPauseReason).where(WorkflowPauseReason.pause_id.in_(pause_ids)).delete( + synchronize_session=False + ) + ) + pauses_deleted = ( + session.query(WorkflowPause) + .where(WorkflowPause.id.in_(pause_ids)) + .delete(synchronize_session=False) + ) + + trigger_logs_deleted = ( + session.query(WorkflowTriggerLog) + .where(WorkflowTriggerLog.workflow_run_id.in_(workflow_run_ids)) + .delete(synchronize_session=False) + ) + + runs_deleted = ( + session.query(WorkflowRun).where(WorkflowRun.id.in_(workflow_run_ids)).delete(synchronize_session=False) + ) + + return { + "runs": runs_deleted, + "node_executions": node_executions_deleted, + "offloads": offloads_deleted, + "app_logs": app_logs_deleted, + "trigger_logs": trigger_logs_deleted, + "pauses": pauses_deleted, + "pause_reasons": pause_reasons_deleted, + } 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 new file mode 100644 index 0000000000..fd756bcddd --- /dev/null +++ b/api/tests/unit_tests/services/test_clear_free_plan_expired_workflow_run_logs.py @@ -0,0 +1,231 @@ +import datetime +from typing import Any + +import pytest + +from services import clear_free_plan_expired_workflow_run_logs as cleanup_module +from services.clear_free_plan_expired_workflow_run_logs import WorkflowRunCleanup, WorkflowRunRow + + +class DummySession: + def __init__(self) -> None: + self.committed = False + + def __enter__(self) -> "DummySession": + return self + + def __exit__(self, exc_type: object, exc: object, tb: object) -> None: + return None + + def commit(self) -> None: + self.committed = True + + +def test_filter_free_tenants_billing_disabled(monkeypatch: pytest.MonkeyPatch) -> None: + cleanup = WorkflowRunCleanup(days=30, batch_size=10) + + monkeypatch.setattr(cleanup_module.dify_config, "BILLING_ENABLED", False) + + def fail_bulk(_: list[str]) -> dict[str, dict[str, Any]]: + raise RuntimeError("should not call") + + monkeypatch.setattr(cleanup_module.BillingService, "get_info_bulk", staticmethod(fail_bulk)) + + tenants = {"t1", "t2"} + free = cleanup._filter_free_tenants(tenants) + + assert free == tenants + + +def test_filter_free_tenants_bulk_mixed(monkeypatch: pytest.MonkeyPatch) -> None: + cleanup = WorkflowRunCleanup(days=30, batch_size=10) + + monkeypatch.setattr(cleanup_module.dify_config, "BILLING_ENABLED", True) + # seed cache to avoid relying on billing service implementation + cleanup.billing_cache["t_free"] = cleanup_module.CloudPlan.SANDBOX + cleanup.billing_cache["t_paid"] = cleanup_module.CloudPlan.TEAM + monkeypatch.setattr( + cleanup_module.BillingService, + "get_info_bulk", + staticmethod(lambda tenant_ids: {tenant_id: {} for tenant_id in tenant_ids}), + ) + + free = cleanup._filter_free_tenants({"t_free", "t_paid", "t_missing"}) + + assert free == {"t_free"} + + +def test_filter_free_tenants_bulk_failure(monkeypatch: pytest.MonkeyPatch) -> None: + cleanup = WorkflowRunCleanup(days=30, batch_size=10) + + monkeypatch.setattr(cleanup_module.dify_config, "BILLING_ENABLED", True) + monkeypatch.setattr( + cleanup_module.BillingService, + "get_info_bulk", + staticmethod(lambda tenant_ids: (_ for _ in ()).throw(RuntimeError("boom"))), + ) + + free = cleanup._filter_free_tenants({"t1", "t2"}) + + assert free == set() + + +def test_run_deletes_only_free_tenants(monkeypatch: pytest.MonkeyPatch) -> None: + cutoff = datetime.datetime.now() + cleanup = WorkflowRunCleanup(days=30, batch_size=10) + + monkeypatch.setattr(cleanup_module.dify_config, "BILLING_ENABLED", True) + cleanup.billing_cache["t_free"] = cleanup_module.CloudPlan.SANDBOX + cleanup.billing_cache["t_paid"] = cleanup_module.CloudPlan.TEAM + monkeypatch.setattr( + cleanup_module.BillingService, + "get_info_bulk", + staticmethod(lambda tenant_ids: {tenant_id: {} for tenant_id in tenant_ids}), + ) + + batches_returned = 0 + + def fake_load_batch( + session: DummySession, last_seen: tuple[datetime.datetime, str] | None + ) -> list[WorkflowRunRow]: + nonlocal batches_returned + if batches_returned > 0: + return [] + batches_returned += 1 + return [ + WorkflowRunRow(id="run-free", tenant_id="t_free", created_at=cutoff), + WorkflowRunRow(id="run-paid", tenant_id="t_paid", created_at=cutoff), + ] + + deleted_ids: list[list[str]] = [] + + def fake_delete_runs(session: DummySession, workflow_run_ids: list[str]) -> dict[str, int]: + deleted_ids.append(list(workflow_run_ids)) + return { + "runs": len(workflow_run_ids), + "node_executions": 0, + "offloads": 0, + "app_logs": 0, + "trigger_logs": 0, + "pauses": 0, + "pause_reasons": 0, + } + + created_sessions: list[DummySession] = [] + + def fake_session_factory(engine: object | None = None) -> DummySession: + session = DummySession() + created_sessions.append(session) + return session + + monkeypatch.setattr(cleanup, "_load_batch", fake_load_batch) + monkeypatch.setattr(cleanup, "_delete_runs", fake_delete_runs) + monkeypatch.setattr(cleanup_module, "Session", fake_session_factory) + + class DummyDB: + engine: object | None = None + + monkeypatch.setattr(cleanup_module, "db", DummyDB()) + + cleanup.run() + + assert deleted_ids == [["run-free"]] + assert created_sessions + assert created_sessions[0].committed is True + + +def test_run_skips_when_no_free_tenants(monkeypatch: pytest.MonkeyPatch) -> None: + cutoff = datetime.datetime.now() + cleanup = WorkflowRunCleanup(days=30, batch_size=10) + + monkeypatch.setattr(cleanup_module.dify_config, "BILLING_ENABLED", True) + monkeypatch.setattr( + cleanup_module.BillingService, + "get_info_bulk", + staticmethod( + lambda tenant_ids: {tenant_id: {"subscription": {"plan": "TEAM"}} for tenant_id in tenant_ids} + ), + ) + + batches_returned = 0 + + def fake_load_batch( + session: DummySession, last_seen: tuple[datetime.datetime, str] | None + ) -> list[WorkflowRunRow]: + nonlocal batches_returned + if batches_returned > 0: + return [] + batches_returned += 1 + return [WorkflowRunRow(id="run-paid", tenant_id="t_paid", created_at=cutoff)] + + delete_called = False + + def fake_delete_runs(session: DummySession, workflow_run_ids: list[str]) -> dict[str, int]: + nonlocal delete_called + delete_called = True + return { + "runs": 0, + "node_executions": 0, + "offloads": 0, + "app_logs": 0, + "trigger_logs": 0, + "pauses": 0, + "pause_reasons": 0 + } + + def fake_session_factory(engine: object | None = None) -> DummySession: # pragma: no cover - simple factory + return DummySession() + + monkeypatch.setattr(cleanup, "_load_batch", fake_load_batch) + monkeypatch.setattr(cleanup, "_delete_runs", fake_delete_runs) + monkeypatch.setattr(cleanup_module, "Session", fake_session_factory) + monkeypatch.setattr(cleanup_module, "db", type("DummyDB", (), {"engine": None})) + + cleanup.run() + + assert delete_called is False + + +def test_run_exits_on_empty_batch(monkeypatch: pytest.MonkeyPatch) -> None: + cleanup = WorkflowRunCleanup(days=30, batch_size=10) + + def fake_load_batch( + session: DummySession, last_seen: tuple[datetime.datetime, str] | None + ) -> list[WorkflowRunRow]: + return [] + + def fake_delete_runs(session: DummySession, workflow_run_ids: list[str]) -> dict[str, int]: + raise AssertionError("should not delete") + + def fake_session_factory(engine: object | None = None) -> DummySession: # pragma: no cover - simple factory + return DummySession() + + monkeypatch.setattr(cleanup, "_load_batch", fake_load_batch) + monkeypatch.setattr(cleanup, "_delete_runs", fake_delete_runs) + monkeypatch.setattr(cleanup_module, "Session", fake_session_factory) + monkeypatch.setattr(cleanup_module, "db", type("DummyDB", (), {"engine": None})) + + cleanup.run() + + +def test_between_sets_window_bounds() -> None: + start_after = datetime.datetime(2024, 5, 1, 0, 0, 0) + end_before = datetime.datetime(2024, 6, 1, 0, 0, 0) + cleanup = WorkflowRunCleanup(days=30, batch_size=10, start_after=start_after, end_before=end_before) + + assert cleanup.window_start == start_after + assert cleanup.window_end == end_before + + +def test_between_requires_both_boundaries() -> None: + with pytest.raises(ValueError): + WorkflowRunCleanup(days=30, batch_size=10, start_after=datetime.datetime.now(), end_before=None) + with pytest.raises(ValueError): + WorkflowRunCleanup(days=30, batch_size=10, start_after=None, end_before=datetime.datetime.now()) + + +def test_between_requires_end_after_start() -> None: + start_after = datetime.datetime(2024, 6, 1, 0, 0, 0) + end_before = datetime.datetime(2024, 5, 1, 0, 0, 0) + with pytest.raises(ValueError): + WorkflowRunCleanup(days=30, batch_size=10, start_after=start_after, end_before=end_before) From d8578dd4ba77a719fa9de4ce339410b55ddd3bda Mon Sep 17 00:00:00 2001 From: hjlarry Date: Thu, 11 Dec 2025 15:15:29 +0800 Subject: [PATCH 02/42] add batch get plan api --- api/services/billing_service.py | 21 +++++++++---------- ...ear_free_plan_expired_workflow_run_logs.py | 7 +------ 2 files changed, 11 insertions(+), 17 deletions(-) diff --git a/api/services/billing_service.py b/api/services/billing_service.py index 4fc61c6c0f..6ffb1f0cb6 100644 --- a/api/services/billing_service.py +++ b/api/services/billing_service.py @@ -31,21 +31,20 @@ class BillingService: @classmethod def get_info_bulk(cls, tenant_ids: list[str]) -> dict[str, dict]: """ - Temporary bulk billing info fetch. Will be replaced by a real batch API. + Bulk billing info fetch via billing API. - Args: - tenant_ids: list of tenant ids + Payload: {"tenant_ids": ["t1", "t2", ...]} (max 200 per request) Returns: - Mapping of tenant_id -> billing info dict + Mapping of tenant_id -> plan """ - result: dict[str, dict] = {} - for tenant_id in tenant_ids: - try: - result[tenant_id] = cls.get_info(tenant_id) - except Exception: - logger.exception("Failed to fetch billing info for tenant %s in bulk mode", tenant_id) - return result + + try: + resp = cls._send_request("POST", "/subscription/plan/batch", json={"tenant_ids": tenant_ids}) + except Exception: + logger.exception("Failed to fetch billing info batch for tenants: %s", tenant_ids) + + return resp.get("data", {}) @classmethod def get_tenant_feature_plan_usage_info(cls, tenant_id: str): 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 98543c2928..5972440732 100644 --- a/api/services/clear_free_plan_expired_workflow_run_logs.py +++ b/api/services/clear_free_plan_expired_workflow_run_logs.py @@ -8,7 +8,6 @@ import sqlalchemy as sa from sqlalchemy import select from sqlalchemy.orm import Session -from configs import dify_config from enums.cloud_plan import CloudPlan from extensions.ext_database import db from models import WorkflowAppLog, WorkflowNodeExecutionModel, WorkflowRun @@ -140,9 +139,6 @@ class WorkflowRunCleanup: return [WorkflowRunRow(id=row.id, tenant_id=row.tenant_id, created_at=row.created_at) for row in rows] def _filter_free_tenants(self, tenant_ids: Iterable[str]) -> set[str]: - if not dify_config.BILLING_ENABLED: - return set(tenant_ids) - tenant_id_list = list(tenant_ids) uncached_tenants = [tenant_id for tenant_id in tenant_id_list if tenant_id not in self.billing_cache] @@ -158,8 +154,7 @@ class WorkflowRunCleanup: info = bulk_info.get(tenant_id) if info: try: - raw_plan = info.get("subscription", {}).get("plan") - plan = CloudPlan(raw_plan) + plan = CloudPlan(info) except Exception: logger.exception("Failed to parse billing plan for tenant %s", tenant_id) else: From 66ee07154d7528a5d0ec415b3a3a67706d401394 Mon Sep 17 00:00:00 2001 From: "autofix-ci[bot]" <114827586+autofix-ci[bot]@users.noreply.github.com> Date: Thu, 11 Dec 2025 07:21:13 +0000 Subject: [PATCH 03/42] [autofix.ci] apply automated fixes --- ...ear_free_plan_expired_workflow_run_logs.py | 18 +++++----- ...ear_free_plan_expired_workflow_run_logs.py | 34 +++++++------------ 2 files changed, 21 insertions(+), 31 deletions(-) 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 5972440732..cca5ab15d2 100644 --- a/api/services/clear_free_plan_expired_workflow_run_logs.py +++ b/api/services/clear_free_plan_expired_workflow_run_logs.py @@ -114,9 +114,7 @@ class WorkflowRunCleanup: click.echo(click.style(summary_message, fg="white")) - def _load_batch( - self, session: Session, last_seen: tuple[datetime.datetime, str] | None - ) -> list[WorkflowRunRow]: + def _load_batch(self, session: Session, last_seen: tuple[datetime.datetime, str] | None) -> list[WorkflowRunRow]: stmt = ( select(WorkflowRun.id, WorkflowRun.tenant_id, WorkflowRun.created_at) .where(WorkflowRun.created_at < self.window_end) @@ -166,7 +164,9 @@ class WorkflowRunCleanup: def _delete_runs(self, session: Session, workflow_run_ids: Sequence[str]) -> dict[str, int]: node_execution_ids = session.scalars( - select(WorkflowNodeExecutionModel.id).where(WorkflowNodeExecutionModel.workflow_run_id.in_(workflow_run_ids)) + select(WorkflowNodeExecutionModel.id).where( + WorkflowNodeExecutionModel.workflow_run_id.in_(workflow_run_ids) + ) ).all() offloads_deleted = 0 @@ -199,14 +199,12 @@ class WorkflowRunCleanup: if pause_ids: pause_reasons_deleted = ( - session.query(WorkflowPauseReason).where(WorkflowPauseReason.pause_id.in_(pause_ids)).delete( - synchronize_session=False - ) + session.query(WorkflowPauseReason) + .where(WorkflowPauseReason.pause_id.in_(pause_ids)) + .delete(synchronize_session=False) ) pauses_deleted = ( - session.query(WorkflowPause) - .where(WorkflowPause.id.in_(pause_ids)) - .delete(synchronize_session=False) + session.query(WorkflowPause).where(WorkflowPause.id.in_(pause_ids)).delete(synchronize_session=False) ) trigger_logs_deleted = ( 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 fd756bcddd..57fb8d4e3f 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 @@ -85,9 +85,7 @@ def test_run_deletes_only_free_tenants(monkeypatch: pytest.MonkeyPatch) -> None: batches_returned = 0 - def fake_load_batch( - session: DummySession, last_seen: tuple[datetime.datetime, str] | None - ) -> list[WorkflowRunRow]: + def fake_load_batch(session: DummySession, last_seen: tuple[datetime.datetime, str] | None) -> list[WorkflowRunRow]: nonlocal batches_returned if batches_returned > 0: return [] @@ -130,7 +128,7 @@ def test_run_deletes_only_free_tenants(monkeypatch: pytest.MonkeyPatch) -> None: cleanup.run() assert deleted_ids == [["run-free"]] - assert created_sessions + assert created_sessions assert created_sessions[0].committed is True @@ -142,16 +140,12 @@ def test_run_skips_when_no_free_tenants(monkeypatch: pytest.MonkeyPatch) -> None monkeypatch.setattr( cleanup_module.BillingService, "get_info_bulk", - staticmethod( - lambda tenant_ids: {tenant_id: {"subscription": {"plan": "TEAM"}} for tenant_id in tenant_ids} - ), + staticmethod(lambda tenant_ids: {tenant_id: {"subscription": {"plan": "TEAM"}} for tenant_id in tenant_ids}), ) batches_returned = 0 - def fake_load_batch( - session: DummySession, last_seen: tuple[datetime.datetime, str] | None - ) -> list[WorkflowRunRow]: + def fake_load_batch(session: DummySession, last_seen: tuple[datetime.datetime, str] | None) -> list[WorkflowRunRow]: nonlocal batches_returned if batches_returned > 0: return [] @@ -164,14 +158,14 @@ def test_run_skips_when_no_free_tenants(monkeypatch: pytest.MonkeyPatch) -> None nonlocal delete_called delete_called = True return { - "runs": 0, - "node_executions": 0, - "offloads": 0, - "app_logs": 0, - "trigger_logs": 0, - "pauses": 0, - "pause_reasons": 0 - } + "runs": 0, + "node_executions": 0, + "offloads": 0, + "app_logs": 0, + "trigger_logs": 0, + "pauses": 0, + "pause_reasons": 0, + } def fake_session_factory(engine: object | None = None) -> DummySession: # pragma: no cover - simple factory return DummySession() @@ -189,9 +183,7 @@ def test_run_skips_when_no_free_tenants(monkeypatch: pytest.MonkeyPatch) -> None def test_run_exits_on_empty_batch(monkeypatch: pytest.MonkeyPatch) -> None: cleanup = WorkflowRunCleanup(days=30, batch_size=10) - def fake_load_batch( - session: DummySession, last_seen: tuple[datetime.datetime, str] | None - ) -> list[WorkflowRunRow]: + def fake_load_batch(session: DummySession, last_seen: tuple[datetime.datetime, str] | None) -> list[WorkflowRunRow]: return [] def fake_delete_runs(session: DummySession, workflow_run_ids: list[str]) -> dict[str, int]: From cd51bfb56816bcf04d078d6e3c27d69fa6544e13 Mon Sep 17 00:00:00 2001 From: hjlarry Date: Thu, 11 Dec 2025 15:39:33 +0800 Subject: [PATCH 04/42] fix CI --- api/services/billing_service.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/api/services/billing_service.py b/api/services/billing_service.py index 6ffb1f0cb6..1fb16836b3 100644 --- a/api/services/billing_service.py +++ b/api/services/billing_service.py @@ -41,10 +41,10 @@ class BillingService: try: resp = cls._send_request("POST", "/subscription/plan/batch", json={"tenant_ids": tenant_ids}) + return resp.get("data", {}) except Exception: logger.exception("Failed to fetch billing info batch for tenants: %s", tenant_ids) - - return resp.get("data", {}) + return {} @classmethod def get_tenant_feature_plan_usage_info(cls, tenant_id: str): From 3c5a96a3ee42573fb16a4e4c29157bbc4f038913 Mon Sep 17 00:00:00 2001 From: hjlarry Date: Thu, 11 Dec 2025 15:44:30 +0800 Subject: [PATCH 05/42] batch size to 200 --- api/commands.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/api/commands.py b/api/commands.py index cdf3a09bb7..9a990459c0 100644 --- a/api/commands.py +++ b/api/commands.py @@ -856,7 +856,7 @@ def clear_free_plan_tenant_expired_logs(days: int, batch: int, tenant_ids: list[ @click.command("clean-workflow-runs", help="Clean expired workflow runs and related data for free tenants.") @click.option("--days", default=30, show_default=True, help="Delete workflow runs created before N days ago.") -@click.option("--batch-size", default=1000, show_default=True, help="Batch size for selecting workflow runs.") +@click.option("--batch-size", default=200, show_default=True, help="Batch size for selecting workflow runs.") @click.option( "--start-after", type=click.DateTime(formats=["%Y-%m-%d", "%Y-%m-%dT%H:%M:%S"]), From 22443df772de02a0c1c0ebc00e9bc78cce94f2eb Mon Sep 17 00:00:00 2001 From: hjlarry Date: Thu, 11 Dec 2025 16:08:22 +0800 Subject: [PATCH 06/42] add batch in get_info_bulk --- api/services/billing_service.py | 24 +++++++++++++++++------- 1 file changed, 17 insertions(+), 7 deletions(-) diff --git a/api/services/billing_service.py b/api/services/billing_service.py index 1fb16836b3..be7bc55f13 100644 --- a/api/services/billing_service.py +++ b/api/services/billing_service.py @@ -1,5 +1,6 @@ import logging import os +from collections.abc import Sequence from typing import Literal import httpx @@ -29,7 +30,7 @@ class BillingService: return billing_info @classmethod - def get_info_bulk(cls, tenant_ids: list[str]) -> dict[str, dict]: + def get_info_bulk(cls, tenant_ids: Sequence[str]) -> dict[str, dict]: """ Bulk billing info fetch via billing API. @@ -39,12 +40,21 @@ class BillingService: Mapping of tenant_id -> plan """ - try: - resp = cls._send_request("POST", "/subscription/plan/batch", json={"tenant_ids": tenant_ids}) - return resp.get("data", {}) - except Exception: - logger.exception("Failed to fetch billing info batch for tenants: %s", tenant_ids) - return {} + results: dict[str, dict] = {} + + chunk_size = 200 + for i in range(0, len(tenant_ids), chunk_size): + chunk = tenant_ids[i : i + chunk_size] + try: + resp = cls._send_request("POST", "/subscription/plan/batch", json={"tenant_ids": chunk}) + data = resp.get("data", {}) if isinstance(resp, dict) else {} + if data: + results.update(data) + except Exception: + logger.exception("Failed to fetch billing info batch for tenants: %s", chunk) + continue + + return results @classmethod def get_tenant_feature_plan_usage_info(cls, tenant_id: str): From 46d824b17f8501d75c5ca33e7586f7c004f62a4e Mon Sep 17 00:00:00 2001 From: hjlarry Date: Thu, 11 Dec 2025 16:56:54 +0800 Subject: [PATCH 07/42] refactor to repo layer --- .../api_workflow_run_repository.py | 19 ++ .../sqlalchemy_api_workflow_run_repository.py | 122 +++++++++++- api/services/billing_service.py | 12 +- ...ear_free_plan_expired_workflow_run_logs.py | 180 +++++------------- 4 files changed, 189 insertions(+), 144 deletions(-) diff --git a/api/repositories/api_workflow_run_repository.py b/api/repositories/api_workflow_run_repository.py index fd547c78ba..dea5c781d0 100644 --- a/api/repositories/api_workflow_run_repository.py +++ b/api/repositories/api_workflow_run_repository.py @@ -253,6 +253,25 @@ class APIWorkflowRunRepository(WorkflowExecutionRepository, Protocol): """ ... + def get_runs_batch_for_cleanup( + self, + start_after: datetime | None, + end_before: datetime, + last_seen: tuple[datetime, str] | None, + batch_size: int, + ) -> Sequence[WorkflowRun]: + """ + Fetch a batch of workflow runs within a time window using keyset pagination for cleanup. + """ + ... + + def delete_runs_with_related(self, run_ids: Sequence[str]) -> dict[str, int]: + """ + Delete workflow runs and their related records (node executions, offloads, app logs, + trigger logs, pauses, pause reasons). + """ + ... + def create_workflow_pause( self, workflow_run_id: str, diff --git a/api/repositories/sqlalchemy_api_workflow_run_repository.py b/api/repositories/sqlalchemy_api_workflow_run_repository.py index b172c6a3ac..84ba62076d 100644 --- a/api/repositories/sqlalchemy_api_workflow_run_repository.py +++ b/api/repositories/sqlalchemy_api_workflow_run_repository.py @@ -40,8 +40,17 @@ from libs.infinite_scroll_pagination import InfiniteScrollPagination from libs.time_parser import get_time_threshold from libs.uuid_utils import uuidv7 from models.enums import WorkflowRunTriggeredFrom -from models.workflow import WorkflowPause as WorkflowPauseModel -from models.workflow import WorkflowPauseReason, WorkflowRun +from models.trigger import WorkflowTriggerLog +from models.workflow import ( + WorkflowAppLog, + WorkflowNodeExecutionModel, + WorkflowNodeExecutionOffload, + WorkflowPauseReason, + WorkflowRun, +) +from models.workflow import ( + WorkflowPause as WorkflowPauseModel, +) from repositories.api_workflow_run_repository import APIWorkflowRunRepository from repositories.entities.workflow_pause import WorkflowPauseEntity from repositories.types import ( @@ -314,6 +323,115 @@ class DifyAPISQLAlchemyWorkflowRunRepository(APIWorkflowRunRepository): logger.info("Total deleted %s workflow runs for app %s", total_deleted, app_id) return total_deleted + def get_runs_batch_for_cleanup( + self, + start_after: datetime | None, + end_before: datetime, + last_seen: tuple[datetime, str] | None, + batch_size: int, + ) -> Sequence[WorkflowRun]: + with self._session_maker() as session: + stmt = ( + select(WorkflowRun) + .where(WorkflowRun.created_at < end_before) + .order_by(WorkflowRun.created_at.asc(), WorkflowRun.id.asc()) + .limit(batch_size) + ) + + if start_after: + stmt = stmt.where(WorkflowRun.created_at >= start_after) + + if last_seen: + stmt = stmt.where( + or_( + WorkflowRun.created_at > last_seen[0], + and_(WorkflowRun.created_at == last_seen[0], WorkflowRun.id > last_seen[1]), + ) + ) + + return session.scalars(stmt).all() + + def delete_runs_with_related(self, run_ids: Sequence[str]) -> dict[str, int]: + if not run_ids: + return { + "runs": 0, + "node_executions": 0, + "offloads": 0, + "app_logs": 0, + "trigger_logs": 0, + "pauses": 0, + "pause_reasons": 0, + } + + with self._session_maker() as session: + node_execution_ids = session.scalars( + select(WorkflowNodeExecutionModel.id).where( + WorkflowNodeExecutionModel.workflow_run_id.in_(run_ids) + ) + ).all() + + offloads_deleted = 0 + if node_execution_ids: + offloads_deleted = ( + session.query(WorkflowNodeExecutionOffload) + .where(WorkflowNodeExecutionOffload.node_execution_id.in_(node_execution_ids)) + .delete(synchronize_session=False) + ) + + node_executions_deleted = 0 + if node_execution_ids: + node_executions_deleted = ( + session.query(WorkflowNodeExecutionModel) + .where(WorkflowNodeExecutionModel.id.in_(node_execution_ids)) + .delete(synchronize_session=False) + ) + + app_logs_deleted = ( + session.query(WorkflowAppLog) + .where(WorkflowAppLog.workflow_run_id.in_(run_ids)) + .delete(synchronize_session=False) + ) + + pause_ids = session.scalars( + select(WorkflowPauseModel.id).where(WorkflowPauseModel.workflow_run_id.in_(run_ids)) + ).all() + pause_reasons_deleted = 0 + pauses_deleted = 0 + + if pause_ids: + pause_reasons_deleted = ( + session.query(WorkflowPauseReason).where(WorkflowPauseReason.pause_id.in_(pause_ids)).delete( + synchronize_session=False + ) + ) + pauses_deleted = ( + session.query(WorkflowPauseModel) + .where(WorkflowPauseModel.id.in_(pause_ids)) + .delete(synchronize_session=False) + ) + + trigger_logs_deleted = ( + session.query(WorkflowTriggerLog) + .where(WorkflowTriggerLog.workflow_run_id.in_(run_ids)) + .delete(synchronize_session=False) + ) + + runs_deleted = ( + session.query(WorkflowRun).where(WorkflowRun.id.in_(run_ids)).delete(synchronize_session=False) + ) + + session.commit() + + return { + "runs": runs_deleted, + "node_executions": node_executions_deleted, + "offloads": offloads_deleted, + "app_logs": app_logs_deleted, + "trigger_logs": trigger_logs_deleted, + "pauses": pauses_deleted, + "pause_reasons": pause_reasons_deleted, + } + def create_workflow_pause( self, workflow_run_id: str, diff --git a/api/services/billing_service.py b/api/services/billing_service.py index be7bc55f13..cd7b5fc389 100644 --- a/api/services/billing_service.py +++ b/api/services/billing_service.py @@ -30,7 +30,7 @@ class BillingService: return billing_info @classmethod - def get_info_bulk(cls, tenant_ids: Sequence[str]) -> dict[str, dict]: + def get_info_bulk(cls, tenant_ids: Sequence[str]) -> dict[str, str]: """ Bulk billing info fetch via billing API. @@ -39,17 +39,17 @@ class BillingService: Returns: Mapping of tenant_id -> plan """ - - results: dict[str, dict] = {} + results: dict[str, str] = {} chunk_size = 200 for i in range(0, len(tenant_ids), chunk_size): chunk = tenant_ids[i : i + chunk_size] try: resp = cls._send_request("POST", "/subscription/plan/batch", json={"tenant_ids": chunk}) - data = resp.get("data", {}) if isinstance(resp, dict) else {} - if data: - results.update(data) + data = resp.get("data", {}) + for tenant_id, plan in data.items(): + if isinstance(plan, str): + results[tenant_id] = plan except Exception: logger.exception("Failed to fetch billing info batch for tenants: %s", chunk) continue 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 cca5ab15d2..68ba7787d3 100644 --- a/api/services/clear_free_plan_expired_workflow_run_logs.py +++ b/api/services/clear_free_plan_expired_workflow_run_logs.py @@ -1,30 +1,19 @@ import datetime import logging -from collections.abc import Iterable, Sequence -from dataclasses import dataclass +from collections.abc import Iterable import click -import sqlalchemy as sa -from sqlalchemy import select -from sqlalchemy.orm import Session +from sqlalchemy.orm import sessionmaker from enums.cloud_plan import CloudPlan from extensions.ext_database import db -from models import WorkflowAppLog, WorkflowNodeExecutionModel, WorkflowRun -from models.trigger import WorkflowTriggerLog -from models.workflow import WorkflowNodeExecutionOffload, WorkflowPause, WorkflowPauseReason +from repositories.api_workflow_run_repository import APIWorkflowRunRepository +from repositories.factory import DifyAPIRepositoryFactory from services.billing_service import BillingService logger = logging.getLogger(__name__) -@dataclass(frozen=True) -class WorkflowRunRow: - id: str - tenant_id: str - created_at: datetime.datetime - - class WorkflowRunCleanup: def __init__( self, @@ -32,6 +21,7 @@ class WorkflowRunCleanup: batch_size: int, start_after: datetime.datetime | None = None, end_before: datetime.datetime | None = None, + repo: APIWorkflowRunRepository | None = None, ): if (start_after is None) ^ (end_before is None): raise ValueError("start_after and end_before must be both set or both omitted.") @@ -45,6 +35,9 @@ class WorkflowRunCleanup: self.batch_size = batch_size self.billing_cache: dict[str, CloudPlan | None] = {} + self.repo = repo or DifyAPIRepositoryFactory.create_api_workflow_run_repository( + sessionmaker(bind=db.engine, expire_on_commit=False) + ) def run(self) -> None: click.echo( @@ -61,46 +54,48 @@ class WorkflowRunCleanup: last_seen: tuple[datetime.datetime, str] | None = None while True: - with Session(db.engine) as session: - run_rows = self._load_batch(session, last_seen) - if not run_rows: - break + run_rows = self.repo.get_runs_batch_for_cleanup( + start_after=self.window_start, + end_before=self.window_end, + last_seen=last_seen, + batch_size=self.batch_size, + ) + if not run_rows: + break - batch_index += 1 - last_seen = (run_rows[-1].created_at, run_rows[-1].id) - tenant_ids = {row.tenant_id for row in run_rows} - free_tenants = self._filter_free_tenants(tenant_ids) - free_run_ids = [row.id for row in run_rows if row.tenant_id in free_tenants] - paid_or_skipped = len(run_rows) - len(free_run_ids) + batch_index += 1 + last_seen = (run_rows[-1].created_at, run_rows[-1].id) + tenant_ids = {row.tenant_id for row in run_rows} + free_tenants = self._filter_free_tenants(tenant_ids) + free_run_ids = [row.id for row in run_rows if row.tenant_id in free_tenants] + paid_or_skipped = len(run_rows) - len(free_run_ids) - if not free_run_ids: - click.echo( - click.style( - f"[batch #{batch_index}] skipped (no sandbox runs in batch, {paid_or_skipped} paid)", - fg="yellow", - ) - ) - continue - - try: - counts = self._delete_runs(session, free_run_ids) - session.commit() - except Exception: - session.rollback() - logger.exception("Failed to delete workflow runs batch ending at %s", last_seen[0]) - raise - - total_runs_deleted += counts["runs"] + if not free_run_ids: click.echo( click.style( - f"[batch #{batch_index}] deleted runs: {counts['runs']} " - f"(nodes {counts['node_executions']}, offloads {counts['offloads']}, " - f"app_logs {counts['app_logs']}, trigger_logs {counts['trigger_logs']}, " - f"pauses {counts['pauses']}, pause_reasons {counts['pause_reasons']}); " - f"skipped {paid_or_skipped} paid/unknown", - fg="green", + f"[batch #{batch_index}] skipped (no sandbox runs in batch, {paid_or_skipped} paid/unknown)", + fg="yellow", ) ) + continue + + try: + counts = self.repo.delete_runs_with_related(free_run_ids) + except Exception: + logger.exception("Failed to delete workflow runs batch ending at %s", last_seen[0]) + raise + + total_runs_deleted += counts["runs"] + click.echo( + click.style( + f"[batch #{batch_index}] deleted runs: {counts['runs']} " + f"(nodes {counts['node_executions']}, offloads {counts['offloads']}, " + f"app_logs {counts['app_logs']}, trigger_logs {counts['trigger_logs']}, " + f"pauses {counts['pauses']}, pause_reasons {counts['pause_reasons']}); " + f"skipped {paid_or_skipped} paid/unknown", + fg="green", + ) + ) if self.window_start: summary_message = ( @@ -114,28 +109,6 @@ class WorkflowRunCleanup: click.echo(click.style(summary_message, fg="white")) - def _load_batch(self, session: Session, last_seen: tuple[datetime.datetime, str] | None) -> list[WorkflowRunRow]: - stmt = ( - select(WorkflowRun.id, WorkflowRun.tenant_id, WorkflowRun.created_at) - .where(WorkflowRun.created_at < self.window_end) - .order_by(WorkflowRun.created_at.asc(), WorkflowRun.id.asc()) - .limit(self.batch_size) - ) - - if self.window_start: - stmt = stmt.where(WorkflowRun.created_at >= self.window_start) - - if last_seen: - stmt = stmt.where( - sa.or_( - WorkflowRun.created_at > last_seen[0], - sa.and_(WorkflowRun.created_at == last_seen[0], WorkflowRun.id > last_seen[1]), - ) - ) - - rows = session.execute(stmt).all() - return [WorkflowRunRow(id=row.id, tenant_id=row.tenant_id, created_at=row.created_at) for row in rows] - def _filter_free_tenants(self, tenant_ids: Iterable[str]) -> set[str]: tenant_id_list = list(tenant_ids) uncached_tenants = [tenant_id for tenant_id in tenant_id_list if tenant_id not in self.billing_cache] @@ -161,68 +134,3 @@ class WorkflowRunCleanup: self.billing_cache[tenant_id] = plan return {tenant_id for tenant_id in tenant_id_list if self.billing_cache.get(tenant_id) == CloudPlan.SANDBOX} - - def _delete_runs(self, session: Session, workflow_run_ids: Sequence[str]) -> dict[str, int]: - node_execution_ids = session.scalars( - select(WorkflowNodeExecutionModel.id).where( - WorkflowNodeExecutionModel.workflow_run_id.in_(workflow_run_ids) - ) - ).all() - - offloads_deleted = 0 - if node_execution_ids: - offloads_deleted = ( - session.query(WorkflowNodeExecutionOffload) - .where(WorkflowNodeExecutionOffload.node_execution_id.in_(node_execution_ids)) - .delete(synchronize_session=False) - ) - - node_executions_deleted = 0 - if node_execution_ids: - node_executions_deleted = ( - session.query(WorkflowNodeExecutionModel) - .where(WorkflowNodeExecutionModel.id.in_(node_execution_ids)) - .delete(synchronize_session=False) - ) - - app_logs_deleted = ( - session.query(WorkflowAppLog) - .where(WorkflowAppLog.workflow_run_id.in_(workflow_run_ids)) - .delete(synchronize_session=False) - ) - - pause_ids = session.scalars( - select(WorkflowPause.id).where(WorkflowPause.workflow_run_id.in_(workflow_run_ids)) - ).all() - pause_reasons_deleted = 0 - pauses_deleted = 0 - - if pause_ids: - pause_reasons_deleted = ( - session.query(WorkflowPauseReason) - .where(WorkflowPauseReason.pause_id.in_(pause_ids)) - .delete(synchronize_session=False) - ) - pauses_deleted = ( - session.query(WorkflowPause).where(WorkflowPause.id.in_(pause_ids)).delete(synchronize_session=False) - ) - - trigger_logs_deleted = ( - session.query(WorkflowTriggerLog) - .where(WorkflowTriggerLog.workflow_run_id.in_(workflow_run_ids)) - .delete(synchronize_session=False) - ) - - runs_deleted = ( - session.query(WorkflowRun).where(WorkflowRun.id.in_(workflow_run_ids)).delete(synchronize_session=False) - ) - - return { - "runs": runs_deleted, - "node_executions": node_executions_deleted, - "offloads": offloads_deleted, - "app_logs": app_logs_deleted, - "trigger_logs": trigger_logs_deleted, - "pauses": pauses_deleted, - "pause_reasons": pause_reasons_deleted, - } From 9616c6bb7d8d849f3c10eaaf3831c0dda7ee05f5 Mon Sep 17 00:00:00 2001 From: "autofix-ci[bot]" <114827586+autofix-ci[bot]@users.noreply.github.com> Date: Thu, 11 Dec 2025 08:59:03 +0000 Subject: [PATCH 08/42] [autofix.ci] apply automated fixes --- .../sqlalchemy_api_workflow_run_repository.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/api/repositories/sqlalchemy_api_workflow_run_repository.py b/api/repositories/sqlalchemy_api_workflow_run_repository.py index 84ba62076d..8a0356ec60 100644 --- a/api/repositories/sqlalchemy_api_workflow_run_repository.py +++ b/api/repositories/sqlalchemy_api_workflow_run_repository.py @@ -365,9 +365,7 @@ class DifyAPISQLAlchemyWorkflowRunRepository(APIWorkflowRunRepository): with self._session_maker() as session: node_execution_ids = session.scalars( - select(WorkflowNodeExecutionModel.id).where( - WorkflowNodeExecutionModel.workflow_run_id.in_(run_ids) - ) + select(WorkflowNodeExecutionModel.id).where(WorkflowNodeExecutionModel.workflow_run_id.in_(run_ids)) ).all() offloads_deleted = 0 @@ -400,9 +398,9 @@ class DifyAPISQLAlchemyWorkflowRunRepository(APIWorkflowRunRepository): if pause_ids: pause_reasons_deleted = ( - session.query(WorkflowPauseReason).where(WorkflowPauseReason.pause_id.in_(pause_ids)).delete( - synchronize_session=False - ) + session.query(WorkflowPauseReason) + .where(WorkflowPauseReason.pause_id.in_(pause_ids)) + .delete(synchronize_session=False) ) pauses_deleted = ( session.query(WorkflowPauseModel) From 8df80e0992e69826d7bf57a21a9f0688b1e65292 Mon Sep 17 00:00:00 2001 From: hjlarry Date: Thu, 11 Dec 2025 17:20:52 +0800 Subject: [PATCH 09/42] fix CI --- ...ear_free_plan_expired_workflow_run_logs.py | 17 +- ...ear_free_plan_expired_workflow_run_logs.py | 198 +++++++----------- 2 files changed, 93 insertions(+), 122 deletions(-) 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 68ba7787d3..922d0ab43c 100644 --- a/api/services/clear_free_plan_expired_workflow_run_logs.py +++ b/api/services/clear_free_plan_expired_workflow_run_logs.py @@ -5,10 +5,10 @@ from collections.abc import Iterable import click from sqlalchemy.orm import sessionmaker +from configs import dify_config from enums.cloud_plan import CloudPlan from extensions.ext_database import db from repositories.api_workflow_run_repository import APIWorkflowRunRepository -from repositories.factory import DifyAPIRepositoryFactory from services.billing_service import BillingService logger = logging.getLogger(__name__) @@ -35,9 +35,15 @@ class WorkflowRunCleanup: self.batch_size = batch_size self.billing_cache: dict[str, CloudPlan | None] = {} - self.repo = repo or DifyAPIRepositoryFactory.create_api_workflow_run_repository( - sessionmaker(bind=db.engine, expire_on_commit=False) - ) + if repo: + self.repo = repo + else: + # Lazy import to avoid circular dependency during module import + from repositories.factory import DifyAPIRepositoryFactory + + self.repo = DifyAPIRepositoryFactory.create_api_workflow_run_repository( + sessionmaker(bind=db.engine, expire_on_commit=False) + ) def run(self) -> None: click.echo( @@ -110,6 +116,9 @@ class WorkflowRunCleanup: click.echo(click.style(summary_message, fg="white")) def _filter_free_tenants(self, tenant_ids: Iterable[str]) -> set[str]: + if not dify_config.BILLING_ENABLED: + return set(tenant_ids) + tenant_id_list = list(tenant_ids) uncached_tenants = [tenant_id for tenant_id in tenant_id_list if tenant_id not in self.billing_cache] 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 57fb8d4e3f..29aefc5014 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 @@ -4,25 +4,53 @@ from typing import Any import pytest from services import clear_free_plan_expired_workflow_run_logs as cleanup_module -from services.clear_free_plan_expired_workflow_run_logs import WorkflowRunCleanup, WorkflowRunRow +from services.clear_free_plan_expired_workflow_run_logs import WorkflowRunCleanup -class DummySession: - def __init__(self) -> None: - self.committed = False +class FakeRun: + def __init__(self, run_id: str, tenant_id: str, created_at: datetime.datetime) -> None: + self.id = run_id + self.tenant_id = tenant_id + self.created_at = created_at - def __enter__(self) -> "DummySession": - return self - def __exit__(self, exc_type: object, exc: object, tb: object) -> None: - return None +class FakeRepo: + def __init__(self, batches: list[list[FakeRun]], delete_result: dict[str, int] | None = None) -> None: + self.batches = batches + self.call_idx = 0 + self.deleted: list[list[str]] = [] + self.delete_result = delete_result or { + "runs": 0, + "node_executions": 0, + "offloads": 0, + "app_logs": 0, + "trigger_logs": 0, + "pauses": 0, + "pause_reasons": 0, + } - def commit(self) -> None: - self.committed = True + def get_runs_batch_for_cleanup( + self, + start_after: datetime.datetime | None, + end_before: datetime.datetime, + last_seen: tuple[datetime.datetime, str] | None, + batch_size: int, + ) -> list[FakeRun]: + if self.call_idx >= len(self.batches): + return [] + batch = self.batches[self.call_idx] + self.call_idx += 1 + return batch + + def delete_runs_with_related(self, run_ids: list[str]) -> dict[str, int]: + self.deleted.append(list(run_ids)) + result = self.delete_result.copy() + result["runs"] = len(run_ids) + return result def test_filter_free_tenants_billing_disabled(monkeypatch: pytest.MonkeyPatch) -> None: - cleanup = WorkflowRunCleanup(days=30, batch_size=10) + cleanup = WorkflowRunCleanup(days=30, batch_size=10, repo=FakeRepo([])) monkeypatch.setattr(cleanup_module.dify_config, "BILLING_ENABLED", False) @@ -38,25 +66,24 @@ def test_filter_free_tenants_billing_disabled(monkeypatch: pytest.MonkeyPatch) - def test_filter_free_tenants_bulk_mixed(monkeypatch: pytest.MonkeyPatch) -> None: - cleanup = WorkflowRunCleanup(days=30, batch_size=10) + cleanup = WorkflowRunCleanup(days=30, batch_size=10, repo=FakeRepo([])) monkeypatch.setattr(cleanup_module.dify_config, "BILLING_ENABLED", True) - # seed cache to avoid relying on billing service implementation cleanup.billing_cache["t_free"] = cleanup_module.CloudPlan.SANDBOX cleanup.billing_cache["t_paid"] = cleanup_module.CloudPlan.TEAM monkeypatch.setattr( cleanup_module.BillingService, "get_info_bulk", - staticmethod(lambda tenant_ids: {tenant_id: {} for tenant_id in tenant_ids}), + staticmethod(lambda tenant_ids: dict.fromkeys(tenant_ids, "sandbox")), ) free = cleanup._filter_free_tenants({"t_free", "t_paid", "t_missing"}) - assert free == {"t_free"} + assert free == {"t_free", "t_missing"} def test_filter_free_tenants_bulk_failure(monkeypatch: pytest.MonkeyPatch) -> None: - cleanup = WorkflowRunCleanup(days=30, batch_size=10) + cleanup = WorkflowRunCleanup(days=30, batch_size=10, repo=FakeRepo([])) monkeypatch.setattr(cleanup_module.dify_config, "BILLING_ENABLED", True) monkeypatch.setattr( @@ -72,7 +99,15 @@ def test_filter_free_tenants_bulk_failure(monkeypatch: pytest.MonkeyPatch) -> No def test_run_deletes_only_free_tenants(monkeypatch: pytest.MonkeyPatch) -> None: cutoff = datetime.datetime.now() - cleanup = WorkflowRunCleanup(days=30, batch_size=10) + repo = FakeRepo( + batches=[ + [ + FakeRun("run-free", "t_free", cutoff), + FakeRun("run-paid", "t_paid", cutoff), + ] + ] + ) + cleanup = WorkflowRunCleanup(days=30, batch_size=10, repo=repo) monkeypatch.setattr(cleanup_module.dify_config, "BILLING_ENABLED", True) cleanup.billing_cache["t_free"] = cleanup_module.CloudPlan.SANDBOX @@ -80,122 +115,33 @@ def test_run_deletes_only_free_tenants(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr( cleanup_module.BillingService, "get_info_bulk", - staticmethod(lambda tenant_ids: {tenant_id: {} for tenant_id in tenant_ids}), + staticmethod(lambda tenant_ids: dict.fromkeys(tenant_ids, "sandbox")), ) - batches_returned = 0 - - def fake_load_batch(session: DummySession, last_seen: tuple[datetime.datetime, str] | None) -> list[WorkflowRunRow]: - nonlocal batches_returned - if batches_returned > 0: - return [] - batches_returned += 1 - return [ - WorkflowRunRow(id="run-free", tenant_id="t_free", created_at=cutoff), - WorkflowRunRow(id="run-paid", tenant_id="t_paid", created_at=cutoff), - ] - - deleted_ids: list[list[str]] = [] - - def fake_delete_runs(session: DummySession, workflow_run_ids: list[str]) -> dict[str, int]: - deleted_ids.append(list(workflow_run_ids)) - return { - "runs": len(workflow_run_ids), - "node_executions": 0, - "offloads": 0, - "app_logs": 0, - "trigger_logs": 0, - "pauses": 0, - "pause_reasons": 0, - } - - created_sessions: list[DummySession] = [] - - def fake_session_factory(engine: object | None = None) -> DummySession: - session = DummySession() - created_sessions.append(session) - return session - - monkeypatch.setattr(cleanup, "_load_batch", fake_load_batch) - monkeypatch.setattr(cleanup, "_delete_runs", fake_delete_runs) - monkeypatch.setattr(cleanup_module, "Session", fake_session_factory) - - class DummyDB: - engine: object | None = None - - monkeypatch.setattr(cleanup_module, "db", DummyDB()) - cleanup.run() - assert deleted_ids == [["run-free"]] - assert created_sessions - assert created_sessions[0].committed is True + assert repo.deleted == [["run-free"]] def test_run_skips_when_no_free_tenants(monkeypatch: pytest.MonkeyPatch) -> None: cutoff = datetime.datetime.now() - cleanup = WorkflowRunCleanup(days=30, batch_size=10) + repo = FakeRepo(batches=[[FakeRun("run-paid", "t_paid", cutoff)]]) + cleanup = WorkflowRunCleanup(days=30, batch_size=10, repo=repo) monkeypatch.setattr(cleanup_module.dify_config, "BILLING_ENABLED", True) monkeypatch.setattr( cleanup_module.BillingService, "get_info_bulk", - staticmethod(lambda tenant_ids: {tenant_id: {"subscription": {"plan": "TEAM"}} for tenant_id in tenant_ids}), + staticmethod(lambda tenant_ids: dict.fromkeys(tenant_ids, "team")), ) - batches_returned = 0 - - def fake_load_batch(session: DummySession, last_seen: tuple[datetime.datetime, str] | None) -> list[WorkflowRunRow]: - nonlocal batches_returned - if batches_returned > 0: - return [] - batches_returned += 1 - return [WorkflowRunRow(id="run-paid", tenant_id="t_paid", created_at=cutoff)] - - delete_called = False - - def fake_delete_runs(session: DummySession, workflow_run_ids: list[str]) -> dict[str, int]: - nonlocal delete_called - delete_called = True - return { - "runs": 0, - "node_executions": 0, - "offloads": 0, - "app_logs": 0, - "trigger_logs": 0, - "pauses": 0, - "pause_reasons": 0, - } - - def fake_session_factory(engine: object | None = None) -> DummySession: # pragma: no cover - simple factory - return DummySession() - - monkeypatch.setattr(cleanup, "_load_batch", fake_load_batch) - monkeypatch.setattr(cleanup, "_delete_runs", fake_delete_runs) - monkeypatch.setattr(cleanup_module, "Session", fake_session_factory) - monkeypatch.setattr(cleanup_module, "db", type("DummyDB", (), {"engine": None})) - cleanup.run() - assert delete_called is False + assert repo.deleted == [] -def test_run_exits_on_empty_batch(monkeypatch: pytest.MonkeyPatch) -> None: - cleanup = WorkflowRunCleanup(days=30, batch_size=10) - - def fake_load_batch(session: DummySession, last_seen: tuple[datetime.datetime, str] | None) -> list[WorkflowRunRow]: - return [] - - def fake_delete_runs(session: DummySession, workflow_run_ids: list[str]) -> dict[str, int]: - raise AssertionError("should not delete") - - def fake_session_factory(engine: object | None = None) -> DummySession: # pragma: no cover - simple factory - return DummySession() - - monkeypatch.setattr(cleanup, "_load_batch", fake_load_batch) - monkeypatch.setattr(cleanup, "_delete_runs", fake_delete_runs) - monkeypatch.setattr(cleanup_module, "Session", fake_session_factory) - monkeypatch.setattr(cleanup_module, "db", type("DummyDB", (), {"engine": None})) +def test_run_exits_on_empty_batch() -> None: + cleanup = WorkflowRunCleanup(days=30, batch_size=10, repo=FakeRepo([])) cleanup.run() @@ -203,7 +149,11 @@ def test_run_exits_on_empty_batch(monkeypatch: pytest.MonkeyPatch) -> None: def test_between_sets_window_bounds() -> None: start_after = datetime.datetime(2024, 5, 1, 0, 0, 0) end_before = datetime.datetime(2024, 6, 1, 0, 0, 0) - cleanup = WorkflowRunCleanup(days=30, batch_size=10, start_after=start_after, end_before=end_before) + cleanup = WorkflowRunCleanup(days=30, + batch_size=10, + start_after=start_after, + end_before=end_before, + repo=FakeRepo([])) assert cleanup.window_start == start_after assert cleanup.window_end == end_before @@ -211,13 +161,25 @@ def test_between_sets_window_bounds() -> None: def test_between_requires_both_boundaries() -> None: with pytest.raises(ValueError): - WorkflowRunCleanup(days=30, batch_size=10, start_after=datetime.datetime.now(), end_before=None) + WorkflowRunCleanup( + days=30, + batch_size=10, + start_after=datetime.datetime.now(), + end_before=None, + repo=FakeRepo([]) + ) with pytest.raises(ValueError): - WorkflowRunCleanup(days=30, batch_size=10, start_after=None, end_before=datetime.datetime.now()) + WorkflowRunCleanup( + days=30, + batch_size=10, + start_after=None, + end_before=datetime.datetime.now(), + repo=FakeRepo([]) + ) def test_between_requires_end_after_start() -> None: start_after = datetime.datetime(2024, 6, 1, 0, 0, 0) end_before = datetime.datetime(2024, 5, 1, 0, 0, 0) with pytest.raises(ValueError): - WorkflowRunCleanup(days=30, batch_size=10, start_after=start_after, end_before=end_before) + WorkflowRunCleanup(days=30, batch_size=10, start_after=start_after, end_before=end_before, repo=FakeRepo([])) From 9ecc584a1c4223ed813b3b10dbde9c8c5da858af Mon Sep 17 00:00:00 2001 From: "autofix-ci[bot]" <114827586+autofix-ci[bot]@users.noreply.github.com> Date: Thu, 11 Dec 2025 09:22:56 +0000 Subject: [PATCH 10/42] [autofix.ci] apply automated fixes --- ...ear_free_plan_expired_workflow_run_logs.py | 20 +++++-------------- 1 file changed, 5 insertions(+), 15 deletions(-) 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 29aefc5014..44d5988bf2 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 @@ -149,11 +149,9 @@ def test_run_exits_on_empty_batch() -> None: def test_between_sets_window_bounds() -> None: start_after = datetime.datetime(2024, 5, 1, 0, 0, 0) end_before = datetime.datetime(2024, 6, 1, 0, 0, 0) - cleanup = WorkflowRunCleanup(days=30, - batch_size=10, - start_after=start_after, - end_before=end_before, - repo=FakeRepo([])) + cleanup = WorkflowRunCleanup( + days=30, batch_size=10, start_after=start_after, end_before=end_before, repo=FakeRepo([]) + ) assert cleanup.window_start == start_after assert cleanup.window_end == end_before @@ -162,19 +160,11 @@ def test_between_sets_window_bounds() -> None: def test_between_requires_both_boundaries() -> None: with pytest.raises(ValueError): WorkflowRunCleanup( - days=30, - batch_size=10, - start_after=datetime.datetime.now(), - end_before=None, - repo=FakeRepo([]) + days=30, batch_size=10, start_after=datetime.datetime.now(), end_before=None, repo=FakeRepo([]) ) with pytest.raises(ValueError): WorkflowRunCleanup( - days=30, - batch_size=10, - start_after=None, - end_before=datetime.datetime.now(), - repo=FakeRepo([]) + days=30, batch_size=10, start_after=None, end_before=datetime.datetime.now(), repo=FakeRepo([]) ) From c5365f89bfb70ab715f7f5c8f41a1dc45340e3c9 Mon Sep 17 00:00:00 2001 From: hjlarry Date: Thu, 11 Dec 2025 17:27:23 +0800 Subject: [PATCH 11/42] add status filter --- .../sqlalchemy_api_workflow_run_repository.py | 12 +++++- ..._sqlalchemy_api_workflow_run_repository.py | 37 +++++++++++++++++++ 2 files changed, 48 insertions(+), 1 deletion(-) diff --git a/api/repositories/sqlalchemy_api_workflow_run_repository.py b/api/repositories/sqlalchemy_api_workflow_run_repository.py index 8a0356ec60..8d5623617d 100644 --- a/api/repositories/sqlalchemy_api_workflow_run_repository.py +++ b/api/repositories/sqlalchemy_api_workflow_run_repository.py @@ -333,7 +333,17 @@ class DifyAPISQLAlchemyWorkflowRunRepository(APIWorkflowRunRepository): with self._session_maker() as session: stmt = ( select(WorkflowRun) - .where(WorkflowRun.created_at < end_before) + .where( + WorkflowRun.created_at < end_before, + WorkflowRun.status.in_( + [ + WorkflowExecutionStatus.SUCCEEDED.value, + WorkflowExecutionStatus.FAILED.value, + WorkflowExecutionStatus.STOPPED.value, + WorkflowExecutionStatus.PARTIAL_SUCCEEDED.value, + ] + ), + ) .order_by(WorkflowRun.created_at.asc(), WorkflowRun.id.asc()) .limit(batch_size) ) diff --git a/api/tests/unit_tests/repositories/test_sqlalchemy_api_workflow_run_repository.py b/api/tests/unit_tests/repositories/test_sqlalchemy_api_workflow_run_repository.py index 0c34676252..f3df3b5483 100644 --- a/api/tests/unit_tests/repositories/test_sqlalchemy_api_workflow_run_repository.py +++ b/api/tests/unit_tests/repositories/test_sqlalchemy_api_workflow_run_repository.py @@ -4,6 +4,7 @@ from datetime import UTC, datetime from unittest.mock import Mock, patch import pytest +from sqlalchemy.dialects import postgresql from sqlalchemy.orm import Session, sessionmaker from core.workflow.enums import WorkflowExecutionStatus @@ -104,6 +105,42 @@ class TestDifyAPISQLAlchemyWorkflowRunRepository: return pause +class TestGetRunsBatchForCleanup(TestDifyAPISQLAlchemyWorkflowRunRepository): + def test_get_runs_batch_for_cleanup_filters_terminal_statuses( + self, repository: DifyAPISQLAlchemyWorkflowRunRepository, mock_session: Mock + ): + scalar_result = Mock() + scalar_result.all.return_value = [] + mock_session.scalars.return_value = scalar_result + + repository.get_runs_batch_for_cleanup( + start_after=None, + end_before=datetime(2024, 1, 1), + last_seen=None, + batch_size=50, + ) + + stmt = mock_session.scalars.call_args[0][0] + compiled_sql = str( + stmt.compile( + dialect=postgresql.dialect(), + compile_kwargs={"literal_binds": True}, + ) + ) + + assert "workflow_runs.status" in compiled_sql + for status in ( + WorkflowExecutionStatus.SUCCEEDED, + WorkflowExecutionStatus.FAILED, + WorkflowExecutionStatus.STOPPED, + WorkflowExecutionStatus.PARTIAL_SUCCEEDED, + ): + assert f"'{status.value}'" in compiled_sql + + assert "'running'" not in compiled_sql + assert "'paused'" not in compiled_sql + + class TestCreateWorkflowPause(TestDifyAPISQLAlchemyWorkflowRunRepository): """Test create_workflow_pause method.""" From ccb55da5e127c6289e45c7a6e7e968fdeeed4109 Mon Sep 17 00:00:00 2001 From: hjlarry Date: Fri, 12 Dec 2025 09:13:37 +0800 Subject: [PATCH 12/42] use sqlalchemy 2.0 style --- .../sqlalchemy_api_workflow_run_repository.py | 49 ++++++++----------- 1 file changed, 21 insertions(+), 28 deletions(-) diff --git a/api/repositories/sqlalchemy_api_workflow_run_repository.py b/api/repositories/sqlalchemy_api_workflow_run_repository.py index 8d5623617d..1dc1a3903d 100644 --- a/api/repositories/sqlalchemy_api_workflow_run_repository.py +++ b/api/repositories/sqlalchemy_api_workflow_run_repository.py @@ -380,25 +380,24 @@ class DifyAPISQLAlchemyWorkflowRunRepository(APIWorkflowRunRepository): offloads_deleted = 0 if node_execution_ids: - offloads_deleted = ( - session.query(WorkflowNodeExecutionOffload) - .where(WorkflowNodeExecutionOffload.node_execution_id.in_(node_execution_ids)) - .delete(synchronize_session=False) + offloads_result = session.execute( + delete(WorkflowNodeExecutionOffload).where( + WorkflowNodeExecutionOffload.node_execution_id.in_(node_execution_ids) + ) ) + offloads_deleted = offloads_result.rowcount or 0 node_executions_deleted = 0 if node_execution_ids: - node_executions_deleted = ( - session.query(WorkflowNodeExecutionModel) - .where(WorkflowNodeExecutionModel.id.in_(node_execution_ids)) - .delete(synchronize_session=False) + node_executions_result = session.execute( + delete(WorkflowNodeExecutionModel).where(WorkflowNodeExecutionModel.id.in_(node_execution_ids)) ) + node_executions_deleted = node_executions_result.rowcount or 0 - app_logs_deleted = ( - session.query(WorkflowAppLog) - .where(WorkflowAppLog.workflow_run_id.in_(run_ids)) - .delete(synchronize_session=False) + app_logs_result = session.execute( + delete(WorkflowAppLog).where(WorkflowAppLog.workflow_run_id.in_(run_ids)) ) + app_logs_deleted = app_logs_result.rowcount or 0 pause_ids = session.scalars( select(WorkflowPauseModel.id).where(WorkflowPauseModel.workflow_run_id.in_(run_ids)) @@ -407,26 +406,20 @@ class DifyAPISQLAlchemyWorkflowRunRepository(APIWorkflowRunRepository): pauses_deleted = 0 if pause_ids: - pause_reasons_deleted = ( - session.query(WorkflowPauseReason) - .where(WorkflowPauseReason.pause_id.in_(pause_ids)) - .delete(synchronize_session=False) - ) - pauses_deleted = ( - session.query(WorkflowPauseModel) - .where(WorkflowPauseModel.id.in_(pause_ids)) - .delete(synchronize_session=False) + pause_reasons_result = session.execute( + delete(WorkflowPauseReason).where(WorkflowPauseReason.pause_id.in_(pause_ids)) ) + pause_reasons_deleted = pause_reasons_result.rowcount or 0 + pauses_result = session.execute(delete(WorkflowPauseModel).where(WorkflowPauseModel.id.in_(pause_ids))) + pauses_deleted = pauses_result.rowcount or 0 - trigger_logs_deleted = ( - session.query(WorkflowTriggerLog) - .where(WorkflowTriggerLog.workflow_run_id.in_(run_ids)) - .delete(synchronize_session=False) + trigger_logs_result = session.execute( + delete(WorkflowTriggerLog).where(WorkflowTriggerLog.workflow_run_id.in_(run_ids)) ) + trigger_logs_deleted = trigger_logs_result.rowcount or 0 - runs_deleted = ( - session.query(WorkflowRun).where(WorkflowRun.id.in_(run_ids)).delete(synchronize_session=False) - ) + runs_result = session.execute(delete(WorkflowRun).where(WorkflowRun.id.in_(run_ids))) + runs_deleted = runs_result.rowcount or 0 session.commit() From 60288d1fd7b716faed31cfd675dc21ea6a0c7460 Mon Sep 17 00:00:00 2001 From: hjlarry Date: Fri, 12 Dec 2025 09:22:31 +0800 Subject: [PATCH 13/42] use the workflow trigger log repo --- .../sqlalchemy_api_workflow_run_repository.py | 7 ++--- ...alchemy_workflow_trigger_log_repository.py | 20 +++++++++++- ..._sqlalchemy_api_workflow_run_repository.py | 27 ++++++++++++++++ ...alchemy_workflow_trigger_log_repository.py | 31 +++++++++++++++++++ 4 files changed, 79 insertions(+), 6 deletions(-) create mode 100644 api/tests/unit_tests/repositories/test_sqlalchemy_workflow_trigger_log_repository.py diff --git a/api/repositories/sqlalchemy_api_workflow_run_repository.py b/api/repositories/sqlalchemy_api_workflow_run_repository.py index 1dc1a3903d..2e462cf3ef 100644 --- a/api/repositories/sqlalchemy_api_workflow_run_repository.py +++ b/api/repositories/sqlalchemy_api_workflow_run_repository.py @@ -40,7 +40,6 @@ from libs.infinite_scroll_pagination import InfiniteScrollPagination from libs.time_parser import get_time_threshold from libs.uuid_utils import uuidv7 from models.enums import WorkflowRunTriggeredFrom -from models.trigger import WorkflowTriggerLog from models.workflow import ( WorkflowAppLog, WorkflowNodeExecutionModel, @@ -53,6 +52,7 @@ from models.workflow import ( ) from repositories.api_workflow_run_repository import APIWorkflowRunRepository from repositories.entities.workflow_pause import WorkflowPauseEntity +from repositories.sqlalchemy_workflow_trigger_log_repository import SQLAlchemyWorkflowTriggerLogRepository from repositories.types import ( AverageInteractionStats, DailyRunsStats, @@ -413,10 +413,7 @@ class DifyAPISQLAlchemyWorkflowRunRepository(APIWorkflowRunRepository): pauses_result = session.execute(delete(WorkflowPauseModel).where(WorkflowPauseModel.id.in_(pause_ids))) pauses_deleted = pauses_result.rowcount or 0 - trigger_logs_result = session.execute( - delete(WorkflowTriggerLog).where(WorkflowTriggerLog.workflow_run_id.in_(run_ids)) - ) - trigger_logs_deleted = trigger_logs_result.rowcount or 0 + trigger_logs_deleted = SQLAlchemyWorkflowTriggerLogRepository(session).delete_by_run_ids(run_ids) runs_result = session.execute(delete(WorkflowRun).where(WorkflowRun.id.in_(run_ids))) runs_deleted = runs_result.rowcount or 0 diff --git a/api/repositories/sqlalchemy_workflow_trigger_log_repository.py b/api/repositories/sqlalchemy_workflow_trigger_log_repository.py index 0d67e286b0..800f63daa4 100644 --- a/api/repositories/sqlalchemy_workflow_trigger_log_repository.py +++ b/api/repositories/sqlalchemy_workflow_trigger_log_repository.py @@ -5,7 +5,7 @@ SQLAlchemy implementation of WorkflowTriggerLogRepository. from collections.abc import Sequence from datetime import UTC, datetime, timedelta -from sqlalchemy import and_, select +from sqlalchemy import and_, delete, select from sqlalchemy.orm import Session from models.enums import WorkflowTriggerStatus @@ -84,3 +84,21 @@ class SQLAlchemyWorkflowTriggerLogRepository(WorkflowTriggerLogRepository): ) return list(self.session.scalars(query).all()) + + def delete_by_run_ids(self, run_ids: Sequence[str]) -> int: + """ + Delete trigger logs associated with the given workflow run ids. + + Args: + run_ids: Collection of workflow run identifiers. + + Returns: + Number of rows deleted. + """ + if not run_ids: + return 0 + + result = self.session.execute( + delete(WorkflowTriggerLog).where(WorkflowTriggerLog.workflow_run_id.in_(run_ids)) + ) + return result.rowcount or 0 diff --git a/api/tests/unit_tests/repositories/test_sqlalchemy_api_workflow_run_repository.py b/api/tests/unit_tests/repositories/test_sqlalchemy_api_workflow_run_repository.py index f3df3b5483..9a8edbb1fe 100644 --- a/api/tests/unit_tests/repositories/test_sqlalchemy_api_workflow_run_repository.py +++ b/api/tests/unit_tests/repositories/test_sqlalchemy_api_workflow_run_repository.py @@ -218,6 +218,33 @@ class TestCreateWorkflowPause(TestDifyAPISQLAlchemyWorkflowRunRepository): ) +class TestDeleteRunsWithRelated(TestDifyAPISQLAlchemyWorkflowRunRepository): + def test_uses_trigger_log_repository( + self, repository: DifyAPISQLAlchemyWorkflowRunRepository, mock_session: Mock + ): + node_ids_result = Mock() + node_ids_result.all.return_value = [] + pause_ids_result = Mock() + pause_ids_result.all.return_value = [] + mock_session.scalars.side_effect = [node_ids_result, pause_ids_result] + + # app_logs delete, runs delete + mock_session.execute.side_effect = [Mock(rowcount=0), Mock(rowcount=1)] + + fake_trigger_repo = Mock() + fake_trigger_repo.delete_by_run_ids.return_value = 3 + + with patch( + "repositories.sqlalchemy_api_workflow_run_repository.SQLAlchemyWorkflowTriggerLogRepository", + return_value=fake_trigger_repo, + ): + counts = repository.delete_runs_with_related(["run-1"]) + + fake_trigger_repo.delete_by_run_ids.assert_called_once_with(["run-1"]) + assert counts["trigger_logs"] == 3 + assert counts["runs"] == 1 + + class TestResumeWorkflowPause(TestDifyAPISQLAlchemyWorkflowRunRepository): """Test resume_workflow_pause method.""" diff --git a/api/tests/unit_tests/repositories/test_sqlalchemy_workflow_trigger_log_repository.py b/api/tests/unit_tests/repositories/test_sqlalchemy_workflow_trigger_log_repository.py new file mode 100644 index 0000000000..d409618211 --- /dev/null +++ b/api/tests/unit_tests/repositories/test_sqlalchemy_workflow_trigger_log_repository.py @@ -0,0 +1,31 @@ +from unittest.mock import Mock + +from sqlalchemy.dialects import postgresql +from sqlalchemy.orm import Session + +from repositories.sqlalchemy_workflow_trigger_log_repository import SQLAlchemyWorkflowTriggerLogRepository + + +def test_delete_by_run_ids_executes_delete(): + session = Mock(spec=Session) + session.execute.return_value = Mock(rowcount=2) + repo = SQLAlchemyWorkflowTriggerLogRepository(session) + + deleted = repo.delete_by_run_ids(["run-1", "run-2"]) + + stmt = session.execute.call_args[0][0] + compiled_sql = str(stmt.compile(dialect=postgresql.dialect(), compile_kwargs={"literal_binds": True})) + assert "workflow_trigger_logs" in compiled_sql + assert "'run-1'" in compiled_sql + assert "'run-2'" in compiled_sql + assert deleted == 2 + + +def test_delete_by_run_ids_empty_short_circuits(): + session = Mock(spec=Session) + repo = SQLAlchemyWorkflowTriggerLogRepository(session) + + deleted = repo.delete_by_run_ids([]) + + session.execute.assert_not_called() + assert deleted == 0 From 03993ed4f7093f46acddafe3b4c4f45e23d1c830 Mon Sep 17 00:00:00 2001 From: hjlarry Date: Fri, 12 Dec 2025 09:26:40 +0800 Subject: [PATCH 14/42] rename the get_runs_batch_for_cleanup --- api/repositories/api_workflow_run_repository.py | 4 ++-- api/repositories/sqlalchemy_api_workflow_run_repository.py | 2 +- api/services/clear_free_plan_expired_workflow_run_logs.py | 2 +- .../test_sqlalchemy_api_workflow_run_repository.py | 6 +++--- .../test_clear_free_plan_expired_workflow_run_logs.py | 2 +- 5 files changed, 8 insertions(+), 8 deletions(-) diff --git a/api/repositories/api_workflow_run_repository.py b/api/repositories/api_workflow_run_repository.py index dea5c781d0..cb9c8921b6 100644 --- a/api/repositories/api_workflow_run_repository.py +++ b/api/repositories/api_workflow_run_repository.py @@ -253,7 +253,7 @@ class APIWorkflowRunRepository(WorkflowExecutionRepository, Protocol): """ ... - def get_runs_batch_for_cleanup( + def get_runs_batch_by_time_range( self, start_after: datetime | None, end_before: datetime, @@ -261,7 +261,7 @@ class APIWorkflowRunRepository(WorkflowExecutionRepository, Protocol): batch_size: int, ) -> Sequence[WorkflowRun]: """ - Fetch a batch of workflow runs within a time window using keyset pagination for cleanup. + Fetch a batch of workflow runs within a time window using keyset pagination. """ ... diff --git a/api/repositories/sqlalchemy_api_workflow_run_repository.py b/api/repositories/sqlalchemy_api_workflow_run_repository.py index 2e462cf3ef..24dd9b5e6b 100644 --- a/api/repositories/sqlalchemy_api_workflow_run_repository.py +++ b/api/repositories/sqlalchemy_api_workflow_run_repository.py @@ -323,7 +323,7 @@ class DifyAPISQLAlchemyWorkflowRunRepository(APIWorkflowRunRepository): logger.info("Total deleted %s workflow runs for app %s", total_deleted, app_id) return total_deleted - def get_runs_batch_for_cleanup( + def get_runs_batch_by_time_range( self, start_after: datetime | None, end_before: datetime, 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 922d0ab43c..6734c0b020 100644 --- a/api/services/clear_free_plan_expired_workflow_run_logs.py +++ b/api/services/clear_free_plan_expired_workflow_run_logs.py @@ -60,7 +60,7 @@ class WorkflowRunCleanup: last_seen: tuple[datetime.datetime, str] | None = None while True: - run_rows = self.repo.get_runs_batch_for_cleanup( + run_rows = self.repo.get_runs_batch_by_time_range( start_after=self.window_start, end_before=self.window_end, last_seen=last_seen, diff --git a/api/tests/unit_tests/repositories/test_sqlalchemy_api_workflow_run_repository.py b/api/tests/unit_tests/repositories/test_sqlalchemy_api_workflow_run_repository.py index 9a8edbb1fe..14d197e0ac 100644 --- a/api/tests/unit_tests/repositories/test_sqlalchemy_api_workflow_run_repository.py +++ b/api/tests/unit_tests/repositories/test_sqlalchemy_api_workflow_run_repository.py @@ -105,15 +105,15 @@ class TestDifyAPISQLAlchemyWorkflowRunRepository: return pause -class TestGetRunsBatchForCleanup(TestDifyAPISQLAlchemyWorkflowRunRepository): - def test_get_runs_batch_for_cleanup_filters_terminal_statuses( +class TestGetRunsBatchByTimeRange(TestDifyAPISQLAlchemyWorkflowRunRepository): + def test_get_runs_batch_by_time_range_filters_terminal_statuses( self, repository: DifyAPISQLAlchemyWorkflowRunRepository, mock_session: Mock ): scalar_result = Mock() scalar_result.all.return_value = [] mock_session.scalars.return_value = scalar_result - repository.get_runs_batch_for_cleanup( + repository.get_runs_batch_by_time_range( start_after=None, end_before=datetime(2024, 1, 1), last_seen=None, 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 44d5988bf2..415bb9b67d 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 @@ -29,7 +29,7 @@ class FakeRepo: "pause_reasons": 0, } - def get_runs_batch_for_cleanup( + def get_runs_batch_by_time_range( self, start_after: datetime.datetime | None, end_before: datetime.datetime, From 59e94c520da91c758040b8e33845ed98028d0da6 Mon Sep 17 00:00:00 2001 From: hjlarry Date: Fri, 12 Dec 2025 09:41:11 +0800 Subject: [PATCH 15/42] add ENABLE_CLEAN_MESSAGES --- api/.env.example | 1 + api/configs/feature/__init__.py | 9 ++++----- docker/.env.example | 1 + docker/docker-compose.yaml | 1 + 4 files changed, 7 insertions(+), 5 deletions(-) diff --git a/api/.env.example b/api/.env.example index 516a119d98..9fb4211d18 100644 --- a/api/.env.example +++ b/api/.env.example @@ -552,6 +552,7 @@ ENABLE_CLEAN_UNUSED_DATASETS_TASK=false ENABLE_CREATE_TIDB_SERVERLESS_TASK=false ENABLE_UPDATE_TIDB_SERVERLESS_STATUS_TASK=false ENABLE_CLEAN_MESSAGES=false +ENABLE_WORKFLOW_RUN_CLEANUP_TASK=false ENABLE_MAIL_CLEAN_DOCUMENT_NOTIFY_TASK=false ENABLE_DATASETS_QUEUE_MONITOR=false ENABLE_CHECK_UPGRADABLE_PLUGIN_TASK=true diff --git a/api/configs/feature/__init__.py b/api/configs/feature/__init__.py index e237470709..094ddbfd2b 100644 --- a/api/configs/feature/__init__.py +++ b/api/configs/feature/__init__.py @@ -1065,6 +1065,10 @@ class CeleryScheduleTasksConfig(BaseSettings): description="Enable clean messages task", default=False, ) + ENABLE_WORKFLOW_RUN_CLEANUP_TASK: bool = Field( + description="Enable scheduled workflow run cleanup task", + default=False, + ) ENABLE_MAIL_CLEAN_DOCUMENT_NOTIFY_TASK: bool = Field( description="Enable mail clean document notify task", default=False, @@ -1116,11 +1120,6 @@ class CeleryScheduleTasksConfig(BaseSettings): default=60 * 60, ) - ENABLE_WORKFLOW_RUN_CLEANUP_TASK: bool = Field( - description="Enable scheduled workflow run cleanup task", - default=False, - ) - class PositionConfig(BaseSettings): POSITION_PROVIDER_PINS: str = Field( diff --git a/docker/.env.example b/docker/.env.example index 85e8b1dc7f..285bf75156 100644 --- a/docker/.env.example +++ b/docker/.env.example @@ -1422,6 +1422,7 @@ ENABLE_CLEAN_UNUSED_DATASETS_TASK=false ENABLE_CREATE_TIDB_SERVERLESS_TASK=false ENABLE_UPDATE_TIDB_SERVERLESS_STATUS_TASK=false ENABLE_CLEAN_MESSAGES=false +ENABLE_WORKFLOW_RUN_CLEANUP_TASK=false ENABLE_MAIL_CLEAN_DOCUMENT_NOTIFY_TASK=false ENABLE_DATASETS_QUEUE_MONITOR=false ENABLE_CHECK_UPGRADABLE_PLUGIN_TASK=true diff --git a/docker/docker-compose.yaml b/docker/docker-compose.yaml index b961f6b216..cb33ff42a8 100644 --- a/docker/docker-compose.yaml +++ b/docker/docker-compose.yaml @@ -627,6 +627,7 @@ x-shared-env: &shared-api-worker-env ENABLE_CREATE_TIDB_SERVERLESS_TASK: ${ENABLE_CREATE_TIDB_SERVERLESS_TASK:-false} ENABLE_UPDATE_TIDB_SERVERLESS_STATUS_TASK: ${ENABLE_UPDATE_TIDB_SERVERLESS_STATUS_TASK:-false} ENABLE_CLEAN_MESSAGES: ${ENABLE_CLEAN_MESSAGES:-false} + ENABLE_WORKFLOW_RUN_CLEANUP_TASK: ${ENABLE_WORKFLOW_RUN_CLEANUP_TASK:-false} ENABLE_MAIL_CLEAN_DOCUMENT_NOTIFY_TASK: ${ENABLE_MAIL_CLEAN_DOCUMENT_NOTIFY_TASK:-false} ENABLE_DATASETS_QUEUE_MONITOR: ${ENABLE_DATASETS_QUEUE_MONITOR:-false} ENABLE_CHECK_UPGRADABLE_PLUGIN_TASK: ${ENABLE_CHECK_UPGRADABLE_PLUGIN_TASK:-true} From e321f7c8555af7b074f308b612677b25d25faeb2 Mon Sep 17 00:00:00 2001 From: "autofix-ci[bot]" <114827586+autofix-ci[bot]@users.noreply.github.com> Date: Fri, 12 Dec 2025 01:43:12 +0000 Subject: [PATCH 16/42] [autofix.ci] apply automated fixes --- api/repositories/sqlalchemy_api_workflow_run_repository.py | 4 +--- .../sqlalchemy_workflow_trigger_log_repository.py | 4 +--- .../test_sqlalchemy_api_workflow_run_repository.py | 4 +--- 3 files changed, 3 insertions(+), 9 deletions(-) diff --git a/api/repositories/sqlalchemy_api_workflow_run_repository.py b/api/repositories/sqlalchemy_api_workflow_run_repository.py index 24dd9b5e6b..4b031f1c9b 100644 --- a/api/repositories/sqlalchemy_api_workflow_run_repository.py +++ b/api/repositories/sqlalchemy_api_workflow_run_repository.py @@ -394,9 +394,7 @@ class DifyAPISQLAlchemyWorkflowRunRepository(APIWorkflowRunRepository): ) node_executions_deleted = node_executions_result.rowcount or 0 - app_logs_result = session.execute( - delete(WorkflowAppLog).where(WorkflowAppLog.workflow_run_id.in_(run_ids)) - ) + app_logs_result = session.execute(delete(WorkflowAppLog).where(WorkflowAppLog.workflow_run_id.in_(run_ids))) app_logs_deleted = app_logs_result.rowcount or 0 pause_ids = session.scalars( diff --git a/api/repositories/sqlalchemy_workflow_trigger_log_repository.py b/api/repositories/sqlalchemy_workflow_trigger_log_repository.py index 800f63daa4..6703a83c46 100644 --- a/api/repositories/sqlalchemy_workflow_trigger_log_repository.py +++ b/api/repositories/sqlalchemy_workflow_trigger_log_repository.py @@ -98,7 +98,5 @@ class SQLAlchemyWorkflowTriggerLogRepository(WorkflowTriggerLogRepository): if not run_ids: return 0 - result = self.session.execute( - delete(WorkflowTriggerLog).where(WorkflowTriggerLog.workflow_run_id.in_(run_ids)) - ) + result = self.session.execute(delete(WorkflowTriggerLog).where(WorkflowTriggerLog.workflow_run_id.in_(run_ids))) return result.rowcount or 0 diff --git a/api/tests/unit_tests/repositories/test_sqlalchemy_api_workflow_run_repository.py b/api/tests/unit_tests/repositories/test_sqlalchemy_api_workflow_run_repository.py index 14d197e0ac..90a3ef6985 100644 --- a/api/tests/unit_tests/repositories/test_sqlalchemy_api_workflow_run_repository.py +++ b/api/tests/unit_tests/repositories/test_sqlalchemy_api_workflow_run_repository.py @@ -219,9 +219,7 @@ class TestCreateWorkflowPause(TestDifyAPISQLAlchemyWorkflowRunRepository): class TestDeleteRunsWithRelated(TestDifyAPISQLAlchemyWorkflowRunRepository): - def test_uses_trigger_log_repository( - self, repository: DifyAPISQLAlchemyWorkflowRunRepository, mock_session: Mock - ): + def test_uses_trigger_log_repository(self, repository: DifyAPISQLAlchemyWorkflowRunRepository, mock_session: Mock): node_ids_result = Mock() node_ids_result.all.return_value = [] pause_ids_result = Mock() From 287efc7decdf1c7e3f9c185f1d33fb00fb90be0e Mon Sep 17 00:00:00 2001 From: hjlarry Date: Fri, 12 Dec 2025 10:01:47 +0800 Subject: [PATCH 17/42] fix CI --- .../sqlalchemy_api_workflow_run_repository.py | 12 ++++++------ .../sqlalchemy_workflow_trigger_log_repository.py | 4 +++- 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/api/repositories/sqlalchemy_api_workflow_run_repository.py b/api/repositories/sqlalchemy_api_workflow_run_repository.py index 4b031f1c9b..95de006d98 100644 --- a/api/repositories/sqlalchemy_api_workflow_run_repository.py +++ b/api/repositories/sqlalchemy_api_workflow_run_repository.py @@ -385,17 +385,17 @@ class DifyAPISQLAlchemyWorkflowRunRepository(APIWorkflowRunRepository): WorkflowNodeExecutionOffload.node_execution_id.in_(node_execution_ids) ) ) - offloads_deleted = offloads_result.rowcount or 0 + offloads_deleted = cast(CursorResult, offloads_result).rowcount or 0 node_executions_deleted = 0 if node_execution_ids: node_executions_result = session.execute( delete(WorkflowNodeExecutionModel).where(WorkflowNodeExecutionModel.id.in_(node_execution_ids)) ) - node_executions_deleted = node_executions_result.rowcount or 0 + node_executions_deleted = cast(CursorResult, node_executions_result).rowcount or 0 app_logs_result = session.execute(delete(WorkflowAppLog).where(WorkflowAppLog.workflow_run_id.in_(run_ids))) - app_logs_deleted = app_logs_result.rowcount or 0 + app_logs_deleted = cast(CursorResult, app_logs_result).rowcount or 0 pause_ids = session.scalars( select(WorkflowPauseModel.id).where(WorkflowPauseModel.workflow_run_id.in_(run_ids)) @@ -407,14 +407,14 @@ class DifyAPISQLAlchemyWorkflowRunRepository(APIWorkflowRunRepository): pause_reasons_result = session.execute( delete(WorkflowPauseReason).where(WorkflowPauseReason.pause_id.in_(pause_ids)) ) - pause_reasons_deleted = pause_reasons_result.rowcount or 0 + pause_reasons_deleted = cast(CursorResult, pause_reasons_result).rowcount or 0 pauses_result = session.execute(delete(WorkflowPauseModel).where(WorkflowPauseModel.id.in_(pause_ids))) - pauses_deleted = pauses_result.rowcount or 0 + pauses_deleted = cast(CursorResult, pauses_result).rowcount or 0 trigger_logs_deleted = SQLAlchemyWorkflowTriggerLogRepository(session).delete_by_run_ids(run_ids) runs_result = session.execute(delete(WorkflowRun).where(WorkflowRun.id.in_(run_ids))) - runs_deleted = runs_result.rowcount or 0 + runs_deleted = cast(CursorResult, runs_result).rowcount or 0 session.commit() diff --git a/api/repositories/sqlalchemy_workflow_trigger_log_repository.py b/api/repositories/sqlalchemy_workflow_trigger_log_repository.py index 6703a83c46..d01c35e5ab 100644 --- a/api/repositories/sqlalchemy_workflow_trigger_log_repository.py +++ b/api/repositories/sqlalchemy_workflow_trigger_log_repository.py @@ -4,8 +4,10 @@ SQLAlchemy implementation of WorkflowTriggerLogRepository. from collections.abc import Sequence from datetime import UTC, datetime, timedelta +from typing import cast from sqlalchemy import and_, delete, select +from sqlalchemy.engine import CursorResult from sqlalchemy.orm import Session from models.enums import WorkflowTriggerStatus @@ -99,4 +101,4 @@ class SQLAlchemyWorkflowTriggerLogRepository(WorkflowTriggerLogRepository): return 0 result = self.session.execute(delete(WorkflowTriggerLog).where(WorkflowTriggerLog.workflow_run_id.in_(run_ids))) - return result.rowcount or 0 + return cast(CursorResult, result).rowcount or 0 From a194d09b0a0e269c3a174809dc6f33bc0b3bd573 Mon Sep 17 00:00:00 2001 From: hjlarry Date: Fri, 12 Dec 2025 12:03:29 +0800 Subject: [PATCH 18/42] refactor the repo and service --- .../api_workflow_run_repository.py | 10 +++- .../sqlalchemy_api_workflow_run_repository.py | 11 +++-- ...ear_free_plan_expired_workflow_run_logs.py | 30 +++++++----- ..._sqlalchemy_api_workflow_run_repository.py | 8 ++-- ...ear_free_plan_expired_workflow_run_logs.py | 48 ++++++++++++------- 5 files changed, 65 insertions(+), 42 deletions(-) diff --git a/api/repositories/api_workflow_run_repository.py b/api/repositories/api_workflow_run_repository.py index cb9c8921b6..d423081bf1 100644 --- a/api/repositories/api_workflow_run_repository.py +++ b/api/repositories/api_workflow_run_repository.py @@ -34,10 +34,12 @@ Example: ``` """ -from collections.abc import Sequence +from collections.abc import Callable, Sequence from datetime import datetime from typing import Protocol +from sqlalchemy.orm import Session + from core.workflow.entities.pause_reason import PauseReason from core.workflow.repositories.workflow_execution_repository import WorkflowExecutionRepository from libs.infinite_scroll_pagination import InfiniteScrollPagination @@ -265,7 +267,11 @@ class APIWorkflowRunRepository(WorkflowExecutionRepository, Protocol): """ ... - def delete_runs_with_related(self, run_ids: Sequence[str]) -> dict[str, int]: + def delete_runs_with_related( + self, + run_ids: Sequence[str], + delete_trigger_logs: Callable[[Session, Sequence[str]], int] | None = None, + ) -> dict[str, int]: """ Delete workflow runs and their related records (node executions, offloads, app logs, trigger logs, pauses, pause reasons). diff --git a/api/repositories/sqlalchemy_api_workflow_run_repository.py b/api/repositories/sqlalchemy_api_workflow_run_repository.py index 95de006d98..1fa89e8f64 100644 --- a/api/repositories/sqlalchemy_api_workflow_run_repository.py +++ b/api/repositories/sqlalchemy_api_workflow_run_repository.py @@ -21,7 +21,7 @@ Implementation Notes: import logging import uuid -from collections.abc import Sequence +from collections.abc import Callable, Sequence from datetime import datetime from decimal import Decimal from typing import Any, cast @@ -52,7 +52,6 @@ from models.workflow import ( ) from repositories.api_workflow_run_repository import APIWorkflowRunRepository from repositories.entities.workflow_pause import WorkflowPauseEntity -from repositories.sqlalchemy_workflow_trigger_log_repository import SQLAlchemyWorkflowTriggerLogRepository from repositories.types import ( AverageInteractionStats, DailyRunsStats, @@ -361,7 +360,11 @@ class DifyAPISQLAlchemyWorkflowRunRepository(APIWorkflowRunRepository): return session.scalars(stmt).all() - def delete_runs_with_related(self, run_ids: Sequence[str]) -> dict[str, int]: + def delete_runs_with_related( + self, + run_ids: Sequence[str], + delete_trigger_logs: Callable[[Session, Sequence[str]], int] | None = None, + ) -> dict[str, int]: if not run_ids: return { "runs": 0, @@ -411,7 +414,7 @@ class DifyAPISQLAlchemyWorkflowRunRepository(APIWorkflowRunRepository): pauses_result = session.execute(delete(WorkflowPauseModel).where(WorkflowPauseModel.id.in_(pause_ids))) pauses_deleted = cast(CursorResult, pauses_result).rowcount or 0 - trigger_logs_deleted = SQLAlchemyWorkflowTriggerLogRepository(session).delete_by_run_ids(run_ids) + trigger_logs_deleted = delete_trigger_logs(session, run_ids) if delete_trigger_logs else 0 runs_result = session.execute(delete(WorkflowRun).where(WorkflowRun.id.in_(run_ids))) runs_deleted = cast(CursorResult, runs_result).rowcount or 0 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 6734c0b020..f51beef923 100644 --- a/api/services/clear_free_plan_expired_workflow_run_logs.py +++ b/api/services/clear_free_plan_expired_workflow_run_logs.py @@ -1,14 +1,15 @@ import datetime import logging -from collections.abc import Iterable +from collections.abc import Iterable, Sequence import click -from sqlalchemy.orm import sessionmaker +from sqlalchemy.orm import Session, sessionmaker from configs import dify_config from enums.cloud_plan import CloudPlan from extensions.ext_database import db from repositories.api_workflow_run_repository import APIWorkflowRunRepository +from repositories.sqlalchemy_workflow_trigger_log_repository import SQLAlchemyWorkflowTriggerLogRepository from services.billing_service import BillingService logger = logging.getLogger(__name__) @@ -21,7 +22,6 @@ class WorkflowRunCleanup: batch_size: int, start_after: datetime.datetime | None = None, end_before: datetime.datetime | None = None, - repo: APIWorkflowRunRepository | None = None, ): if (start_after is None) ^ (end_before is None): raise ValueError("start_after and end_before must be both set or both omitted.") @@ -35,15 +35,12 @@ class WorkflowRunCleanup: self.batch_size = batch_size self.billing_cache: dict[str, CloudPlan | None] = {} - if repo: - self.repo = repo - else: - # Lazy import to avoid circular dependency during module import - from repositories.factory import DifyAPIRepositoryFactory + # Lazy import to avoid circular dependency during module import + from repositories.factory import DifyAPIRepositoryFactory - self.repo = DifyAPIRepositoryFactory.create_api_workflow_run_repository( - sessionmaker(bind=db.engine, expire_on_commit=False) - ) + self.workflow_run_repo: APIWorkflowRunRepository = DifyAPIRepositoryFactory.create_api_workflow_run_repository( + sessionmaker(bind=db.engine, expire_on_commit=False) + ) def run(self) -> None: click.echo( @@ -60,7 +57,7 @@ class WorkflowRunCleanup: last_seen: tuple[datetime.datetime, str] | None = None while True: - run_rows = self.repo.get_runs_batch_by_time_range( + run_rows = self.workflow_run_repo.get_runs_batch_by_time_range( start_after=self.window_start, end_before=self.window_end, last_seen=last_seen, @@ -86,7 +83,10 @@ class WorkflowRunCleanup: continue try: - counts = self.repo.delete_runs_with_related(free_run_ids) + counts = self.workflow_run_repo.delete_runs_with_related( + free_run_ids, + delete_trigger_logs=self._delete_trigger_logs, + ) except Exception: logger.exception("Failed to delete workflow runs batch ending at %s", last_seen[0]) raise @@ -143,3 +143,7 @@ class WorkflowRunCleanup: self.billing_cache[tenant_id] = plan return {tenant_id for tenant_id in tenant_id_list if self.billing_cache.get(tenant_id) == CloudPlan.SANDBOX} + + 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) diff --git a/api/tests/unit_tests/repositories/test_sqlalchemy_api_workflow_run_repository.py b/api/tests/unit_tests/repositories/test_sqlalchemy_api_workflow_run_repository.py index 90a3ef6985..8b81b45d67 100644 --- a/api/tests/unit_tests/repositories/test_sqlalchemy_api_workflow_run_repository.py +++ b/api/tests/unit_tests/repositories/test_sqlalchemy_api_workflow_run_repository.py @@ -232,11 +232,9 @@ class TestDeleteRunsWithRelated(TestDifyAPISQLAlchemyWorkflowRunRepository): fake_trigger_repo = Mock() fake_trigger_repo.delete_by_run_ids.return_value = 3 - with patch( - "repositories.sqlalchemy_api_workflow_run_repository.SQLAlchemyWorkflowTriggerLogRepository", - return_value=fake_trigger_repo, - ): - counts = repository.delete_runs_with_related(["run-1"]) + counts = repository.delete_runs_with_related( + ["run-1"], delete_trigger_logs=lambda session, run_ids: fake_trigger_repo.delete_by_run_ids(run_ids) + ) fake_trigger_repo.delete_by_run_ids.assert_called_once_with(["run-1"]) assert counts["trigger_logs"] == 3 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 415bb9b67d..9ca9aa9208 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 @@ -3,6 +3,7 @@ from typing import Any import pytest +import repositories.factory as repo_factory_module from services import clear_free_plan_expired_workflow_run_logs as cleanup_module from services.clear_free_plan_expired_workflow_run_logs import WorkflowRunCleanup @@ -42,15 +43,24 @@ class FakeRepo: self.call_idx += 1 return batch - def delete_runs_with_related(self, run_ids: list[str]) -> dict[str, int]: + def delete_runs_with_related(self, run_ids: list[str], delete_trigger_logs=None) -> dict[str, int]: self.deleted.append(list(run_ids)) result = self.delete_result.copy() result["runs"] = len(run_ids) return result +def create_cleanup(monkeypatch: pytest.MonkeyPatch, repo: FakeRepo, **kwargs: Any) -> WorkflowRunCleanup: + monkeypatch.setattr( + repo_factory_module.DifyAPIRepositoryFactory, + "create_api_workflow_run_repository", + classmethod(lambda _cls, session_maker: repo), + ) + return WorkflowRunCleanup(**kwargs) + + def test_filter_free_tenants_billing_disabled(monkeypatch: pytest.MonkeyPatch) -> None: - cleanup = WorkflowRunCleanup(days=30, batch_size=10, repo=FakeRepo([])) + cleanup = create_cleanup(monkeypatch, repo=FakeRepo([]), days=30, batch_size=10) monkeypatch.setattr(cleanup_module.dify_config, "BILLING_ENABLED", False) @@ -66,7 +76,7 @@ def test_filter_free_tenants_billing_disabled(monkeypatch: pytest.MonkeyPatch) - def test_filter_free_tenants_bulk_mixed(monkeypatch: pytest.MonkeyPatch) -> None: - cleanup = WorkflowRunCleanup(days=30, batch_size=10, repo=FakeRepo([])) + cleanup = create_cleanup(monkeypatch, repo=FakeRepo([]), days=30, batch_size=10) monkeypatch.setattr(cleanup_module.dify_config, "BILLING_ENABLED", True) cleanup.billing_cache["t_free"] = cleanup_module.CloudPlan.SANDBOX @@ -83,7 +93,7 @@ def test_filter_free_tenants_bulk_mixed(monkeypatch: pytest.MonkeyPatch) -> None def test_filter_free_tenants_bulk_failure(monkeypatch: pytest.MonkeyPatch) -> None: - cleanup = WorkflowRunCleanup(days=30, batch_size=10, repo=FakeRepo([])) + cleanup = create_cleanup(monkeypatch, repo=FakeRepo([]), days=30, batch_size=10) monkeypatch.setattr(cleanup_module.dify_config, "BILLING_ENABLED", True) monkeypatch.setattr( @@ -107,7 +117,7 @@ def test_run_deletes_only_free_tenants(monkeypatch: pytest.MonkeyPatch) -> None: ] ] ) - cleanup = WorkflowRunCleanup(days=30, batch_size=10, repo=repo) + cleanup = create_cleanup(monkeypatch, repo=repo, days=30, batch_size=10) monkeypatch.setattr(cleanup_module.dify_config, "BILLING_ENABLED", True) cleanup.billing_cache["t_free"] = cleanup_module.CloudPlan.SANDBOX @@ -126,7 +136,7 @@ def test_run_deletes_only_free_tenants(monkeypatch: pytest.MonkeyPatch) -> None: def test_run_skips_when_no_free_tenants(monkeypatch: pytest.MonkeyPatch) -> None: cutoff = datetime.datetime.now() repo = FakeRepo(batches=[[FakeRun("run-paid", "t_paid", cutoff)]]) - cleanup = WorkflowRunCleanup(days=30, batch_size=10, repo=repo) + cleanup = create_cleanup(monkeypatch, repo=repo, days=30, batch_size=10) monkeypatch.setattr(cleanup_module.dify_config, "BILLING_ENABLED", True) monkeypatch.setattr( @@ -140,36 +150,38 @@ def test_run_skips_when_no_free_tenants(monkeypatch: pytest.MonkeyPatch) -> None assert repo.deleted == [] -def test_run_exits_on_empty_batch() -> None: - cleanup = WorkflowRunCleanup(days=30, batch_size=10, repo=FakeRepo([])) +def test_run_exits_on_empty_batch(monkeypatch: pytest.MonkeyPatch) -> None: + cleanup = create_cleanup(monkeypatch, repo=FakeRepo([]), days=30, batch_size=10) cleanup.run() -def test_between_sets_window_bounds() -> None: +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) - cleanup = WorkflowRunCleanup( - days=30, batch_size=10, start_after=start_after, end_before=end_before, repo=FakeRepo([]) + cleanup = create_cleanup( + monkeypatch, repo=FakeRepo([]), days=30, batch_size=10, start_after=start_after, end_before=end_before ) assert cleanup.window_start == start_after assert cleanup.window_end == end_before -def test_between_requires_both_boundaries() -> None: +def test_between_requires_both_boundaries(monkeypatch: pytest.MonkeyPatch) -> None: with pytest.raises(ValueError): - WorkflowRunCleanup( - days=30, batch_size=10, start_after=datetime.datetime.now(), end_before=None, repo=FakeRepo([]) + create_cleanup( + monkeypatch, repo=FakeRepo([]), days=30, batch_size=10, start_after=datetime.datetime.now(), end_before=None ) with pytest.raises(ValueError): - WorkflowRunCleanup( - days=30, batch_size=10, start_after=None, end_before=datetime.datetime.now(), repo=FakeRepo([]) + create_cleanup( + monkeypatch, repo=FakeRepo([]), days=30, batch_size=10, start_after=None, end_before=datetime.datetime.now() ) -def test_between_requires_end_after_start() -> None: +def test_between_requires_end_after_start(monkeypatch: pytest.MonkeyPatch) -> None: start_after = datetime.datetime(2024, 6, 1, 0, 0, 0) end_before = datetime.datetime(2024, 5, 1, 0, 0, 0) with pytest.raises(ValueError): - WorkflowRunCleanup(days=30, batch_size=10, start_after=start_after, end_before=end_before, repo=FakeRepo([])) + create_cleanup( + monkeypatch, repo=FakeRepo([]), days=30, batch_size=10, start_after=start_after, end_before=end_before + ) From e2cceee024a8e30825d1b9ec0ff8e479d07e4d7c Mon Sep 17 00:00:00 2001 From: hjlarry Date: Fri, 12 Dec 2025 12:27:14 +0800 Subject: [PATCH 19/42] refactor the workflowNodeExecution --- .../api_workflow_run_repository.py | 3 +- ..._api_workflow_node_execution_repository.py | 61 ++++++++++++++++++- .../sqlalchemy_api_workflow_run_repository.py | 31 +++------- ...ear_free_plan_expired_workflow_run_logs.py | 25 ++++++-- ..._sqlalchemy_api_workflow_run_repository.py | 7 ++- ...ear_free_plan_expired_workflow_run_logs.py | 22 +++++-- 6 files changed, 113 insertions(+), 36 deletions(-) diff --git a/api/repositories/api_workflow_run_repository.py b/api/repositories/api_workflow_run_repository.py index d423081bf1..22c3ea7130 100644 --- a/api/repositories/api_workflow_run_repository.py +++ b/api/repositories/api_workflow_run_repository.py @@ -269,7 +269,8 @@ class APIWorkflowRunRepository(WorkflowExecutionRepository, Protocol): def delete_runs_with_related( self, - run_ids: Sequence[str], + runs: Sequence[WorkflowRun], + delete_node_executions: Callable[[Session, Sequence[WorkflowRun]], tuple[int, int]] | None = None, delete_trigger_logs: Callable[[Session, Sequence[str]], int] | None = None, ) -> dict[str, int]: """ diff --git a/api/repositories/sqlalchemy_api_workflow_node_execution_repository.py b/api/repositories/sqlalchemy_api_workflow_node_execution_repository.py index 7e2173acdd..9d0a42f518 100644 --- a/api/repositories/sqlalchemy_api_workflow_node_execution_repository.py +++ b/api/repositories/sqlalchemy_api_workflow_node_execution_repository.py @@ -7,13 +7,13 @@ using SQLAlchemy 2.0 style queries for WorkflowNodeExecutionModel operations. from collections.abc import Sequence from datetime import datetime -from typing import cast +from typing import TypedDict, cast -from sqlalchemy import asc, delete, desc, select +from sqlalchemy import asc, delete, desc, select, tuple_ from sqlalchemy.engine import CursorResult from sqlalchemy.orm import Session, sessionmaker -from models.workflow import WorkflowNodeExecutionModel +from models.workflow import WorkflowNodeExecutionModel, WorkflowNodeExecutionOffload from repositories.api_workflow_node_execution_repository import DifyAPIWorkflowNodeExecutionRepository @@ -290,3 +290,58 @@ class DifyAPISQLAlchemyWorkflowNodeExecutionRepository(DifyAPIWorkflowNodeExecut result = cast(CursorResult, session.execute(stmt)) session.commit() return result.rowcount + + class _RunContext(TypedDict): + run_id: str + tenant_id: str + app_id: str + workflow_id: str + triggered_from: str + + @staticmethod + def delete_by_runs(session: Session, runs: Sequence[_RunContext]) -> tuple[int, int]: + """ + Delete node executions (and offloads) for the given workflow runs using indexed columns. + + Uses the composite index on (tenant_id, app_id, workflow_id, triggered_from, workflow_run_id) + by filtering on those columns with tuple IN. + """ + if not runs: + return 0, 0 + + tuple_values = [ + (run["tenant_id"], run["app_id"], run["workflow_id"], run["triggered_from"], run["run_id"]) for run in runs + ] + + node_execution_ids = session.scalars( + select(WorkflowNodeExecutionModel.id).where( + tuple_( + WorkflowNodeExecutionModel.tenant_id, + WorkflowNodeExecutionModel.app_id, + WorkflowNodeExecutionModel.workflow_id, + WorkflowNodeExecutionModel.triggered_from, + WorkflowNodeExecutionModel.workflow_run_id, + ).in_(tuple_values) + ) + ).all() + + if not node_execution_ids: + return 0, 0 + + offloads_deleted = cast( + CursorResult, + session.execute( + delete(WorkflowNodeExecutionOffload).where( + WorkflowNodeExecutionOffload.node_execution_id.in_(node_execution_ids) + ) + ), + ).rowcount or 0 + + node_executions_deleted = cast( + CursorResult, + session.execute( + delete(WorkflowNodeExecutionModel).where(WorkflowNodeExecutionModel.id.in_(node_execution_ids)) + ), + ).rowcount or 0 + + return node_executions_deleted, offloads_deleted diff --git a/api/repositories/sqlalchemy_api_workflow_run_repository.py b/api/repositories/sqlalchemy_api_workflow_run_repository.py index 1fa89e8f64..f081124a80 100644 --- a/api/repositories/sqlalchemy_api_workflow_run_repository.py +++ b/api/repositories/sqlalchemy_api_workflow_run_repository.py @@ -42,8 +42,6 @@ from libs.uuid_utils import uuidv7 from models.enums import WorkflowRunTriggeredFrom from models.workflow import ( WorkflowAppLog, - WorkflowNodeExecutionModel, - WorkflowNodeExecutionOffload, WorkflowPauseReason, WorkflowRun, ) @@ -362,10 +360,11 @@ class DifyAPISQLAlchemyWorkflowRunRepository(APIWorkflowRunRepository): def delete_runs_with_related( self, - run_ids: Sequence[str], + runs: Sequence[WorkflowRun], + delete_node_executions: Callable[[Session, Sequence[WorkflowRun]], tuple[int, int]] | None = None, delete_trigger_logs: Callable[[Session, Sequence[str]], int] | None = None, ) -> dict[str, int]: - if not run_ids: + if not runs: return { "runs": 0, "node_executions": 0, @@ -377,25 +376,11 @@ class DifyAPISQLAlchemyWorkflowRunRepository(APIWorkflowRunRepository): } with self._session_maker() as session: - node_execution_ids = session.scalars( - select(WorkflowNodeExecutionModel.id).where(WorkflowNodeExecutionModel.workflow_run_id.in_(run_ids)) - ).all() - - offloads_deleted = 0 - if node_execution_ids: - offloads_result = session.execute( - delete(WorkflowNodeExecutionOffload).where( - WorkflowNodeExecutionOffload.node_execution_id.in_(node_execution_ids) - ) - ) - offloads_deleted = cast(CursorResult, offloads_result).rowcount or 0 - - node_executions_deleted = 0 - if node_execution_ids: - node_executions_result = session.execute( - delete(WorkflowNodeExecutionModel).where(WorkflowNodeExecutionModel.id.in_(node_execution_ids)) - ) - node_executions_deleted = cast(CursorResult, node_executions_result).rowcount or 0 + run_ids = [run.id for run in runs] + if delete_node_executions: + node_executions_deleted, offloads_deleted = delete_node_executions(session, runs) + else: + node_executions_deleted, offloads_deleted = 0, 0 app_logs_result = session.execute(delete(WorkflowAppLog).where(WorkflowAppLog.workflow_run_id.in_(run_ids))) app_logs_deleted = cast(CursorResult, app_logs_result).rowcount or 0 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 f51beef923..cd5f0d208d 100644 --- a/api/services/clear_free_plan_expired_workflow_run_logs.py +++ b/api/services/clear_free_plan_expired_workflow_run_logs.py @@ -9,6 +9,9 @@ from configs import dify_config from enums.cloud_plan import CloudPlan from extensions.ext_database import db from repositories.api_workflow_run_repository import APIWorkflowRunRepository +from repositories.sqlalchemy_api_workflow_node_execution_repository import ( + DifyAPISQLAlchemyWorkflowNodeExecutionRepository, +) from repositories.sqlalchemy_workflow_trigger_log_repository import SQLAlchemyWorkflowTriggerLogRepository from services.billing_service import BillingService @@ -70,10 +73,10 @@ class WorkflowRunCleanup: last_seen = (run_rows[-1].created_at, run_rows[-1].id) tenant_ids = {row.tenant_id for row in run_rows} free_tenants = self._filter_free_tenants(tenant_ids) - free_run_ids = [row.id for row in run_rows if row.tenant_id in free_tenants] - paid_or_skipped = len(run_rows) - len(free_run_ids) + free_runs = [row for row in run_rows if row.tenant_id in free_tenants] + paid_or_skipped = len(run_rows) - len(free_runs) - if not free_run_ids: + if not free_runs: click.echo( click.style( f"[batch #{batch_index}] skipped (no sandbox runs in batch, {paid_or_skipped} paid/unknown)", @@ -84,7 +87,8 @@ class WorkflowRunCleanup: try: counts = self.workflow_run_repo.delete_runs_with_related( - free_run_ids, + free_runs, + delete_node_executions=self._delete_node_executions, delete_trigger_logs=self._delete_trigger_logs, ) except Exception: @@ -147,3 +151,16 @@ class WorkflowRunCleanup: 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) + + def _delete_node_executions(self, session: Session, runs: Sequence[object]) -> tuple[int, int]: + run_contexts = [ + { + "run_id": run.id, + "tenant_id": run.tenant_id, + "app_id": run.app_id, + "workflow_id": run.workflow_id, + "triggered_from": run.triggered_from, + } + for run in runs + ] + return DifyAPISQLAlchemyWorkflowNodeExecutionRepository.delete_by_runs(session, run_contexts) diff --git a/api/tests/unit_tests/repositories/test_sqlalchemy_api_workflow_run_repository.py b/api/tests/unit_tests/repositories/test_sqlalchemy_api_workflow_run_repository.py index 8b81b45d67..ef3ac29519 100644 --- a/api/tests/unit_tests/repositories/test_sqlalchemy_api_workflow_run_repository.py +++ b/api/tests/unit_tests/repositories/test_sqlalchemy_api_workflow_run_repository.py @@ -232,11 +232,16 @@ class TestDeleteRunsWithRelated(TestDifyAPISQLAlchemyWorkflowRunRepository): fake_trigger_repo = Mock() fake_trigger_repo.delete_by_run_ids.return_value = 3 + run = Mock(id="run-1", tenant_id="t1", app_id="a1", workflow_id="w1", triggered_from="tf") counts = repository.delete_runs_with_related( - ["run-1"], delete_trigger_logs=lambda session, run_ids: fake_trigger_repo.delete_by_run_ids(run_ids) + [run], + delete_node_executions=lambda session, runs: (2, 1), + delete_trigger_logs=lambda session, run_ids: fake_trigger_repo.delete_by_run_ids(run_ids), ) fake_trigger_repo.delete_by_run_ids.assert_called_once_with(["run-1"]) + assert counts["node_executions"] == 2 + assert counts["offloads"] == 1 assert counts["trigger_logs"] == 3 assert counts["runs"] == 1 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 9ca9aa9208..913a8f7ff4 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 @@ -9,9 +9,20 @@ from services.clear_free_plan_expired_workflow_run_logs import WorkflowRunCleanu class FakeRun: - def __init__(self, run_id: str, tenant_id: str, created_at: datetime.datetime) -> None: + def __init__( + self, + run_id: str, + tenant_id: str, + created_at: datetime.datetime, + app_id: str = "app-1", + workflow_id: str = "wf-1", + triggered_from: str = "workflow-run", + ) -> None: self.id = run_id self.tenant_id = tenant_id + self.app_id = app_id + self.workflow_id = workflow_id + self.triggered_from = triggered_from self.created_at = created_at @@ -43,10 +54,12 @@ class FakeRepo: self.call_idx += 1 return batch - def delete_runs_with_related(self, run_ids: list[str], delete_trigger_logs=None) -> dict[str, int]: - self.deleted.append(list(run_ids)) + def delete_runs_with_related(self, runs: list[FakeRun], + delete_node_executions=None, + delete_trigger_logs=None) -> dict[str, int]: + self.deleted.append([run.id for run in runs]) result = self.delete_result.copy() - result["runs"] = len(run_ids) + result["runs"] = len(runs) return result @@ -56,6 +69,7 @@ def create_cleanup(monkeypatch: pytest.MonkeyPatch, repo: FakeRepo, **kwargs: An "create_api_workflow_run_repository", classmethod(lambda _cls, session_maker: repo), ) + kwargs.pop("repo", None) return WorkflowRunCleanup(**kwargs) From 568cd9bdbcf6f0826c35fae1f68c321213420843 Mon Sep 17 00:00:00 2001 From: "autofix-ci[bot]" <114827586+autofix-ci[bot]@users.noreply.github.com> Date: Fri, 12 Dec 2025 04:31:09 +0000 Subject: [PATCH 20/42] [autofix.ci] apply automated fixes --- ..._api_workflow_node_execution_repository.py | 34 +++++++++++-------- ...ear_free_plan_expired_workflow_run_logs.py | 6 ++-- 2 files changed, 23 insertions(+), 17 deletions(-) diff --git a/api/repositories/sqlalchemy_api_workflow_node_execution_repository.py b/api/repositories/sqlalchemy_api_workflow_node_execution_repository.py index 9d0a42f518..ab116611b8 100644 --- a/api/repositories/sqlalchemy_api_workflow_node_execution_repository.py +++ b/api/repositories/sqlalchemy_api_workflow_node_execution_repository.py @@ -328,20 +328,26 @@ class DifyAPISQLAlchemyWorkflowNodeExecutionRepository(DifyAPIWorkflowNodeExecut if not node_execution_ids: return 0, 0 - offloads_deleted = cast( - CursorResult, - session.execute( - delete(WorkflowNodeExecutionOffload).where( - WorkflowNodeExecutionOffload.node_execution_id.in_(node_execution_ids) - ) - ), - ).rowcount or 0 + offloads_deleted = ( + cast( + CursorResult, + session.execute( + delete(WorkflowNodeExecutionOffload).where( + WorkflowNodeExecutionOffload.node_execution_id.in_(node_execution_ids) + ) + ), + ).rowcount + or 0 + ) - node_executions_deleted = cast( - CursorResult, - session.execute( - delete(WorkflowNodeExecutionModel).where(WorkflowNodeExecutionModel.id.in_(node_execution_ids)) - ), - ).rowcount or 0 + node_executions_deleted = ( + cast( + CursorResult, + session.execute( + delete(WorkflowNodeExecutionModel).where(WorkflowNodeExecutionModel.id.in_(node_execution_ids)) + ), + ).rowcount + or 0 + ) return node_executions_deleted, offloads_deleted 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 913a8f7ff4..3728386e34 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 @@ -54,9 +54,9 @@ class FakeRepo: self.call_idx += 1 return batch - def delete_runs_with_related(self, runs: list[FakeRun], - delete_node_executions=None, - delete_trigger_logs=None) -> dict[str, int]: + def delete_runs_with_related( + self, runs: list[FakeRun], delete_node_executions=None, delete_trigger_logs=None + ) -> dict[str, int]: self.deleted.append([run.id for run in runs]) result = self.delete_result.copy() result["runs"] = len(runs) From 2e6d71c8614368bbf70b8e6eac56cce2680d370a Mon Sep 17 00:00:00 2001 From: hjlarry Date: Fri, 12 Dec 2025 13:51:57 +0800 Subject: [PATCH 21/42] fix CI --- ..._api_workflow_node_execution_repository.py | 4 ++-- ...ear_free_plan_expired_workflow_run_logs.py | 20 ++++++++++++------- ...ear_free_plan_expired_workflow_run_logs.py | 9 +-------- 3 files changed, 16 insertions(+), 17 deletions(-) diff --git a/api/repositories/sqlalchemy_api_workflow_node_execution_repository.py b/api/repositories/sqlalchemy_api_workflow_node_execution_repository.py index ab116611b8..004bd7cad9 100644 --- a/api/repositories/sqlalchemy_api_workflow_node_execution_repository.py +++ b/api/repositories/sqlalchemy_api_workflow_node_execution_repository.py @@ -291,7 +291,7 @@ class DifyAPISQLAlchemyWorkflowNodeExecutionRepository(DifyAPIWorkflowNodeExecut session.commit() return result.rowcount - class _RunContext(TypedDict): + class RunContext(TypedDict): run_id: str tenant_id: str app_id: str @@ -299,7 +299,7 @@ class DifyAPISQLAlchemyWorkflowNodeExecutionRepository(DifyAPIWorkflowNodeExecut triggered_from: str @staticmethod - def delete_by_runs(session: Session, runs: Sequence[_RunContext]) -> tuple[int, int]: + def delete_by_runs(session: Session, runs: Sequence[RunContext]) -> tuple[int, int]: """ Delete node executions (and offloads) for the given workflow runs using indexed columns. 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 cd5f0d208d..17cce15311 100644 --- a/api/services/clear_free_plan_expired_workflow_run_logs.py +++ b/api/services/clear_free_plan_expired_workflow_run_logs.py @@ -8,6 +8,7 @@ from sqlalchemy.orm import Session, sessionmaker from configs import dify_config from enums.cloud_plan import CloudPlan from extensions.ext_database import db +from models.workflow import WorkflowRun from repositories.api_workflow_run_repository import APIWorkflowRunRepository from repositories.sqlalchemy_api_workflow_node_execution_repository import ( DifyAPISQLAlchemyWorkflowNodeExecutionRepository, @@ -25,6 +26,7 @@ class WorkflowRunCleanup: batch_size: int, start_after: datetime.datetime | None = None, end_before: datetime.datetime | None = None, + workflow_run_repo: APIWorkflowRunRepository | None = None, ): if (start_after is None) ^ (end_before is None): raise ValueError("start_after and end_before must be both set or both omitted.") @@ -38,12 +40,16 @@ class WorkflowRunCleanup: self.batch_size = batch_size self.billing_cache: dict[str, CloudPlan | None] = {} - # Lazy import to avoid circular dependency during module import - from repositories.factory import DifyAPIRepositoryFactory + if workflow_run_repo: + self.workflow_run_repo = workflow_run_repo + else: + # Lazy import to avoid circular dependencies during module import + from repositories.factory import DifyAPIRepositoryFactory - self.workflow_run_repo: APIWorkflowRunRepository = DifyAPIRepositoryFactory.create_api_workflow_run_repository( - sessionmaker(bind=db.engine, expire_on_commit=False) - ) + session_maker = sessionmaker(bind=db.engine, expire_on_commit=False) + self.workflow_run_repo: APIWorkflowRunRepository = ( + DifyAPIRepositoryFactory.create_api_workflow_run_repository(session_maker) + ) def run(self) -> None: click.echo( @@ -152,8 +158,8 @@ class WorkflowRunCleanup: trigger_repo = SQLAlchemyWorkflowTriggerLogRepository(session) return trigger_repo.delete_by_run_ids(run_ids) - def _delete_node_executions(self, session: Session, runs: Sequence[object]) -> tuple[int, int]: - run_contexts = [ + def _delete_node_executions(self, session: Session, runs: Sequence[WorkflowRun]) -> tuple[int, int]: + run_contexts: list[DifyAPISQLAlchemyWorkflowNodeExecutionRepository.RunContext] = [ { "run_id": run.id, "tenant_id": run.tenant_id, 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 3728386e34..a1685fcfb0 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 @@ -3,7 +3,6 @@ from typing import Any import pytest -import repositories.factory as repo_factory_module from services import clear_free_plan_expired_workflow_run_logs as cleanup_module from services.clear_free_plan_expired_workflow_run_logs import WorkflowRunCleanup @@ -64,13 +63,7 @@ class FakeRepo: def create_cleanup(monkeypatch: pytest.MonkeyPatch, repo: FakeRepo, **kwargs: Any) -> WorkflowRunCleanup: - monkeypatch.setattr( - repo_factory_module.DifyAPIRepositoryFactory, - "create_api_workflow_run_repository", - classmethod(lambda _cls, session_maker: repo), - ) - kwargs.pop("repo", None) - return WorkflowRunCleanup(**kwargs) + return WorkflowRunCleanup(workflow_run_repo=repo, **kwargs) def test_filter_free_tenants_billing_disabled(monkeypatch: pytest.MonkeyPatch) -> None: From 2f635173a871e77c098331ac06722e589adffc23 Mon Sep 17 00:00:00 2001 From: hjlarry Date: Fri, 12 Dec 2025 14:07:19 +0800 Subject: [PATCH 22/42] fix mypy --- api/services/clear_free_plan_expired_workflow_run_logs.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) 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 17cce15311..1fe2bad2d0 100644 --- a/api/services/clear_free_plan_expired_workflow_run_logs.py +++ b/api/services/clear_free_plan_expired_workflow_run_logs.py @@ -40,6 +40,7 @@ class WorkflowRunCleanup: self.batch_size = batch_size self.billing_cache: dict[str, CloudPlan | None] = {} + self.workflow_run_repo: APIWorkflowRunRepository if workflow_run_repo: self.workflow_run_repo = workflow_run_repo else: @@ -47,9 +48,7 @@ class WorkflowRunCleanup: from repositories.factory import DifyAPIRepositoryFactory session_maker = sessionmaker(bind=db.engine, expire_on_commit=False) - self.workflow_run_repo: APIWorkflowRunRepository = ( - DifyAPIRepositoryFactory.create_api_workflow_run_repository(session_maker) - ) + self.workflow_run_repo = DifyAPIRepositoryFactory.create_api_workflow_run_repository(session_maker) def run(self) -> None: click.echo( From 2eaceed4c5a5a5472938eba34e0e63b7453bccf6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9D=9E=E6=B3=95=E6=93=8D=E4=BD=9C?= Date: Mon, 15 Dec 2025 15:38:41 +0800 Subject: [PATCH 23/42] Update api/migrations/versions/2025_12_10_1504-8a7f2ad7c23e_add_workflow_runs_created_at_idx.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- ..._12_10_1504-8a7f2ad7c23e_add_workflow_runs_created_at_idx.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/api/migrations/versions/2025_12_10_1504-8a7f2ad7c23e_add_workflow_runs_created_at_idx.py b/api/migrations/versions/2025_12_10_1504-8a7f2ad7c23e_add_workflow_runs_created_at_idx.py index 7968429ca8..67c855d054 100644 --- a/api/migrations/versions/2025_12_10_1504-8a7f2ad7c23e_add_workflow_runs_created_at_idx.py +++ b/api/migrations/versions/2025_12_10_1504-8a7f2ad7c23e_add_workflow_runs_created_at_idx.py @@ -6,8 +6,6 @@ Create Date: 2025-12-10 15:04:00.000000 """ from alembic import op -import sqlalchemy as sa - # revision identifiers, used by Alembic. revision = "8a7f2ad7c23e" down_revision = "d57accd375ae" From 5f957a115aecc875ea3cf1c5d76633d46c336dbc Mon Sep 17 00:00:00 2001 From: hjlarry Date: Mon, 15 Dec 2025 16:17:05 +0800 Subject: [PATCH 24/42] 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) From f5952b388493010cd4714f629f7c8618bb03ce7e Mon Sep 17 00:00:00 2001 From: "autofix-ci[bot]" <114827586+autofix-ci[bot]@users.noreply.github.com> Date: Mon, 15 Dec 2025 08:19:04 +0000 Subject: [PATCH 25/42] [autofix.ci] apply automated fixes --- api/services/clear_free_plan_expired_workflow_run_logs.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) 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 b5e5f40823..55603795b8 100644 --- a/api/services/clear_free_plan_expired_workflow_run_logs.py +++ b/api/services/clear_free_plan_expired_workflow_run_logs.py @@ -150,8 +150,7 @@ class WorkflowRunCleanup: ) else: summary_message = ( - f"Cleanup complete. Deleted {total_runs_deleted} workflow runs " - f"before {self.window_end.isoformat()}" + f"Cleanup complete. Deleted {total_runs_deleted} workflow runs before {self.window_end.isoformat()}" ) summary_color = "white" From c42e7c8a97a4ac2cdb524c5cd87f7f59bdd8a64f Mon Sep 17 00:00:00 2001 From: hjlarry Date: Tue, 16 Dec 2025 15:16:04 +0800 Subject: [PATCH 26/42] add grace_deadline logic --- api/configs/feature/__init__.py | 7 +++ api/services/billing_service.py | 36 ++++++++--- ...ear_free_plan_expired_workflow_run_logs.py | 61 ++++++++++++++----- ...ear_free_plan_expired_workflow_run_logs.py | 51 +++++++++++++--- 4 files changed, 124 insertions(+), 31 deletions(-) diff --git a/api/configs/feature/__init__.py b/api/configs/feature/__init__.py index 2e50077b46..16e5de3d4c 100644 --- a/api/configs/feature/__init__.py +++ b/api/configs/feature/__init__.py @@ -647,6 +647,13 @@ class BillingConfig(BaseSettings): default=False, ) + BILLING_FREE_PLAN_GRACE_PERIOD_DAYS: NonNegativeInt = Field( + description=( + "Extra grace period in days applied after a tenant leaves a paid plan before being treated as free." + ), + default=21, + ) + class UpdateConfig(BaseSettings): """ diff --git a/api/services/billing_service.py b/api/services/billing_service.py index cd7b5fc389..b449ada26f 100644 --- a/api/services/billing_service.py +++ b/api/services/billing_service.py @@ -4,6 +4,7 @@ from collections.abc import Sequence from typing import Literal import httpx +from pydantic import BaseModel, ValidationError from tenacity import retry, retry_if_exception_type, stop_before_delay, wait_fixed from werkzeug.exceptions import InternalServerError @@ -16,6 +17,11 @@ from models import Account, TenantAccountJoin, TenantAccountRole logger = logging.getLogger(__name__) +class TenantPlanInfo(BaseModel): + plan: CloudPlan + expiration_date: int + + class BillingService: base_url = os.environ.get("BILLING_API_URL", "BILLING_API_URL") secret_key = os.environ.get("BILLING_API_SECRET_KEY", "BILLING_API_SECRET_KEY") @@ -30,32 +36,46 @@ class BillingService: return billing_info @classmethod - def get_info_bulk(cls, tenant_ids: Sequence[str]) -> dict[str, str]: + def get_info_bulk(cls, tenant_ids: Sequence[str]) -> dict[str, TenantPlanInfo]: """ Bulk billing info fetch via billing API. Payload: {"tenant_ids": ["t1", "t2", ...]} (max 200 per request) Returns: - Mapping of tenant_id -> plan + Mapping of tenant_id -> TenantPlanInfo(plan + expiration timestamp) """ - results: dict[str, str] = {} + results: dict[str, TenantPlanInfo] = {} chunk_size = 200 for i in range(0, len(tenant_ids), chunk_size): chunk = tenant_ids[i : i + chunk_size] try: resp = cls._send_request("POST", "/subscription/plan/batch", json={"tenant_ids": chunk}) - data = resp.get("data", {}) - for tenant_id, plan in data.items(): - if isinstance(plan, str): - results[tenant_id] = plan + results.update(cls._parse_bulk_response(chunk, resp)) except Exception: logger.exception("Failed to fetch billing info batch for tenants: %s", chunk) - continue + raise return results + @classmethod + def _parse_bulk_response(cls, expected_ids: Sequence[str], response: dict) -> dict[str, TenantPlanInfo]: + data = response.get("data") + if not isinstance(data, dict): + raise ValueError("Billing API response missing 'data' object.") + + parsed: dict[str, TenantPlanInfo] = {} + for tenant_id in expected_ids: + payload = data.get(tenant_id) + + try: + parsed[tenant_id] = TenantPlanInfo.model_validate(payload) + except ValidationError as exc: + raise ValueError(f"Invalid billing info for tenant {tenant_id}") from exc + + return parsed + @classmethod def get_tenant_feature_plan_usage_info(cls, tenant_id: str): params = {"tenant_id": tenant_id} 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 55603795b8..c3fbd6600a 100644 --- a/api/services/clear_free_plan_expired_workflow_run_logs.py +++ b/api/services/clear_free_plan_expired_workflow_run_logs.py @@ -14,7 +14,7 @@ from repositories.sqlalchemy_api_workflow_node_execution_repository import ( DifyAPISQLAlchemyWorkflowNodeExecutionRepository, ) from repositories.sqlalchemy_workflow_trigger_log_repository import SQLAlchemyWorkflowTriggerLogRepository -from services.billing_service import BillingService +from services.billing_service import BillingService, TenantPlanInfo logger = logging.getLogger(__name__) @@ -40,8 +40,9 @@ class WorkflowRunCleanup: raise ValueError("end_before must be greater than start_after.") self.batch_size = batch_size - self.billing_cache: dict[str, CloudPlan | None] = {} + self.billing_cache: dict[str, TenantPlanInfo | None] = {} self.dry_run = dry_run + self.free_plan_grace_period_days = dify_config.BILLING_FREE_PLAN_GRACE_PERIOD_DAYS self.workflow_run_repo: APIWorkflowRunRepository if workflow_run_repo: self.workflow_run_repo = workflow_run_repo @@ -157,10 +158,14 @@ class WorkflowRunCleanup: 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: - return set(tenant_ids) - tenant_id_list = list(tenant_ids) + + if not dify_config.BILLING_ENABLED: + return set(tenant_id_list) + + if not tenant_id_list: + return set() + uncached_tenants = [tenant_id for tenant_id in tenant_id_list if tenant_id not in self.billing_cache] if uncached_tenants: @@ -171,19 +176,47 @@ class WorkflowRunCleanup: logger.exception("Failed to fetch billing plans in bulk for tenants: %s", uncached_tenants) for tenant_id in uncached_tenants: - plan: CloudPlan | None = None info = bulk_info.get(tenant_id) - if info: - try: - plan = CloudPlan(info) - except Exception: - logger.exception("Failed to parse billing plan for tenant %s", tenant_id) - else: + if info is None: logger.warning("Missing billing info for tenant %s in bulk resp; treating as non-free", tenant_id) + self.billing_cache[tenant_id] = info - self.billing_cache[tenant_id] = plan + eligible_free_tenants: set[str] = set() + for tenant_id in tenant_id_list: + info = self.billing_cache.get(tenant_id) + if not info: + continue - return {tenant_id for tenant_id in tenant_id_list if self.billing_cache.get(tenant_id) == CloudPlan.SANDBOX} + if info.plan != CloudPlan.SANDBOX: + continue + + if self._is_within_grace_period(tenant_id, info): + continue + + eligible_free_tenants.add(tenant_id) + + return eligible_free_tenants + + def _expiration_datetime(self, tenant_id: str, expiration_value: int) -> datetime.datetime | None: + if expiration_value < 0: + return None + + try: + return datetime.datetime.fromtimestamp(expiration_value, datetime.UTC) + except (OverflowError, OSError, ValueError): + logger.exception("Failed to parse expiration timestamp for tenant %s", tenant_id) + return None + + def _is_within_grace_period(self, tenant_id: str, info: TenantPlanInfo) -> bool: + if self.free_plan_grace_period_days <= 0: + return False + + expiration_at = self._expiration_datetime(tenant_id, info.expiration_date) + if expiration_at is None: + return False + + grace_deadline = expiration_at + datetime.timedelta(days=self.free_plan_grace_period_days) + return datetime.datetime.now(datetime.UTC) < grace_deadline def _delete_trigger_logs(self, session: Session, run_ids: Sequence[str]) -> int: trigger_repo = SQLAlchemyWorkflowTriggerLogRepository(session) 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 24a07e7937..66cd7ff8c9 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 @@ -4,6 +4,7 @@ from typing import Any import pytest from services import clear_free_plan_expired_workflow_run_logs as cleanup_module +from services.billing_service import TenantPlanInfo from services.clear_free_plan_expired_workflow_run_logs import WorkflowRunCleanup @@ -62,7 +63,18 @@ class FakeRepo: return result -def create_cleanup(monkeypatch: pytest.MonkeyPatch, repo: FakeRepo, **kwargs: Any) -> WorkflowRunCleanup: +def plan_info(plan: str, expiration: int) -> TenantPlanInfo: + return TenantPlanInfo(plan=plan, expiration_date=expiration) + + +def create_cleanup( + monkeypatch: pytest.MonkeyPatch, + repo: FakeRepo, + *, + grace_period_days: int = 0, + **kwargs: Any, +) -> WorkflowRunCleanup: + monkeypatch.setattr(cleanup_module.dify_config, "BILLING_FREE_PLAN_GRACE_PERIOD_DAYS", grace_period_days) return WorkflowRunCleanup(workflow_run_repo=repo, **kwargs) @@ -71,7 +83,7 @@ def test_filter_free_tenants_billing_disabled(monkeypatch: pytest.MonkeyPatch) - monkeypatch.setattr(cleanup_module.dify_config, "BILLING_ENABLED", False) - def fail_bulk(_: list[str]) -> dict[str, dict[str, Any]]: + def fail_bulk(_: list[str]) -> dict[str, TenantPlanInfo]: raise RuntimeError("should not call") monkeypatch.setattr(cleanup_module.BillingService, "get_info_bulk", staticmethod(fail_bulk)) @@ -86,12 +98,12 @@ def test_filter_free_tenants_bulk_mixed(monkeypatch: pytest.MonkeyPatch) -> None cleanup = create_cleanup(monkeypatch, repo=FakeRepo([]), days=30, batch_size=10) monkeypatch.setattr(cleanup_module.dify_config, "BILLING_ENABLED", True) - cleanup.billing_cache["t_free"] = cleanup_module.CloudPlan.SANDBOX - cleanup.billing_cache["t_paid"] = cleanup_module.CloudPlan.TEAM + cleanup.billing_cache["t_free"] = plan_info("sandbox", -1) + cleanup.billing_cache["t_paid"] = plan_info("team", -1) monkeypatch.setattr( cleanup_module.BillingService, "get_info_bulk", - staticmethod(lambda tenant_ids: dict.fromkeys(tenant_ids, "sandbox")), + staticmethod(lambda tenant_ids: {tenant_id: plan_info("sandbox", -1) for tenant_id in tenant_ids}), ) free = cleanup._filter_free_tenants({"t_free", "t_paid", "t_missing"}) @@ -99,6 +111,27 @@ def test_filter_free_tenants_bulk_mixed(monkeypatch: pytest.MonkeyPatch) -> None assert free == {"t_free", "t_missing"} +def test_filter_free_tenants_respects_grace_period(monkeypatch: pytest.MonkeyPatch) -> None: + cleanup = create_cleanup(monkeypatch, repo=FakeRepo([]), days=30, batch_size=10, grace_period_days=45) + + monkeypatch.setattr(cleanup_module.dify_config, "BILLING_ENABLED", True) + now = datetime.datetime.now(datetime.UTC) + within_grace_ts = int((now - datetime.timedelta(days=10)).timestamp()) + outside_grace_ts = int((now - datetime.timedelta(days=90)).timestamp()) + + def fake_bulk(_: list[str]) -> dict[str, TenantPlanInfo]: + return { + "recently_downgraded": plan_info("sandbox", within_grace_ts), + "long_sandbox": plan_info("sandbox", outside_grace_ts), + } + + monkeypatch.setattr(cleanup_module.BillingService, "get_info_bulk", staticmethod(fake_bulk)) + + free = cleanup._filter_free_tenants({"recently_downgraded", "long_sandbox"}) + + assert free == {"long_sandbox"} + + def test_filter_free_tenants_bulk_failure(monkeypatch: pytest.MonkeyPatch) -> None: cleanup = create_cleanup(monkeypatch, repo=FakeRepo([]), days=30, batch_size=10) @@ -127,12 +160,12 @@ def test_run_deletes_only_free_tenants(monkeypatch: pytest.MonkeyPatch) -> None: cleanup = create_cleanup(monkeypatch, repo=repo, days=30, batch_size=10) monkeypatch.setattr(cleanup_module.dify_config, "BILLING_ENABLED", True) - cleanup.billing_cache["t_free"] = cleanup_module.CloudPlan.SANDBOX - cleanup.billing_cache["t_paid"] = cleanup_module.CloudPlan.TEAM + cleanup.billing_cache["t_free"] = plan_info("sandbox", -1) + cleanup.billing_cache["t_paid"] = plan_info("team", -1) monkeypatch.setattr( cleanup_module.BillingService, "get_info_bulk", - staticmethod(lambda tenant_ids: dict.fromkeys(tenant_ids, "sandbox")), + staticmethod(lambda tenant_ids: {tenant_id: plan_info("sandbox", -1) for tenant_id in tenant_ids}), ) cleanup.run() @@ -149,7 +182,7 @@ def test_run_skips_when_no_free_tenants(monkeypatch: pytest.MonkeyPatch) -> None monkeypatch.setattr( cleanup_module.BillingService, "get_info_bulk", - staticmethod(lambda tenant_ids: dict.fromkeys(tenant_ids, "team")), + staticmethod(lambda tenant_ids: {tenant_id: plan_info("team", 1893456000) for tenant_id in tenant_ids}), ) cleanup.run() From 774286797d9c577b621c53bc7ded3b43b536340f Mon Sep 17 00:00:00 2001 From: hjlarry Date: Wed, 17 Dec 2025 10:04:11 +0800 Subject: [PATCH 27/42] add batch limit --- api/schedule/clean_workflow_runs_task.py | 4 +--- api/services/clear_free_plan_expired_workflow_run_logs.py | 3 +++ 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/api/schedule/clean_workflow_runs_task.py b/api/schedule/clean_workflow_runs_task.py index b59dc7f823..0f34c534ab 100644 --- a/api/schedule/clean_workflow_runs_task.py +++ b/api/schedule/clean_workflow_runs_task.py @@ -4,10 +4,8 @@ import app from configs import dify_config from services.clear_free_plan_expired_workflow_run_logs import WorkflowRunCleanup -CLEANUP_QUEUE = "retention" - -@app.celery.task(queue=CLEANUP_QUEUE) +@app.celery.task(queue="retention") def clean_workflow_runs_task() -> None: """ Scheduled cleanup for workflow runs and related records (sandbox tenants only). 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 c3fbd6600a..0589bb4868 100644 --- a/api/services/clear_free_plan_expired_workflow_run_logs.py +++ b/api/services/clear_free_plan_expired_workflow_run_logs.py @@ -39,6 +39,9 @@ class WorkflowRunCleanup: if self.window_start and self.window_end <= self.window_start: raise ValueError("end_before must be greater than start_after.") + if batch_size <= 0: + raise ValueError("batch_size must be greater than 0.") + self.batch_size = batch_size self.billing_cache: dict[str, TenantPlanInfo | None] = {} self.dry_run = dry_run From b819552a7c9b670542f02b6119bb280983f0efe8 Mon Sep 17 00:00:00 2001 From: hjlarry Date: Wed, 17 Dec 2025 10:10:40 +0800 Subject: [PATCH 28/42] print run time --- api/commands.py | 15 +++++++++++++-- api/schedule/clean_workflow_runs_task.py | 14 +++++++++++++- 2 files changed, 26 insertions(+), 3 deletions(-) diff --git a/api/commands.py b/api/commands.py index 16554dbedc..49b3bff2e9 100644 --- a/api/commands.py +++ b/api/commands.py @@ -887,7 +887,10 @@ def clean_workflow_runs( if (start_after is None) ^ (end_before is None): raise click.UsageError("--start-after and --end-before must be provided together.") - click.echo(click.style("Starting workflow run cleanup.", fg="white")) + start_time = datetime.datetime.now(datetime.UTC) + click.echo( + click.style(f"Starting workflow run cleanup at {start_time.isoformat()}.", fg="white") + ) WorkflowRunCleanup( days=days, @@ -897,7 +900,15 @@ def clean_workflow_runs( dry_run=dry_run, ).run() - click.echo(click.style("Workflow run cleanup completed.", fg="green")) + end_time = datetime.datetime.now(datetime.UTC) + elapsed = end_time - start_time + click.echo( + click.style( + f"Workflow run cleanup completed. start={start_time.isoformat()} " + f"end={end_time.isoformat()} duration={elapsed}", + fg="green", + ) + ) @click.option("-f", "--force", is_flag=True, help="Skip user confirmation and force the command to execute.") diff --git a/api/schedule/clean_workflow_runs_task.py b/api/schedule/clean_workflow_runs_task.py index 0f34c534ab..dde5d46606 100644 --- a/api/schedule/clean_workflow_runs_task.py +++ b/api/schedule/clean_workflow_runs_task.py @@ -1,3 +1,5 @@ +from datetime import UTC, datetime + import click import app @@ -18,6 +20,8 @@ def clean_workflow_runs_task() -> None: ) ) + start_time = datetime.now(UTC) + WorkflowRunCleanup( days=dify_config.WORKFLOW_LOG_RETENTION_DAYS, batch_size=dify_config.WORKFLOW_LOG_CLEANUP_BATCH_SIZE, @@ -25,4 +29,12 @@ def clean_workflow_runs_task() -> None: end_before=None, ).run() - click.echo(click.style("Scheduled workflow run cleanup finished.", fg="green")) + end_time = datetime.now(UTC) + elapsed = end_time - start_time + click.echo( + click.style( + f"Scheduled workflow run cleanup finished. start={start_time.isoformat()} " + f"end={end_time.isoformat()} duration={elapsed}", + fg="green", + ) + ) From 7bdaab432b6379c89c05b2bb09d9b036f97a0014 Mon Sep 17 00:00:00 2001 From: "autofix-ci[bot]" <114827586+autofix-ci[bot]@users.noreply.github.com> Date: Wed, 17 Dec 2025 02:12:46 +0000 Subject: [PATCH 29/42] [autofix.ci] apply automated fixes --- api/commands.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/api/commands.py b/api/commands.py index 49b3bff2e9..d12941068e 100644 --- a/api/commands.py +++ b/api/commands.py @@ -888,9 +888,7 @@ def clean_workflow_runs( raise click.UsageError("--start-after and --end-before must be provided together.") start_time = datetime.datetime.now(datetime.UTC) - click.echo( - click.style(f"Starting workflow run cleanup at {start_time.isoformat()}.", fg="white") - ) + click.echo(click.style(f"Starting workflow run cleanup at {start_time.isoformat()}.", fg="white")) WorkflowRunCleanup( days=days, From a80194fe060f0273336fe80e16db5634a510bef1 Mon Sep 17 00:00:00 2001 From: hjlarry Date: Wed, 17 Dec 2025 14:46:56 +0800 Subject: [PATCH 30/42] use SANDBOX_RECORDS_CLEAN_GRACEFUL_PERIOD env --- api/configs/feature/__init__.py | 19 ++++++++++++------- ...ear_free_plan_expired_workflow_run_logs.py | 2 +- ...ear_free_plan_expired_workflow_run_logs.py | 6 +++++- 3 files changed, 18 insertions(+), 9 deletions(-) diff --git a/api/configs/feature/__init__.py b/api/configs/feature/__init__.py index 0ca317ee44..8fea64b841 100644 --- a/api/configs/feature/__init__.py +++ b/api/configs/feature/__init__.py @@ -647,13 +647,6 @@ class BillingConfig(BaseSettings): default=False, ) - BILLING_FREE_PLAN_GRACE_PERIOD_DAYS: NonNegativeInt = Field( - description=( - "Extra grace period in days applied after a tenant leaves a paid plan before being treated as free." - ), - default=21, - ) - class UpdateConfig(BaseSettings): """ @@ -1159,6 +1152,17 @@ class CeleryScheduleTasksConfig(BaseSettings): ) +class SandboxRecordsCleanConfig(BaseSettings): + SANDBOX_RECORDS_CLEAN_GRACEFUL_PERIOD: NonNegativeInt = Field( + description="Graceful period in days for sandbox records clean after subscription expiration", + default=21, + ) + SANDBOX_RECORDS_CLEAN_BATCH_SIZE: PositiveInt = Field( + description="Maximum number of records to process in each batch", + default=1000, + ) + + class PositionConfig(BaseSettings): POSITION_PROVIDER_PINS: str = Field( description="Comma-separated list of pinned model providers", @@ -1306,6 +1310,7 @@ class FeatureConfig( PositionConfig, RagEtlConfig, RepositoryConfig, + SandboxRecordsCleanConfig, SecurityConfig, TenantIsolatedTaskQueueConfig, ToolConfig, 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 0589bb4868..6c17e9ddb4 100644 --- a/api/services/clear_free_plan_expired_workflow_run_logs.py +++ b/api/services/clear_free_plan_expired_workflow_run_logs.py @@ -45,7 +45,7 @@ class WorkflowRunCleanup: self.batch_size = batch_size self.billing_cache: dict[str, TenantPlanInfo | None] = {} self.dry_run = dry_run - self.free_plan_grace_period_days = dify_config.BILLING_FREE_PLAN_GRACE_PERIOD_DAYS + self.free_plan_grace_period_days = dify_config.SANDBOX_RECORDS_CLEAN_GRACEFUL_PERIOD self.workflow_run_repo: APIWorkflowRunRepository if workflow_run_repo: self.workflow_run_repo = workflow_run_repo 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 66cd7ff8c9..aca7f6f669 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 @@ -74,7 +74,11 @@ def create_cleanup( grace_period_days: int = 0, **kwargs: Any, ) -> WorkflowRunCleanup: - monkeypatch.setattr(cleanup_module.dify_config, "BILLING_FREE_PLAN_GRACE_PERIOD_DAYS", grace_period_days) + monkeypatch.setattr( + cleanup_module.dify_config, + "SANDBOX_RECORDS_CLEAN_GRACEFUL_PERIOD", + grace_period_days, + ) return WorkflowRunCleanup(workflow_run_repo=repo, **kwargs) From fbeb35fbb595b8c2a33d2dac633cfaaa9ecc3de0 Mon Sep 17 00:00:00 2001 From: hjlarry Date: Wed, 17 Dec 2025 14:54:48 +0800 Subject: [PATCH 31/42] should be SANDBOX_EXPIRED_RECORDS_CLEAN_GRACEFUL_PERIOD --- api/configs/feature/__init__.py | 6 +++--- api/services/clear_free_plan_expired_workflow_run_logs.py | 2 +- .../test_clear_free_plan_expired_workflow_run_logs.py | 2 +- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/api/configs/feature/__init__.py b/api/configs/feature/__init__.py index 8fea64b841..c270e17ccc 100644 --- a/api/configs/feature/__init__.py +++ b/api/configs/feature/__init__.py @@ -1152,12 +1152,12 @@ class CeleryScheduleTasksConfig(BaseSettings): ) -class SandboxRecordsCleanConfig(BaseSettings): - SANDBOX_RECORDS_CLEAN_GRACEFUL_PERIOD: NonNegativeInt = Field( +class SandboxExpiredRecordsCleanConfig(BaseSettings): + SANDBOX_EXPIRED_RECORDS_CLEAN_GRACEFUL_PERIOD: NonNegativeInt = Field( description="Graceful period in days for sandbox records clean after subscription expiration", default=21, ) - SANDBOX_RECORDS_CLEAN_BATCH_SIZE: PositiveInt = Field( + SANDBOX_EXPIRED_RECORDS_CLEAN_BATCH_SIZE: PositiveInt = Field( description="Maximum number of records to process in each batch", default=1000, ) 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 6c17e9ddb4..ed829b853d 100644 --- a/api/services/clear_free_plan_expired_workflow_run_logs.py +++ b/api/services/clear_free_plan_expired_workflow_run_logs.py @@ -45,7 +45,7 @@ class WorkflowRunCleanup: self.batch_size = batch_size self.billing_cache: dict[str, TenantPlanInfo | None] = {} self.dry_run = dry_run - self.free_plan_grace_period_days = dify_config.SANDBOX_RECORDS_CLEAN_GRACEFUL_PERIOD + self.free_plan_grace_period_days = dify_config.SANDBOX_EXPIRED_RECORDS_CLEAN_GRACEFUL_PERIOD self.workflow_run_repo: APIWorkflowRunRepository if workflow_run_repo: self.workflow_run_repo = workflow_run_repo 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 aca7f6f669..c2cfcd9811 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 @@ -76,7 +76,7 @@ def create_cleanup( ) -> WorkflowRunCleanup: monkeypatch.setattr( cleanup_module.dify_config, - "SANDBOX_RECORDS_CLEAN_GRACEFUL_PERIOD", + "SANDBOX_EXPIRED_RECORDS_CLEAN_GRACEFUL_PERIOD", grace_period_days, ) return WorkflowRunCleanup(workflow_run_repo=repo, **kwargs) From fc2d0245126e3d450aeae0a5bd1ee82d585288ad Mon Sep 17 00:00:00 2001 From: hjlarry Date: Wed, 17 Dec 2025 15:09:46 +0800 Subject: [PATCH 32/42] add SANDBOX_EXPIRED_RECORDS_RETENTION_DAYS --- api/configs/feature/__init__.py | 6 +++++- api/schedule/clean_workflow_runs_task.py | 9 ++++++--- 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/api/configs/feature/__init__.py b/api/configs/feature/__init__.py index c270e17ccc..608dbe4022 100644 --- a/api/configs/feature/__init__.py +++ b/api/configs/feature/__init__.py @@ -1161,6 +1161,10 @@ class SandboxExpiredRecordsCleanConfig(BaseSettings): description="Maximum number of records to process in each batch", default=1000, ) + SANDBOX_EXPIRED_RECORDS_RETENTION_DAYS: PositiveInt = Field( + description="Retention days for sandbox expired workflow_run records and message records", + default=30, + ) class PositionConfig(BaseSettings): @@ -1310,7 +1314,7 @@ class FeatureConfig( PositionConfig, RagEtlConfig, RepositoryConfig, - SandboxRecordsCleanConfig, + SandboxExpiredRecordsCleanConfig, SecurityConfig, TenantIsolatedTaskQueueConfig, ToolConfig, diff --git a/api/schedule/clean_workflow_runs_task.py b/api/schedule/clean_workflow_runs_task.py index dde5d46606..96c2f9afea 100644 --- a/api/schedule/clean_workflow_runs_task.py +++ b/api/schedule/clean_workflow_runs_task.py @@ -14,8 +14,11 @@ def clean_workflow_runs_task() -> None: """ click.echo( click.style( - f"Scheduled workflow run cleanup starting: cutoff={dify_config.WORKFLOW_LOG_RETENTION_DAYS} days, " - f"batch={dify_config.WORKFLOW_LOG_CLEANUP_BATCH_SIZE}", + ( + "Scheduled workflow run cleanup starting: " + f"cutoff={dify_config.SANDBOX_EXPIRED_RECORDS_RETENTION_DAYS} days, " + f"batch={dify_config.WORKFLOW_LOG_CLEANUP_BATCH_SIZE}" + ), fg="green", ) ) @@ -23,7 +26,7 @@ def clean_workflow_runs_task() -> None: start_time = datetime.now(UTC) WorkflowRunCleanup( - days=dify_config.WORKFLOW_LOG_RETENTION_DAYS, + days=dify_config.SANDBOX_EXPIRED_RECORDS_RETENTION_DAYS, batch_size=dify_config.WORKFLOW_LOG_CLEANUP_BATCH_SIZE, start_after=None, end_before=None, From eafb853764f0a46111f8fd5c52f6aa0ce897c3cb Mon Sep 17 00:00:00 2001 From: hjlarry Date: Wed, 17 Dec 2025 15:10:29 +0800 Subject: [PATCH 33/42] use SANDBOX_EXPIRED_RECORDS_CLEAN_BATCH_SIZE --- api/schedule/clean_workflow_runs_task.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/api/schedule/clean_workflow_runs_task.py b/api/schedule/clean_workflow_runs_task.py index 96c2f9afea..540a46f33d 100644 --- a/api/schedule/clean_workflow_runs_task.py +++ b/api/schedule/clean_workflow_runs_task.py @@ -17,7 +17,7 @@ def clean_workflow_runs_task() -> None: ( "Scheduled workflow run cleanup starting: " f"cutoff={dify_config.SANDBOX_EXPIRED_RECORDS_RETENTION_DAYS} days, " - f"batch={dify_config.WORKFLOW_LOG_CLEANUP_BATCH_SIZE}" + f"batch={dify_config.SANDBOX_EXPIRED_RECORDS_CLEAN_BATCH_SIZE}" ), fg="green", ) @@ -27,7 +27,7 @@ def clean_workflow_runs_task() -> None: WorkflowRunCleanup( days=dify_config.SANDBOX_EXPIRED_RECORDS_RETENTION_DAYS, - batch_size=dify_config.WORKFLOW_LOG_CLEANUP_BATCH_SIZE, + batch_size=dify_config.SANDBOX_EXPIRED_RECORDS_CLEAN_BATCH_SIZE, start_after=None, end_before=None, ).run() From 0e9b6f63e097d5731a294573e00e3a0160274563 Mon Sep 17 00:00:00 2001 From: hjlarry Date: Wed, 17 Dec 2025 15:16:29 +0800 Subject: [PATCH 34/42] change migration file version --- ...7_1504-8a7f2ad7c23e_add_workflow_runs_created_at_idx.py} | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) rename api/migrations/versions/{2025_12_10_1504-8a7f2ad7c23e_add_workflow_runs_created_at_idx.py => 2025_12_17_1504-8a7f2ad7c23e_add_workflow_runs_created_at_idx.py} (86%) diff --git a/api/migrations/versions/2025_12_10_1504-8a7f2ad7c23e_add_workflow_runs_created_at_idx.py b/api/migrations/versions/2025_12_17_1504-8a7f2ad7c23e_add_workflow_runs_created_at_idx.py similarity index 86% rename from api/migrations/versions/2025_12_10_1504-8a7f2ad7c23e_add_workflow_runs_created_at_idx.py rename to api/migrations/versions/2025_12_17_1504-8a7f2ad7c23e_add_workflow_runs_created_at_idx.py index 67c855d054..2acd00fc47 100644 --- a/api/migrations/versions/2025_12_10_1504-8a7f2ad7c23e_add_workflow_runs_created_at_idx.py +++ b/api/migrations/versions/2025_12_17_1504-8a7f2ad7c23e_add_workflow_runs_created_at_idx.py @@ -1,14 +1,14 @@ """Add index on workflow_runs.created_at Revision ID: 8a7f2ad7c23e -Revises: d57accd375ae -Create Date: 2025-12-10 15:04:00.000000 +Revises: 03ea244985ce +Create Date: 2025-12-17 15:04:00.000000 """ from alembic import op # revision identifiers, used by Alembic. revision = "8a7f2ad7c23e" -down_revision = "d57accd375ae" +down_revision = "03ea244985ce" branch_labels = None depends_on = None From b8c1f4c705248657f0b819cb9b94ec3dca513da3 Mon Sep 17 00:00:00 2001 From: hjlarry Date: Wed, 17 Dec 2025 15:25:07 +0800 Subject: [PATCH 35/42] update .env.example --- api/.env.example | 5 +++++ docker/.env.example | 5 +++++ docker/docker-compose.yaml | 3 +++ 3 files changed, 13 insertions(+) diff --git a/api/.env.example b/api/.env.example index 11af6a795a..d96927caf0 100644 --- a/api/.env.example +++ b/api/.env.example @@ -582,6 +582,11 @@ WORKFLOW_SCHEDULE_POLLER_BATCH_SIZE=100 # Maximum number of scheduled workflows to dispatch per tick (0 for unlimited) WORKFLOW_SCHEDULE_MAX_DISPATCH_PER_TICK=0 +# Sandbox expired records clean configuration +SANDBOX_EXPIRED_RECORDS_CLEAN_GRACEFUL_PERIOD=21 +SANDBOX_EXPIRED_RECORDS_CLEAN_BATCH_SIZE=1000 +SANDBOX_EXPIRED_RECORDS_RETENTION_DAYS=30 + # Position configuration POSITION_TOOL_PINS= POSITION_TOOL_INCLUDES= diff --git a/docker/.env.example b/docker/.env.example index e47eea7241..604c41b3a3 100644 --- a/docker/.env.example +++ b/docker/.env.example @@ -1465,6 +1465,11 @@ WORKFLOW_SCHEDULE_POLLER_INTERVAL=1 WORKFLOW_SCHEDULE_POLLER_BATCH_SIZE=100 WORKFLOW_SCHEDULE_MAX_DISPATCH_PER_TICK=0 +# Sandbox expired records clean configuration +SANDBOX_EXPIRED_RECORDS_CLEAN_GRACEFUL_PERIOD=21 +SANDBOX_EXPIRED_RECORDS_CLEAN_BATCH_SIZE=1000 +SANDBOX_EXPIRED_RECORDS_RETENTION_DAYS=30 + # Tenant isolated task queue configuration TENANT_ISOLATED_TASK_CONCURRENCY=1 diff --git a/docker/docker-compose.yaml b/docker/docker-compose.yaml index 2b8d009aa4..51cec2ddbd 100644 --- a/docker/docker-compose.yaml +++ b/docker/docker-compose.yaml @@ -656,6 +656,9 @@ x-shared-env: &shared-api-worker-env WORKFLOW_SCHEDULE_POLLER_INTERVAL: ${WORKFLOW_SCHEDULE_POLLER_INTERVAL:-1} WORKFLOW_SCHEDULE_POLLER_BATCH_SIZE: ${WORKFLOW_SCHEDULE_POLLER_BATCH_SIZE:-100} WORKFLOW_SCHEDULE_MAX_DISPATCH_PER_TICK: ${WORKFLOW_SCHEDULE_MAX_DISPATCH_PER_TICK:-0} + SANDBOX_EXPIRED_RECORDS_CLEAN_GRACEFUL_PERIOD: ${SANDBOX_EXPIRED_RECORDS_CLEAN_GRACEFUL_PERIOD:-21} + SANDBOX_EXPIRED_RECORDS_CLEAN_BATCH_SIZE: ${SANDBOX_EXPIRED_RECORDS_CLEAN_BATCH_SIZE:-1000} + SANDBOX_EXPIRED_RECORDS_RETENTION_DAYS: ${SANDBOX_EXPIRED_RECORDS_RETENTION_DAYS:-30} TENANT_ISOLATED_TASK_CONCURRENCY: ${TENANT_ISOLATED_TASK_CONCURRENCY:-1} ANNOTATION_IMPORT_FILE_SIZE_LIMIT: ${ANNOTATION_IMPORT_FILE_SIZE_LIMIT:-2} ANNOTATION_IMPORT_MAX_RECORDS: ${ANNOTATION_IMPORT_MAX_RECORDS:-10000} From 7bccd8be443ed3acd1fcc7e1891843e65055b941 Mon Sep 17 00:00:00 2001 From: hjlarry Date: Thu, 18 Dec 2025 15:29:31 +0800 Subject: [PATCH 36/42] seprate public part to another PR --- api/.env.example | 5 --- api/configs/feature/__init__.py | 15 ------- api/services/billing_service.py | 41 ------------------- ...ear_free_plan_expired_workflow_run_logs.py | 2 +- docker/.env.example | 5 --- docker/docker-compose.yaml | 3 -- 6 files changed, 1 insertion(+), 70 deletions(-) diff --git a/api/.env.example b/api/.env.example index c2bfa9be6a..c537ecdc27 100644 --- a/api/.env.example +++ b/api/.env.example @@ -582,11 +582,6 @@ WORKFLOW_SCHEDULE_POLLER_BATCH_SIZE=100 # Maximum number of scheduled workflows to dispatch per tick (0 for unlimited) WORKFLOW_SCHEDULE_MAX_DISPATCH_PER_TICK=0 -# Sandbox expired records clean configuration -SANDBOX_EXPIRED_RECORDS_CLEAN_GRACEFUL_PERIOD=21 -SANDBOX_EXPIRED_RECORDS_CLEAN_BATCH_SIZE=1000 -SANDBOX_EXPIRED_RECORDS_RETENTION_DAYS=30 - # Position configuration POSITION_TOOL_PINS= POSITION_TOOL_INCLUDES= diff --git a/api/configs/feature/__init__.py b/api/configs/feature/__init__.py index 1391d2915e..b854293367 100644 --- a/api/configs/feature/__init__.py +++ b/api/configs/feature/__init__.py @@ -1152,21 +1152,6 @@ class CeleryScheduleTasksConfig(BaseSettings): ) -class SandboxExpiredRecordsCleanConfig(BaseSettings): - SANDBOX_EXPIRED_RECORDS_CLEAN_GRACEFUL_PERIOD: NonNegativeInt = Field( - description="Graceful period in days for sandbox records clean after subscription expiration", - default=21, - ) - SANDBOX_EXPIRED_RECORDS_CLEAN_BATCH_SIZE: PositiveInt = Field( - description="Maximum number of records to process in each batch", - default=1000, - ) - SANDBOX_EXPIRED_RECORDS_RETENTION_DAYS: PositiveInt = Field( - description="Retention days for sandbox expired workflow_run records and message records", - default=30, - ) - - class PositionConfig(BaseSettings): POSITION_PROVIDER_PINS: str = Field( description="Comma-separated list of pinned model providers", diff --git a/api/services/billing_service.py b/api/services/billing_service.py index 2666956d46..3d7cb6cc8d 100644 --- a/api/services/billing_service.py +++ b/api/services/billing_service.py @@ -38,47 +38,6 @@ class BillingService: billing_info = cls._send_request("GET", "/subscription/info", params=params) return billing_info - @classmethod - def get_info_bulk(cls, tenant_ids: Sequence[str]) -> dict[str, TenantPlanInfo]: - """ - Bulk billing info fetch via billing API. - - Payload: {"tenant_ids": ["t1", "t2", ...]} (max 200 per request) - - Returns: - Mapping of tenant_id -> TenantPlanInfo(plan + expiration timestamp) - """ - results: dict[str, TenantPlanInfo] = {} - - chunk_size = 200 - for i in range(0, len(tenant_ids), chunk_size): - chunk = tenant_ids[i : i + chunk_size] - try: - resp = cls._send_request("POST", "/subscription/plan/batch", json={"tenant_ids": chunk}) - results.update(cls._parse_bulk_response(chunk, resp)) - except Exception: - logger.exception("Failed to fetch billing info batch for tenants: %s", chunk) - raise - - return results - - @classmethod - def _parse_bulk_response(cls, expected_ids: Sequence[str], response: dict) -> dict[str, TenantPlanInfo]: - data = response.get("data") - if not isinstance(data, dict): - raise ValueError("Billing API response missing 'data' object.") - - parsed: dict[str, TenantPlanInfo] = {} - for tenant_id in expected_ids: - payload = data.get(tenant_id) - - try: - parsed[tenant_id] = TenantPlanInfo.model_validate(payload) - except ValidationError as exc: - raise ValueError(f"Invalid billing info for tenant {tenant_id}") from exc - - return parsed - @classmethod def get_tenant_feature_plan_usage_info(cls, tenant_id: str): params = {"tenant_id": tenant_id} 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 ed829b853d..459cc5f46f 100644 --- a/api/services/clear_free_plan_expired_workflow_run_logs.py +++ b/api/services/clear_free_plan_expired_workflow_run_logs.py @@ -173,7 +173,7 @@ class WorkflowRunCleanup: if uncached_tenants: try: - bulk_info = BillingService.get_info_bulk(uncached_tenants) + bulk_info = BillingService.get_plan_bulk(uncached_tenants) except Exception: bulk_info = {} logger.exception("Failed to fetch billing plans in bulk for tenants: %s", uncached_tenants) diff --git a/docker/.env.example b/docker/.env.example index fd0def2b1a..fb24944e37 100644 --- a/docker/.env.example +++ b/docker/.env.example @@ -1468,11 +1468,6 @@ WORKFLOW_SCHEDULE_POLLER_INTERVAL=1 WORKFLOW_SCHEDULE_POLLER_BATCH_SIZE=100 WORKFLOW_SCHEDULE_MAX_DISPATCH_PER_TICK=0 -# Sandbox expired records clean configuration -SANDBOX_EXPIRED_RECORDS_CLEAN_GRACEFUL_PERIOD=21 -SANDBOX_EXPIRED_RECORDS_CLEAN_BATCH_SIZE=1000 -SANDBOX_EXPIRED_RECORDS_RETENTION_DAYS=30 - # Tenant isolated task queue configuration TENANT_ISOLATED_TASK_CONCURRENCY=1 diff --git a/docker/docker-compose.yaml b/docker/docker-compose.yaml index d90db26755..35c3ef1149 100644 --- a/docker/docker-compose.yaml +++ b/docker/docker-compose.yaml @@ -657,9 +657,6 @@ x-shared-env: &shared-api-worker-env WORKFLOW_SCHEDULE_POLLER_INTERVAL: ${WORKFLOW_SCHEDULE_POLLER_INTERVAL:-1} WORKFLOW_SCHEDULE_POLLER_BATCH_SIZE: ${WORKFLOW_SCHEDULE_POLLER_BATCH_SIZE:-100} WORKFLOW_SCHEDULE_MAX_DISPATCH_PER_TICK: ${WORKFLOW_SCHEDULE_MAX_DISPATCH_PER_TICK:-0} - SANDBOX_EXPIRED_RECORDS_CLEAN_GRACEFUL_PERIOD: ${SANDBOX_EXPIRED_RECORDS_CLEAN_GRACEFUL_PERIOD:-21} - SANDBOX_EXPIRED_RECORDS_CLEAN_BATCH_SIZE: ${SANDBOX_EXPIRED_RECORDS_CLEAN_BATCH_SIZE:-1000} - SANDBOX_EXPIRED_RECORDS_RETENTION_DAYS: ${SANDBOX_EXPIRED_RECORDS_RETENTION_DAYS:-30} TENANT_ISOLATED_TASK_CONCURRENCY: ${TENANT_ISOLATED_TASK_CONCURRENCY:-1} ANNOTATION_IMPORT_FILE_SIZE_LIMIT: ${ANNOTATION_IMPORT_FILE_SIZE_LIMIT:-2} ANNOTATION_IMPORT_MAX_RECORDS: ${ANNOTATION_IMPORT_MAX_RECORDS:-10000} From f02da9cae2e9aad1e3a91cd363faa5338bc78a53 Mon Sep 17 00:00:00 2001 From: hjlarry Date: Thu, 18 Dec 2025 15:40:38 +0800 Subject: [PATCH 37/42] fix CI --- api/services/clear_free_plan_expired_workflow_run_logs.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 459cc5f46f..6510aa27c9 100644 --- a/api/services/clear_free_plan_expired_workflow_run_logs.py +++ b/api/services/clear_free_plan_expired_workflow_run_logs.py @@ -14,7 +14,7 @@ from repositories.sqlalchemy_api_workflow_node_execution_repository import ( DifyAPISQLAlchemyWorkflowNodeExecutionRepository, ) from repositories.sqlalchemy_workflow_trigger_log_repository import SQLAlchemyWorkflowTriggerLogRepository -from services.billing_service import BillingService, TenantPlanInfo +from services.billing_service import BillingService, SubscriptionPlan logger = logging.getLogger(__name__) @@ -210,7 +210,7 @@ class WorkflowRunCleanup: logger.exception("Failed to parse expiration timestamp for tenant %s", tenant_id) return None - def _is_within_grace_period(self, tenant_id: str, info: TenantPlanInfo) -> bool: + def _is_within_grace_period(self, tenant_id: str, info: SubscriptionPlan) -> bool: if self.free_plan_grace_period_days <= 0: return False From a5647749aa710adcacc2398d522d287cf2a98922 Mon Sep 17 00:00:00 2001 From: hjlarry Date: Thu, 18 Dec 2025 16:00:55 +0800 Subject: [PATCH 38/42] fix CI --- .../clear_free_plan_expired_workflow_run_logs.py | 2 +- .../test_clear_free_plan_expired_workflow_run_logs.py | 10 +++++----- 2 files changed, 6 insertions(+), 6 deletions(-) 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 6510aa27c9..588a1fcf11 100644 --- a/api/services/clear_free_plan_expired_workflow_run_logs.py +++ b/api/services/clear_free_plan_expired_workflow_run_logs.py @@ -43,7 +43,7 @@ class WorkflowRunCleanup: raise ValueError("batch_size must be greater than 0.") self.batch_size = batch_size - self.billing_cache: dict[str, TenantPlanInfo | None] = {} + self.billing_cache: dict[str, SubscriptionPlan | 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 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 c2cfcd9811..5fb70811b8 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 @@ -4,7 +4,7 @@ from typing import Any import pytest from services import clear_free_plan_expired_workflow_run_logs as cleanup_module -from services.billing_service import TenantPlanInfo +from services.billing_service import SubscriptionPlan from services.clear_free_plan_expired_workflow_run_logs import WorkflowRunCleanup @@ -63,8 +63,8 @@ class FakeRepo: return result -def plan_info(plan: str, expiration: int) -> TenantPlanInfo: - return TenantPlanInfo(plan=plan, expiration_date=expiration) +def plan_info(plan: str, expiration: int) -> SubscriptionPlan: + return SubscriptionPlan(plan=plan, expiration_date=expiration) def create_cleanup( @@ -87,7 +87,7 @@ def test_filter_free_tenants_billing_disabled(monkeypatch: pytest.MonkeyPatch) - monkeypatch.setattr(cleanup_module.dify_config, "BILLING_ENABLED", False) - def fail_bulk(_: list[str]) -> dict[str, TenantPlanInfo]: + def fail_bulk(_: list[str]) -> dict[str, SubscriptionPlan]: raise RuntimeError("should not call") monkeypatch.setattr(cleanup_module.BillingService, "get_info_bulk", staticmethod(fail_bulk)) @@ -123,7 +123,7 @@ def test_filter_free_tenants_respects_grace_period(monkeypatch: pytest.MonkeyPat within_grace_ts = int((now - datetime.timedelta(days=10)).timestamp()) outside_grace_ts = int((now - datetime.timedelta(days=90)).timestamp()) - def fake_bulk(_: list[str]) -> dict[str, TenantPlanInfo]: + def fake_bulk(_: list[str]) -> dict[str, SubscriptionPlan]: return { "recently_downgraded": plan_info("sandbox", within_grace_ts), "long_sandbox": plan_info("sandbox", outside_grace_ts), From e82e1045f30884db5317e9ba6795ccef6edc8f77 Mon Sep 17 00:00:00 2001 From: hjlarry Date: Thu, 18 Dec 2025 16:19:52 +0800 Subject: [PATCH 39/42] fix CI --- .../clear_free_plan_expired_workflow_run_logs.py | 5 +++-- ...test_clear_free_plan_expired_workflow_run_logs.py | 12 ++++++------ 2 files changed, 9 insertions(+), 8 deletions(-) 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 588a1fcf11..a6962c6053 100644 --- a/api/services/clear_free_plan_expired_workflow_run_logs.py +++ b/api/services/clear_free_plan_expired_workflow_run_logs.py @@ -190,7 +190,7 @@ class WorkflowRunCleanup: if not info: continue - if info.plan != CloudPlan.SANDBOX: + if info.get("plan") != CloudPlan.SANDBOX: continue if self._is_within_grace_period(tenant_id, info): @@ -214,7 +214,8 @@ class WorkflowRunCleanup: if self.free_plan_grace_period_days <= 0: return False - expiration_at = self._expiration_datetime(tenant_id, info.expiration_date) + expiration_value = info.get("expiration_date", -1) + expiration_at = self._expiration_datetime(tenant_id, expiration_value) if expiration_at is None: return False 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 5fb70811b8..6f0fde2956 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 @@ -90,7 +90,7 @@ def test_filter_free_tenants_billing_disabled(monkeypatch: pytest.MonkeyPatch) - def fail_bulk(_: list[str]) -> dict[str, SubscriptionPlan]: raise RuntimeError("should not call") - monkeypatch.setattr(cleanup_module.BillingService, "get_info_bulk", staticmethod(fail_bulk)) + monkeypatch.setattr(cleanup_module.BillingService, "get_plan_bulk", staticmethod(fail_bulk)) tenants = {"t1", "t2"} free = cleanup._filter_free_tenants(tenants) @@ -106,7 +106,7 @@ def test_filter_free_tenants_bulk_mixed(monkeypatch: pytest.MonkeyPatch) -> None cleanup.billing_cache["t_paid"] = plan_info("team", -1) monkeypatch.setattr( cleanup_module.BillingService, - "get_info_bulk", + "get_plan_bulk", staticmethod(lambda tenant_ids: {tenant_id: plan_info("sandbox", -1) for tenant_id in tenant_ids}), ) @@ -129,7 +129,7 @@ def test_filter_free_tenants_respects_grace_period(monkeypatch: pytest.MonkeyPat "long_sandbox": plan_info("sandbox", outside_grace_ts), } - monkeypatch.setattr(cleanup_module.BillingService, "get_info_bulk", staticmethod(fake_bulk)) + monkeypatch.setattr(cleanup_module.BillingService, "get_plan_bulk", staticmethod(fake_bulk)) free = cleanup._filter_free_tenants({"recently_downgraded", "long_sandbox"}) @@ -142,7 +142,7 @@ def test_filter_free_tenants_bulk_failure(monkeypatch: pytest.MonkeyPatch) -> No monkeypatch.setattr(cleanup_module.dify_config, "BILLING_ENABLED", True) monkeypatch.setattr( cleanup_module.BillingService, - "get_info_bulk", + "get_plan_bulk", staticmethod(lambda tenant_ids: (_ for _ in ()).throw(RuntimeError("boom"))), ) @@ -168,7 +168,7 @@ def test_run_deletes_only_free_tenants(monkeypatch: pytest.MonkeyPatch) -> None: cleanup.billing_cache["t_paid"] = plan_info("team", -1) monkeypatch.setattr( cleanup_module.BillingService, - "get_info_bulk", + "get_plan_bulk", staticmethod(lambda tenant_ids: {tenant_id: plan_info("sandbox", -1) for tenant_id in tenant_ids}), ) @@ -185,7 +185,7 @@ def test_run_skips_when_no_free_tenants(monkeypatch: pytest.MonkeyPatch) -> None monkeypatch.setattr(cleanup_module.dify_config, "BILLING_ENABLED", True) monkeypatch.setattr( cleanup_module.BillingService, - "get_info_bulk", + "get_plan_bulk", staticmethod(lambda tenant_ids: {tenant_id: plan_info("team", 1893456000) for tenant_id in tenant_ids}), ) From 3bcaabceae297f0db04a5606d12a63b964e21efa Mon Sep 17 00:00:00 2001 From: hjlarry Date: Thu, 18 Dec 2025 16:31:11 +0800 Subject: [PATCH 40/42] use (created_at,id) for index --- ...d7c23e_add_workflow_runs_created_at_idx.py | 27 ---------------- .../versions/2025_12_18_1630-905527cc8fd3_.py | 32 +++++++++++++++++++ api/models/workflow.py | 1 + 3 files changed, 33 insertions(+), 27 deletions(-) delete mode 100644 api/migrations/versions/2025_12_17_1504-8a7f2ad7c23e_add_workflow_runs_created_at_idx.py create mode 100644 api/migrations/versions/2025_12_18_1630-905527cc8fd3_.py diff --git a/api/migrations/versions/2025_12_17_1504-8a7f2ad7c23e_add_workflow_runs_created_at_idx.py b/api/migrations/versions/2025_12_17_1504-8a7f2ad7c23e_add_workflow_runs_created_at_idx.py deleted file mode 100644 index 2acd00fc47..0000000000 --- a/api/migrations/versions/2025_12_17_1504-8a7f2ad7c23e_add_workflow_runs_created_at_idx.py +++ /dev/null @@ -1,27 +0,0 @@ -"""Add index on workflow_runs.created_at - -Revision ID: 8a7f2ad7c23e -Revises: 03ea244985ce -Create Date: 2025-12-17 15:04:00.000000 -""" - -from alembic import op -# revision identifiers, used by Alembic. -revision = "8a7f2ad7c23e" -down_revision = "03ea244985ce" -branch_labels = None -depends_on = None - - -def upgrade(): - with op.batch_alter_table("workflow_runs", schema=None) as batch_op: - batch_op.create_index( - batch_op.f("workflow_runs_created_at_idx"), - ["created_at"], - unique=False, - ) - - -def downgrade(): - with op.batch_alter_table("workflow_runs", schema=None) as batch_op: - batch_op.drop_index(batch_op.f("workflow_runs_created_at_idx")) diff --git a/api/migrations/versions/2025_12_18_1630-905527cc8fd3_.py b/api/migrations/versions/2025_12_18_1630-905527cc8fd3_.py new file mode 100644 index 0000000000..d7a59c281e --- /dev/null +++ b/api/migrations/versions/2025_12_18_1630-905527cc8fd3_.py @@ -0,0 +1,32 @@ +"""empty message + +Revision ID: 905527cc8fd3 +Revises: 03ea244985ce +Create Date: 2025-12-18 16:30:02.462084 + +""" +from alembic import op +import models as models +import sqlalchemy as sa +from sqlalchemy.dialects import postgresql + +# revision identifiers, used by Alembic. +revision = '905527cc8fd3' +down_revision = '03ea244985ce' +branch_labels = None +depends_on = None + + +def upgrade(): + # ### commands auto generated by Alembic - please adjust! ### + with op.batch_alter_table('workflow_runs', schema=None) as batch_op: + batch_op.create_index('workflow_run_created_at_id_idx', ['created_at', 'id'], unique=False) + + # ### end Alembic commands ### + + +def downgrade(): + # ### commands auto generated by Alembic - please adjust! ### + with op.batch_alter_table('workflow_runs', schema=None) as batch_op: + batch_op.drop_index('workflow_run_created_at_id_idx') + # ### end Alembic commands ### diff --git a/api/models/workflow.py b/api/models/workflow.py index 853d5afefc..90c99d49d6 100644 --- a/api/models/workflow.py +++ b/api/models/workflow.py @@ -595,6 +595,7 @@ class WorkflowRun(Base): __table_args__ = ( sa.PrimaryKeyConstraint("id", name="workflow_run_pkey"), sa.Index("workflow_run_triggerd_from_idx", "tenant_id", "app_id", "triggered_from"), + sa.Index("workflow_run_created_at_id_idx", "created_at", "id"), ) id: Mapped[str] = mapped_column(StringUUID, default=lambda: str(uuid4())) From 0d30eef055d1e69e9372e0e8c330407cc8eca9fd Mon Sep 17 00:00:00 2001 From: hjlarry Date: Thu, 18 Dec 2025 16:45:32 +0800 Subject: [PATCH 41/42] add clean up whitelist --- ...ear_free_plan_expired_workflow_run_logs.py | 23 +++++++++++++++ ...ear_free_plan_expired_workflow_run_logs.py | 29 +++++++++++++++++++ 2 files changed, 52 insertions(+) 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 a6962c6053..5b01a9caf6 100644 --- a/api/services/clear_free_plan_expired_workflow_run_logs.py +++ b/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) 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 6f0fde2956..d0921c1b5d 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 @@ -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) From ba389dd5c1119a63195f3421684de2cebc78a033 Mon Sep 17 00:00:00 2001 From: hjlarry Date: Wed, 24 Dec 2025 09:19:56 +0800 Subject: [PATCH 42/42] fix migrate msg --- api/migrations/versions/2025_12_18_1630-905527cc8fd3_.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/api/migrations/versions/2025_12_18_1630-905527cc8fd3_.py b/api/migrations/versions/2025_12_18_1630-905527cc8fd3_.py index d7a59c281e..e8b20674e8 100644 --- a/api/migrations/versions/2025_12_18_1630-905527cc8fd3_.py +++ b/api/migrations/versions/2025_12_18_1630-905527cc8fd3_.py @@ -1,4 +1,4 @@ -"""empty message +"""add workflow_run_created_at_id_idx Revision ID: 905527cc8fd3 Revises: 03ea244985ce