Compare commits

...
9 changed files with 1619 additions and 84 deletions
+67 -5
View File
@@ -10,7 +10,7 @@ from extensions.ext_database import db
from libs.datetime_utils import naive_utc_now
from services.clear_free_plan_tenant_expired_logs import ClearFreePlanTenantExpiredLogs
from services.retention.conversation.messages_clean_policy import create_message_clean_policy
from services.retention.conversation.messages_clean_service import MessagesCleanService
from services.retention.conversation.messages_clean_service import MessagesCleanScanStrategy, MessagesCleanService
from services.retention.workflow_run.clear_free_plan_expired_workflow_run_logs import WorkflowRunCleanup
from tasks.remove_app_and_related_data_task import delete_draft_variables_batch
@@ -664,7 +664,44 @@ def cleanup_orphaned_draft_variables(
default=None,
help="Relative upper bound in days ago (exclusive). Required for relative mode.",
)
@click.option("--batch-size", default=1000, show_default=True, help="Batch size for selecting messages.")
@click.option("--batch-size", default=1000, show_default=True, help="Backward-compatible default batch size.")
@click.option(
"--candidate-batch-size",
default=None,
type=int,
help="Initial number of global candidate messages to scan per batch. Defaults to --batch-size.",
)
@click.option(
"--max-candidate-batch-size",
default=None,
type=int,
help="Maximum number of global candidate messages to scan per batch.",
)
@click.option(
"--delete-batch-size",
default=None,
type=int,
help="Maximum number of messages to delete per transaction. Defaults to --batch-size.",
)
@click.option(
"--per-app-batch-size",
default=None,
type=int,
help="Maximum number of messages to scan per eligible app turn. Defaults to --batch-size.",
)
@click.option(
"--app-page-size",
default=500,
show_default=True,
help="Number of apps to inspect per eligible-app discovery query.",
)
@click.option(
"--scan-strategy",
type=click.Choice(["auto", "global", "eligible_apps"]),
default="auto",
show_default=True,
help="Message cleanup scan strategy.",
)
@click.option(
"--graceful-period",
default=21,
@@ -674,6 +711,12 @@ def cleanup_orphaned_draft_variables(
@click.option("--dry-run", is_flag=True, default=False, help="Show messages logs would be cleaned without deleting")
def clean_expired_messages(
batch_size: int,
candidate_batch_size: int | None,
max_candidate_batch_size: int | None,
delete_batch_size: int | None,
per_app_batch_size: int | None,
app_page_size: int,
scan_strategy: MessagesCleanScanStrategy,
graceful_period: int,
start_from: datetime.datetime | None,
end_before: datetime.datetime | None,
@@ -732,6 +775,10 @@ def clean_expired_messages(
else:
task_label = "custom"
effective_candidate_batch_size = candidate_batch_size if candidate_batch_size is not None else batch_size
effective_delete_batch_size = delete_batch_size if delete_batch_size is not None else batch_size
effective_per_app_batch_size = per_app_batch_size if per_app_batch_size is not None else batch_size
# Create and run the cleanup service
if abs_mode:
assert start_from is not None
@@ -740,7 +787,12 @@ def clean_expired_messages(
policy=policy,
start_from=start_from,
end_before=end_before,
batch_size=batch_size,
batch_size=effective_candidate_batch_size,
max_candidate_batch_size=max_candidate_batch_size,
delete_batch_size=effective_delete_batch_size,
per_app_batch_size=effective_per_app_batch_size,
app_page_size=app_page_size,
scan_strategy=scan_strategy,
dry_run=dry_run,
task_label=task_label,
)
@@ -749,7 +801,12 @@ def clean_expired_messages(
service = MessagesCleanService.from_days(
policy=policy,
days=before_days,
batch_size=batch_size,
batch_size=effective_candidate_batch_size,
max_candidate_batch_size=max_candidate_batch_size,
delete_batch_size=effective_delete_batch_size,
per_app_batch_size=effective_per_app_batch_size,
app_page_size=app_page_size,
scan_strategy=scan_strategy,
dry_run=dry_run,
task_label=task_label,
)
@@ -761,7 +818,12 @@ def clean_expired_messages(
policy=policy,
start_from=now - datetime.timedelta(days=from_days_ago),
end_before=now - datetime.timedelta(days=before_days),
batch_size=batch_size,
batch_size=effective_candidate_batch_size,
max_candidate_batch_size=max_candidate_batch_size,
delete_batch_size=effective_delete_batch_size,
per_app_batch_size=effective_per_app_batch_size,
app_page_size=app_page_size,
scan_strategy=scan_strategy,
dry_run=dry_run,
task_label=task_label,
)
+24
View File
@@ -1460,6 +1460,30 @@ class SandboxExpiredRecordsCleanConfig(BaseSettings):
description="Maximum number of records to process in each batch",
default=1000,
)
SANDBOX_EXPIRED_RECORDS_CLEAN_CANDIDATE_BATCH_SIZE: PositiveInt = Field(
description="Initial number of message candidates to scan in each expired records cleanup batch",
default=10000,
)
SANDBOX_EXPIRED_RECORDS_CLEAN_CANDIDATE_BATCH_MAX_SIZE: PositiveInt = Field(
description="Maximum number of message candidates to scan in each expired records cleanup batch",
default=50000,
)
SANDBOX_EXPIRED_RECORDS_CLEAN_DELETE_BATCH_SIZE: PositiveInt = Field(
description="Maximum number of records to delete in each expired records cleanup transaction",
default=1000,
)
SANDBOX_EXPIRED_RECORDS_CLEAN_PER_APP_BATCH_SIZE: PositiveInt = Field(
description="Maximum number of messages to scan per eligible app turn in each expired records cleanup batch",
default=1000,
)
SANDBOX_EXPIRED_RECORDS_CLEAN_APP_PAGE_SIZE: PositiveInt = Field(
description="Maximum number of apps to inspect per expired records cleanup eligibility discovery query",
default=500,
)
SANDBOX_EXPIRED_RECORDS_CLEAN_SCAN_STRATEGY: str = Field(
description="Message cleanup scan strategy: auto, global, or eligible_apps",
default="global",
)
SANDBOX_EXPIRED_RECORDS_CLEAN_BATCH_MAX_INTERVAL: PositiveInt = Field(
description="Maximum interval in milliseconds between batches",
default=200,
+6 -1
View File
@@ -40,7 +40,12 @@ def clean_messages():
service = MessagesCleanService.from_days(
policy=policy,
days=dify_config.SANDBOX_EXPIRED_RECORDS_RETENTION_DAYS,
batch_size=dify_config.SANDBOX_EXPIRED_RECORDS_CLEAN_BATCH_SIZE,
batch_size=dify_config.SANDBOX_EXPIRED_RECORDS_CLEAN_CANDIDATE_BATCH_SIZE,
max_candidate_batch_size=dify_config.SANDBOX_EXPIRED_RECORDS_CLEAN_CANDIDATE_BATCH_MAX_SIZE,
delete_batch_size=dify_config.SANDBOX_EXPIRED_RECORDS_CLEAN_DELETE_BATCH_SIZE,
per_app_batch_size=dify_config.SANDBOX_EXPIRED_RECORDS_CLEAN_PER_APP_BATCH_SIZE,
app_page_size=dify_config.SANDBOX_EXPIRED_RECORDS_CLEAN_APP_PAGE_SIZE,
scan_strategy=dify_config.SANDBOX_EXPIRED_RECORDS_CLEAN_SCAN_STRATEGY,
)
stats = service.run()
@@ -0,0 +1,180 @@
import datetime
from collections import deque
from collections.abc import Sequence
from dataclasses import dataclass
from typing import cast
import sqlalchemy as sa
from sqlalchemy import select, tuple_
from sqlalchemy.orm import Session, sessionmaker
from models.model import App, Message
from services.retention.conversation.messages_clean_policy import (
EligibleAppMessagesCleanPolicy,
SimpleMessage,
)
@dataclass
class AppCursorState:
app_id: str
tenant_id: str
cursor: tuple[datetime.datetime, str] | None = None
@dataclass
class EligibleAppScanBatch:
messages: list[SimpleMessage]
app_to_tenant: dict[str, str]
app_fetches: int
exhausted_apps: int
class EligibleAppRoundRobinScanner:
"""
Scans messages only for apps whose tenants are eligible for sandbox cleanup.
App discovery walks the App table in stable id order and asks the policy to filter eligible tenants.
Message scanning then uses a per-app (created_at, id) cursor and puts non-exhausted apps back at the
end of the queue, so a single large eligible app cannot monopolize the whole job.
"""
_policy: EligibleAppMessagesCleanPolicy
_start_from: datetime.datetime | None
_end_before: datetime.datetime
_app_page_size: int
_per_app_batch_size: int
_app_queue: deque[AppCursorState]
_last_seen_app_id: str | None
_app_pages_exhausted: bool
scanned_apps: int
eligible_apps: int
empty_apps: int
def __init__(
self,
*,
policy: EligibleAppMessagesCleanPolicy,
start_from: datetime.datetime | None,
end_before: datetime.datetime,
app_page_size: int,
per_app_batch_size: int,
) -> None:
self._policy = policy
self._start_from = start_from
self._end_before = end_before
self._app_page_size = app_page_size
self._per_app_batch_size = per_app_batch_size
self._app_queue = deque()
self._last_seen_app_id = None
self._app_pages_exhausted = False
self.scanned_apps = 0
self.eligible_apps = 0
self.empty_apps = 0
def fetch_batch(
self,
session_factory: sessionmaker[Session],
*,
target_message_count: int,
) -> EligibleAppScanBatch:
messages: list[SimpleMessage] = []
app_to_tenant: dict[str, str] = {}
app_fetches = 0
exhausted_apps = 0
while len(messages) < target_message_count and self._ensure_app_queue(session_factory):
app_state = self._app_queue.popleft()
fetch_limit = min(self._per_app_batch_size, target_message_count - len(messages))
with session_factory() as session:
fetched_messages = self._fetch_messages_for_app(session, app_state, limit=fetch_limit)
app_fetches += 1
if not fetched_messages:
exhausted_apps += 1
self.empty_apps += 1
continue
messages.extend(fetched_messages)
app_to_tenant[app_state.app_id] = app_state.tenant_id
app_state.cursor = (fetched_messages[-1].created_at, fetched_messages[-1].id)
if len(fetched_messages) == fetch_limit:
self._app_queue.append(app_state)
else:
exhausted_apps += 1
return EligibleAppScanBatch(
messages=messages,
app_to_tenant=app_to_tenant,
app_fetches=app_fetches,
exhausted_apps=exhausted_apps,
)
def _ensure_app_queue(self, session_factory: sessionmaker[Session]) -> bool:
while not self._app_queue and not self._app_pages_exhausted:
rows = self._fetch_next_app_page(session_factory)
if rows:
self._enqueue_eligible_apps(rows)
return bool(self._app_queue)
def _fetch_next_app_page(self, session_factory: sessionmaker[Session]) -> list[tuple[str, str]]:
app_stmt = select(App.id, App.tenant_id).order_by(App.id).limit(self._app_page_size)
if self._last_seen_app_id is not None:
app_stmt = app_stmt.where(App.id > self._last_seen_app_id)
with session_factory() as session:
rows = cast(list[tuple[str, str]], list(session.execute(app_stmt).all()))
if not rows:
self._app_pages_exhausted = True
return []
self._last_seen_app_id = rows[-1][0]
self.scanned_apps += len(rows)
return rows
def _enqueue_eligible_apps(self, rows: Sequence[tuple[str, str]]) -> None:
app_to_tenant: dict[str, str] = dict(rows)
eligible_app_to_tenant = self._policy.filter_app_to_tenant(app_to_tenant)
self.eligible_apps += len(eligible_app_to_tenant)
for app_id, tenant_id in rows:
if eligible_app_to_tenant.get(app_id) == tenant_id:
self._app_queue.append(AppCursorState(app_id=app_id, tenant_id=tenant_id))
def _fetch_messages_for_app(
self,
session: Session,
app_state: AppCursorState,
*,
limit: int,
) -> list[SimpleMessage]:
msg_stmt = (
select(Message.id, Message.app_id, Message.created_at)
.where(
Message.app_id == app_state.app_id,
Message.created_at < self._end_before,
)
.order_by(Message.created_at, Message.id)
.limit(limit)
)
if self._start_from:
msg_stmt = msg_stmt.where(Message.created_at >= self._start_from)
if app_state.cursor:
msg_stmt = msg_stmt.where(
tuple_(Message.created_at, Message.id)
> tuple_(
sa.literal(app_state.cursor[0], type_=sa.DateTime()),
sa.literal(app_state.cursor[1], type_=Message.id.type),
)
)
rows = cast(Sequence[tuple[str, str, datetime.datetime]], session.execute(msg_stmt).all())
return self._rows_to_simple_messages(rows)
@staticmethod
def _rows_to_simple_messages(rows: Sequence[tuple[str, str, datetime.datetime]]) -> list[SimpleMessage]:
return [SimpleMessage(id=msg_id, app_id=app_id, created_at=created_at) for msg_id, app_id, created_at in rows]
@@ -2,7 +2,7 @@ import datetime
import logging
from collections.abc import Callable, Sequence
from dataclasses import dataclass
from typing import Protocol, override
from typing import Protocol, override, runtime_checkable
from configs import dify_config
from enums.cloud_plan import CloudPlan
@@ -43,6 +43,35 @@ class MessagesCleanPolicy(Protocol):
...
@runtime_checkable
class EligibleAppMessagesCleanPolicy(MessagesCleanPolicy, Protocol):
"""
Policy extension for cleanup strategies that discover eligible apps before scanning messages.
Discovery may use a job-level plan cache to reduce billing traffic, but delete-time revalidation
must fetch fresh plan data so tenants upgraded during a long cleanup run are not removed by stale state.
"""
def filter_app_to_tenant(
self,
app_to_tenant: dict[str, str],
) -> dict[str, str]:
"""
Return only apps whose current cached tenant plan is eligible for cleanup.
"""
...
def revalidate_message_ids(
self,
messages: Sequence[SimpleMessage],
app_to_tenant: dict[str, str],
) -> Sequence[str]:
"""
Re-check message tenant plans without the job-level cache immediately before deletion.
"""
...
class BillingDisabledPolicy(MessagesCleanPolicy):
"""
Policy for community or enterpriseedition (billing disabled).
@@ -70,17 +99,29 @@ class BillingSandboxPolicy(MessagesCleanPolicy):
- Safe default: if tenant mapping or plan is missing, do NOT delete
"""
_graceful_period_days: int
_tenant_whitelist: Sequence[str]
_tenant_whitelist_set: set[str]
_plan_provider: Callable[[Sequence[str]], dict[str, SubscriptionPlan]]
_fresh_plan_provider: Callable[[Sequence[str]], dict[str, SubscriptionPlan]]
_current_timestamp: int | None
_plan_cache: dict[str, SubscriptionPlan | None]
def __init__(
self,
plan_provider: Callable[[Sequence[str]], dict[str, SubscriptionPlan]],
fresh_plan_provider: Callable[[Sequence[str]], dict[str, SubscriptionPlan]] | None = None,
graceful_period_days: int = 21,
tenant_whitelist: Sequence[str] | None = None,
current_timestamp: int | None = None,
) -> None:
self._graceful_period_days = graceful_period_days
self._tenant_whitelist: Sequence[str] = tenant_whitelist or []
self._tenant_whitelist_set = set(self._tenant_whitelist)
self._plan_provider = plan_provider
self._fresh_plan_provider = fresh_plan_provider or plan_provider
self._current_timestamp = current_timestamp
self._plan_cache = {}
@override
def filter_message_ids(
@@ -101,9 +142,12 @@ class BillingSandboxPolicy(MessagesCleanPolicy):
if not messages or not app_to_tenant:
return []
# Get unique tenant_ids and fetch subscription plans
tenant_ids = list(set(app_to_tenant.values()))
tenant_plans = self._plan_provider(tenant_ids)
# Get unique tenant_ids and fetch subscription plans. Plans are cached for the whole
# policy lifetime because message cleanup evaluates many adjacent batches from the same apps.
tenant_ids = sorted(
{tenant_id for tenant_id in app_to_tenant.values() if tenant_id not in self._tenant_whitelist_set}
)
tenant_plans = self._get_tenant_plans(tenant_ids)
if not tenant_plans:
return []
@@ -115,6 +159,119 @@ class BillingSandboxPolicy(MessagesCleanPolicy):
tenant_plans=tenant_plans,
)
def filter_app_to_tenant(
self,
app_to_tenant: dict[str, str],
) -> dict[str, str]:
"""
Return apps that belong to tenants currently eligible for sandbox cleanup.
This method is used during app discovery and can use the job-level plan cache. Deletion still calls
revalidate_message_ids(), which bypasses this cache.
"""
if not app_to_tenant:
return {}
eligible_tenant_ids = self._eligible_tenant_ids(
list(app_to_tenant.values()),
use_fresh_provider=False,
)
return {app_id: tenant_id for app_id, tenant_id in app_to_tenant.items() if tenant_id in eligible_tenant_ids}
def revalidate_message_ids(
self,
messages: Sequence[SimpleMessage],
app_to_tenant: dict[str, str],
) -> Sequence[str]:
"""
Re-check sandbox eligibility with fresh billing data immediately before deletion.
"""
if not messages or not app_to_tenant:
return []
tenant_ids = sorted(
{tenant_id for tenant_id in app_to_tenant.values() if tenant_id not in self._tenant_whitelist_set}
)
tenant_plans = self._get_fresh_tenant_plans(tenant_ids)
if not tenant_plans:
return []
return self._filter_expired_sandbox_messages(
messages=messages,
app_to_tenant=app_to_tenant,
tenant_plans=tenant_plans,
)
def _eligible_tenant_ids(self, tenant_ids: Sequence[str], *, use_fresh_provider: bool) -> set[str]:
unique_tenant_ids = sorted(set(tenant_ids) - self._tenant_whitelist_set)
if not unique_tenant_ids:
return set()
tenant_plans = (
self._get_fresh_tenant_plans(unique_tenant_ids)
if use_fresh_provider
else self._get_tenant_plans(unique_tenant_ids)
)
return set(
self._tenant_ids_for_expired_sandbox_plans(
tenant_plans,
)
)
def _get_tenant_plans(self, tenant_ids: Sequence[str]) -> dict[str, SubscriptionPlan]:
"""
Return cached subscription plans for tenant ids.
Missing billing responses are cached as None and remain a safe non-delete decision.
Provider exceptions still propagate so transient billing failures do not silently change cleanup behavior.
"""
unique_tenant_ids = sorted(set(tenant_ids))
missing_tenant_ids = [tenant_id for tenant_id in unique_tenant_ids if tenant_id not in self._plan_cache]
if missing_tenant_ids:
fetched_plans = self._plan_provider(missing_tenant_ids)
for tenant_id in missing_tenant_ids:
self._plan_cache[tenant_id] = fetched_plans.get(tenant_id)
plans: dict[str, SubscriptionPlan] = {}
for tenant_id in unique_tenant_ids:
plan = self._plan_cache.get(tenant_id)
if plan is not None:
plans[tenant_id] = plan
return plans
def _get_fresh_tenant_plans(self, tenant_ids: Sequence[str]) -> dict[str, SubscriptionPlan]:
"""
Fetch subscription plans without the job-level cache.
This is intentionally used on the delete path to avoid deleting messages for tenants that upgraded
while a long-running cleanup job was scanning.
"""
unique_tenant_ids = sorted(set(tenant_ids))
if not unique_tenant_ids:
return {}
return self._fresh_plan_provider(unique_tenant_ids)
def _tenant_ids_for_expired_sandbox_plans(self, tenant_plans: dict[str, SubscriptionPlan]) -> list[str]:
current_timestamp = self._current_timestamp
if current_timestamp is None:
current_timestamp = int(datetime.datetime.now(datetime.UTC).timestamp())
eligible_tenant_ids: list[str] = []
graceful_period_seconds = self._graceful_period_days * 24 * 60 * 60
for tenant_id, tenant_plan in tenant_plans.items():
plan = str(tenant_plan["plan"])
expiration_date = int(tenant_plan["expiration_date"])
if plan != CloudPlan.SANDBOX:
continue
if expiration_date == -1 or current_timestamp - expiration_date > graceful_period_seconds:
eligible_tenant_ids.append(tenant_id)
return eligible_tenant_ids
def _filter_expired_sandbox_messages(
self,
messages: Sequence[SimpleMessage],
@@ -152,7 +309,7 @@ class BillingSandboxPolicy(MessagesCleanPolicy):
continue
# Skip tenant messages in whitelist
if tenant_id in self._tenant_whitelist:
if tenant_id in self._tenant_whitelist_set:
continue
# Get subscription plan for this tenant
@@ -201,6 +358,7 @@ def create_message_clean_policy(
# Billing enabled - fetch whitelist from BillingService
tenant_whitelist = BillingService.get_expired_subscription_cleanup_whitelist()
plan_provider = BillingService.get_plan_bulk_with_cache
fresh_plan_provider = BillingService.get_plan_bulk
logger.info(
"create_message_clean_policy: billing enabled, using BillingSandboxPolicy "
@@ -211,6 +369,7 @@ def create_message_clean_policy(
return BillingSandboxPolicy(
plan_provider=plan_provider,
fresh_plan_provider=fresh_plan_provider,
graceful_period_days=graceful_period_days,
tenant_whitelist=tenant_whitelist,
current_timestamp=current_timestamp,
@@ -1,9 +1,10 @@
import datetime
import logging
import math
import random
import time
from collections.abc import Sequence
from typing import TYPE_CHECKING, TypedDict, cast
from collections.abc import Iterator, Sequence
from typing import TYPE_CHECKING, Literal, TypedDict, cast
import sqlalchemy as sa
from sqlalchemy import delete, select, tuple_
@@ -25,7 +26,9 @@ from models.model import (
MessageFile,
)
from models.web import SavedMessage
from services.retention.conversation.messages_clean_app_scan import EligibleAppRoundRobinScanner
from services.retention.conversation.messages_clean_policy import (
EligibleAppMessagesCleanPolicy,
MessagesCleanPolicy,
SimpleMessage,
)
@@ -37,6 +40,11 @@ if TYPE_CHECKING:
from opentelemetry.metrics import Counter, Histogram
_MIN_ELIGIBLE_HIT_RATE = 0.005
_HIT_RATE_EMA_ALPHA = 0.3
MessagesCleanScanStrategy = Literal["auto", "global", "eligible_apps"]
class MessagesCleanupMetrics:
"""
Records low-cardinality OpenTelemetry metrics for expired message cleanup jobs.
@@ -165,12 +173,24 @@ class MessagesCleanStatsDict(TypedDict):
total_deleted: int
class MessageDeleteResultDict(TypedDict):
messages_deleted: int
chunks: int
relations_ms: int
messages_ms: int
class MessagesCleanService:
"""
Service for cleaning expired messages based on retention policies.
Compatible with non cloud edition (billing disabled): all messages in the time range will be deleted.
If billing is enabled: only sandbox plan tenant messages are deleted (with whitelist and grace period support).
Billing-disabled cleanup uses the global message cursor because every message in the time range is eligible.
Billing-enabled cleanup can use an eligible-app round-robin strategy: discover apps whose tenants are sandbox
and past grace, scan messages through per-app cursors, then revalidate tenant plans immediately before deletion.
The global cursor remains available as a rollback strategy and for non-billing policies.
"""
def __init__(
@@ -179,6 +199,11 @@ class MessagesCleanService:
end_before: datetime.datetime,
start_from: datetime.datetime | None = None,
batch_size: int = 1000,
max_candidate_batch_size: int | None = None,
delete_batch_size: int | None = None,
per_app_batch_size: int | None = None,
app_page_size: int = 500,
scan_strategy: MessagesCleanScanStrategy = "auto",
dry_run: bool = False,
task_label: str = "custom",
) -> None:
@@ -189,14 +214,43 @@ class MessagesCleanService:
policy: The policy that determines which messages to delete
end_before: End time (exclusive) of the range
start_from: Optional start time (inclusive) of the range
batch_size: Number of messages to process per batch
batch_size: Initial number of candidate messages to scan per batch
max_candidate_batch_size: Maximum number of candidate messages to scan per batch
delete_batch_size: Maximum number of messages to delete per transaction
per_app_batch_size: Maximum number of messages to fetch per eligible app turn
app_page_size: Number of apps to inspect per eligibility discovery query
scan_strategy: "auto", "global", or "eligible_apps"
dry_run: Whether to perform a dry run (no actual deletion)
task_label: Optional task label for retention metrics
"""
if batch_size <= 0:
raise ValueError(f"batch_size ({batch_size}) must be greater than 0")
if max_candidate_batch_size is not None and max_candidate_batch_size <= 0:
raise ValueError(f"max_candidate_batch_size ({max_candidate_batch_size}) must be greater than 0")
if delete_batch_size is not None and delete_batch_size <= 0:
raise ValueError(f"delete_batch_size ({delete_batch_size}) must be greater than 0")
if per_app_batch_size is not None and per_app_batch_size <= 0:
raise ValueError(f"per_app_batch_size ({per_app_batch_size}) must be greater than 0")
if app_page_size <= 0:
raise ValueError(f"app_page_size ({app_page_size}) must be greater than 0")
if scan_strategy not in ("auto", "global", "eligible_apps"):
raise ValueError(f"scan_strategy ({scan_strategy}) must be one of: auto, global, eligible_apps")
self._policy = policy
self._end_before = end_before
self._start_from = start_from
self._batch_size = batch_size
self._candidate_batch_size = batch_size
self._delete_batch_size = delete_batch_size or batch_size
self._max_candidate_batch_size = max(max_candidate_batch_size or batch_size, self._candidate_batch_size)
self._per_app_batch_size = per_app_batch_size or batch_size
self._app_page_size = app_page_size
self._scan_strategy = scan_strategy
self._dry_run = dry_run
self._metrics = MessagesCleanupMetrics(
dry_run=dry_run,
@@ -211,6 +265,11 @@ class MessagesCleanService:
start_from: datetime.datetime,
end_before: datetime.datetime,
batch_size: int = 1000,
max_candidate_batch_size: int | None = None,
delete_batch_size: int | None = None,
per_app_batch_size: int | None = None,
app_page_size: int = 500,
scan_strategy: MessagesCleanScanStrategy = "auto",
dry_run: bool = False,
task_label: str = "custom",
) -> "MessagesCleanService":
@@ -223,7 +282,12 @@ class MessagesCleanService:
policy: The policy that determines which messages to delete
start_from: Start time (inclusive) of the range
end_before: End time (exclusive) of the range
batch_size: Number of messages to process per batch
batch_size: Initial number of candidate messages to scan per batch
max_candidate_batch_size: Maximum number of candidate messages to scan per batch
delete_batch_size: Maximum number of messages to delete per transaction
per_app_batch_size: Maximum number of messages to fetch per eligible app turn
app_page_size: Number of apps to inspect per eligibility discovery query
scan_strategy: "auto", "global", or "eligible_apps"
dry_run: Whether to perform a dry run (no actual deletion)
task_label: Optional task label for retention metrics
@@ -239,11 +303,30 @@ class MessagesCleanService:
if batch_size <= 0:
raise ValueError(f"batch_size ({batch_size}) must be greater than 0")
if max_candidate_batch_size is not None and max_candidate_batch_size <= 0:
raise ValueError(f"max_candidate_batch_size ({max_candidate_batch_size}) must be greater than 0")
if delete_batch_size is not None and delete_batch_size <= 0:
raise ValueError(f"delete_batch_size ({delete_batch_size}) must be greater than 0")
if per_app_batch_size is not None and per_app_batch_size <= 0:
raise ValueError(f"per_app_batch_size ({per_app_batch_size}) must be greater than 0")
if app_page_size <= 0:
raise ValueError(f"app_page_size ({app_page_size}) must be greater than 0")
logger.info(
"clean_messages: start_from=%s, end_before=%s, batch_size=%s, policy=%s",
"clean_messages: start_from=%s, end_before=%s, batch_size=%s, "
"max_candidate_batch_size=%s, delete_batch_size=%s, per_app_batch_size=%s, "
"app_page_size=%s, scan_strategy=%s, policy=%s",
start_from,
end_before,
batch_size,
max_candidate_batch_size,
delete_batch_size,
per_app_batch_size,
app_page_size,
scan_strategy,
policy.__class__.__name__,
)
@@ -252,6 +335,11 @@ class MessagesCleanService:
end_before=end_before,
start_from=start_from,
batch_size=batch_size,
max_candidate_batch_size=max_candidate_batch_size,
delete_batch_size=delete_batch_size,
per_app_batch_size=per_app_batch_size,
app_page_size=app_page_size,
scan_strategy=scan_strategy,
dry_run=dry_run,
task_label=task_label,
)
@@ -262,6 +350,11 @@ class MessagesCleanService:
policy: MessagesCleanPolicy,
days: int = 30,
batch_size: int = 1000,
max_candidate_batch_size: int | None = None,
delete_batch_size: int | None = None,
per_app_batch_size: int | None = None,
app_page_size: int = 500,
scan_strategy: MessagesCleanScanStrategy = "auto",
dry_run: bool = False,
task_label: str = "custom",
) -> "MessagesCleanService":
@@ -271,7 +364,12 @@ class MessagesCleanService:
Args:
policy: The policy that determines which messages to delete
days: Number of days to look back from now
batch_size: Number of messages to process per batch
batch_size: Initial number of candidate messages to scan per batch
max_candidate_batch_size: Maximum number of candidate messages to scan per batch
delete_batch_size: Maximum number of messages to delete per transaction
per_app_batch_size: Maximum number of messages to fetch per eligible app turn
app_page_size: Number of apps to inspect per eligibility discovery query
scan_strategy: "auto", "global", or "eligible_apps"
dry_run: Whether to perform a dry run (no actual deletion)
task_label: Optional task label for retention metrics
@@ -287,13 +385,32 @@ class MessagesCleanService:
if batch_size <= 0:
raise ValueError(f"batch_size ({batch_size}) must be greater than 0")
if max_candidate_batch_size is not None and max_candidate_batch_size <= 0:
raise ValueError(f"max_candidate_batch_size ({max_candidate_batch_size}) must be greater than 0")
if delete_batch_size is not None and delete_batch_size <= 0:
raise ValueError(f"delete_batch_size ({delete_batch_size}) must be greater than 0")
if per_app_batch_size is not None and per_app_batch_size <= 0:
raise ValueError(f"per_app_batch_size ({per_app_batch_size}) must be greater than 0")
if app_page_size <= 0:
raise ValueError(f"app_page_size ({app_page_size}) must be greater than 0")
end_before = naive_utc_now() - datetime.timedelta(days=days)
logger.info(
"clean_messages: days=%s, end_before=%s, batch_size=%s, policy=%s",
"clean_messages: days=%s, end_before=%s, batch_size=%s, "
"max_candidate_batch_size=%s, delete_batch_size=%s, per_app_batch_size=%s, "
"app_page_size=%s, scan_strategy=%s, policy=%s",
days,
end_before,
batch_size,
max_candidate_batch_size,
delete_batch_size,
per_app_batch_size,
app_page_size,
scan_strategy,
policy.__class__.__name__,
)
@@ -302,6 +419,11 @@ class MessagesCleanService:
end_before=end_before,
start_from=None,
batch_size=batch_size,
max_candidate_batch_size=max_candidate_batch_size,
delete_batch_size=delete_batch_size,
per_app_batch_size=per_app_batch_size,
app_page_size=app_page_size,
scan_strategy=scan_strategy,
dry_run=dry_run,
task_label=task_label,
)
@@ -316,6 +438,8 @@ class MessagesCleanService:
status = "success"
run_start = time.monotonic()
try:
if self._should_use_eligible_app_strategy():
return self._clean_messages_by_eligible_apps()
return self._clean_messages_by_time_range()
except Exception:
status = "failed"
@@ -326,6 +450,18 @@ class MessagesCleanService:
job_duration_seconds=time.monotonic() - run_start,
)
def _should_use_eligible_app_strategy(self) -> bool:
if self._scan_strategy == "global":
return False
if isinstance(self._policy, EligibleAppMessagesCleanPolicy):
return True
if self._scan_strategy == "eligible_apps":
raise ValueError("scan_strategy=eligible_apps requires an eligible-app cleanup policy")
return False
def _clean_messages_by_time_range(self) -> MessagesCleanStatsDict:
"""
Clean messages within a time range using cursor-based pagination.
@@ -333,10 +469,10 @@ class MessagesCleanService:
Time range is [start_from, end_before)
Steps:
1. Iterate messages using cursor pagination (by created_at, id)
2. Query app_id -> tenant_id mapping
1. Iterate candidate messages using cursor pagination (by created_at, id)
2. Resolve app_id -> tenant_id mapping with a job-level cache
3. Delegate to policy to determine which messages to delete
4. Batch delete messages and their relations
4. Delete messages and their relations in small chunks
Returns:
Dict with statistics: batches, filtered_messages, total_deleted
@@ -348,19 +484,23 @@ class MessagesCleanService:
"total_deleted": 0,
}
# Cursor-based pagination using (created_at, id) to avoid infinite loops
# and ensure proper ordering with time-based filtering
_cursor: tuple[datetime.datetime, str] | None = None
cursor: tuple[datetime.datetime, str] | None = None
app_to_tenant_cache: dict[str, str | None] = {}
current_candidate_batch_size = self._candidate_batch_size
smoothed_hit_rate: float | None = None
session_factory = sessionmaker(bind=db.engine, expire_on_commit=False)
logger.info(
"clean_messages: start cleaning messages (dry_run=%s), start_from=%s, end_before=%s",
"clean_messages: start cleaning messages (dry_run=%s), start_from=%s, end_before=%s, "
"candidate_batch_size=%s, max_candidate_batch_size=%s, delete_batch_size=%s",
self._dry_run,
self._start_from,
self._end_before,
self._candidate_batch_size,
self._max_candidate_batch_size,
self._delete_batch_size,
)
max_batch_interval_ms = dify_config.SANDBOX_EXPIRED_RECORDS_CLEAN_BATCH_MAX_INTERVAL
while True:
stats["batches"] += 1
batch_start = time.monotonic()
@@ -369,25 +509,25 @@ class MessagesCleanService:
batch_deleted_messages = 0
# Step 1: Fetch a batch of messages using cursor
with sessionmaker(bind=db.engine, expire_on_commit=False).begin() as session:
with session_factory.begin() as session:
fetch_messages_start = time.monotonic()
msg_stmt = (
select(Message.id, Message.app_id, Message.created_at)
.where(Message.created_at < self._end_before)
.order_by(Message.created_at, Message.id)
.limit(self._batch_size)
.limit(current_candidate_batch_size)
)
if self._start_from:
msg_stmt = msg_stmt.where(Message.created_at >= self._start_from)
# Apply cursor condition: (created_at, id) > (last_created_at, last_message_id)
if _cursor:
if cursor:
msg_stmt = msg_stmt.where(
tuple_(Message.created_at, Message.id)
> tuple_(
sa.literal(_cursor[0], type_=sa.DateTime()),
sa.literal(_cursor[1], type_=Message.id.type),
sa.literal(cursor[0], type_=sa.DateTime()),
sa.literal(cursor[1], type_=Message.id.type),
)
)
@@ -397,13 +537,13 @@ class MessagesCleanService:
for msg_id, app_id, msg_created_at in raw_messages
]
logger.info(
"clean_messages (batch %s): fetched %s messages in %sms",
"clean_messages (batch %s): fetched %s candidate messages with limit %s in %sms",
stats["batches"],
len(messages),
current_candidate_batch_size,
int((time.monotonic() - fetch_messages_start) * 1000),
)
# Track total messages fetched across all batches
stats["total_messages"] += len(messages)
batch_scanned_messages = len(messages)
@@ -417,29 +557,50 @@ class MessagesCleanService:
)
break
# Update cursor to the last message's (created_at, id)
_cursor = (messages[-1].created_at, messages[-1].id)
# Advance by candidate rows before policy filtering. This avoids retrying the same paid/unknown slice.
cursor = (messages[-1].created_at, messages[-1].id)
# Step 2: Extract app_ids and query tenant_ids
app_ids = list({msg.app_id for msg in messages})
if not app_ids:
logger.info("clean_messages (batch %s): no app_ids found, skip", stats["batches"])
smoothed_hit_rate, current_candidate_batch_size = self._adjust_candidate_batch_size(
smoothed_hit_rate=smoothed_hit_rate,
candidate_count=batch_scanned_messages,
eligible_count=0,
)
self._metrics.record_batch(
scanned_messages=batch_scanned_messages,
filtered_messages=batch_filtered_messages,
deleted_messages=batch_deleted_messages,
batch_duration_seconds=time.monotonic() - batch_start,
)
continue
fetch_apps_start = time.monotonic()
app_stmt = select(App.id, App.tenant_id).where(App.id.in_(app_ids))
apps = list(session.execute(app_stmt).all())
app_to_tenant, app_cache_misses, apps_found = self._load_app_to_tenant_mapping(
session=session,
app_ids=app_ids,
app_to_tenant_cache=app_to_tenant_cache,
)
logger.info(
"clean_messages (batch %s): fetched %s apps for %s app_ids in %sms",
"clean_messages (batch %s): resolved %s apps for %s app_ids (cache_misses=%s, found=%s) in %sms",
stats["batches"],
len(apps),
len(app_to_tenant),
len(app_ids),
app_cache_misses,
apps_found,
int((time.monotonic() - fetch_apps_start) * 1000),
)
if not apps:
if not app_to_tenant:
logger.info("clean_messages (batch %s): no apps found, skip", stats["batches"])
smoothed_hit_rate, current_candidate_batch_size = self._adjust_candidate_batch_size(
smoothed_hit_rate=smoothed_hit_rate,
candidate_count=batch_scanned_messages,
eligible_count=0,
)
self._metrics.record_batch(
scanned_messages=batch_scanned_messages,
filtered_messages=batch_filtered_messages,
@@ -448,19 +609,27 @@ class MessagesCleanService:
)
continue
# Build app_id -> tenant_id mapping
app_to_tenant: dict[str, str] = {app.id: app.tenant_id for app in apps}
# Step 3: Delegate to policy to determine which messages to delete
policy_start = time.monotonic()
message_ids_to_delete = self._policy.filter_message_ids(messages, app_to_tenant)
message_ids_to_delete = list(self._policy.filter_message_ids(messages, app_to_tenant))
batch_filtered_messages = len(message_ids_to_delete)
stats["filtered_messages"] += batch_filtered_messages
smoothed_hit_rate, next_candidate_batch_size = self._adjust_candidate_batch_size(
smoothed_hit_rate=smoothed_hit_rate,
candidate_count=batch_scanned_messages,
eligible_count=batch_filtered_messages,
)
logger.info(
"clean_messages (batch %s): policy selected %s/%s messages in %sms",
"clean_messages (batch %s): policy selected %s/%s messages in %sms "
"(smoothed_hit_rate=%.4f, next_candidate_batch_size=%s)",
stats["batches"],
len(message_ids_to_delete),
batch_filtered_messages,
len(messages),
int((time.monotonic() - policy_start) * 1000),
smoothed_hit_rate,
next_candidate_batch_size,
)
current_candidate_batch_size = next_candidate_batch_size
if not message_ids_to_delete:
logger.info("clean_messages (batch %s): no messages to delete, skip", stats["batches"])
@@ -472,47 +641,32 @@ class MessagesCleanService:
)
continue
stats["filtered_messages"] += len(message_ids_to_delete)
batch_filtered_messages = len(message_ids_to_delete)
# Step 4: Batch delete messages and their relations
if not self._dry_run:
with sessionmaker(bind=db.engine, expire_on_commit=False).begin() as session:
delete_relations_start = time.monotonic()
# Delete related records first
self._batch_delete_message_relations(session, message_ids_to_delete)
delete_relations_ms = int((time.monotonic() - delete_relations_start) * 1000)
delete_result = self._delete_messages_in_chunks(
session_factory=session_factory,
message_ids=message_ids_to_delete,
)
# Delete messages
delete_messages_start = time.monotonic()
delete_stmt = delete(Message).where(Message.id.in_(message_ids_to_delete))
delete_result = cast(CursorResult, session.execute(delete_stmt))
messages_deleted = delete_result.rowcount
delete_messages_ms = int((time.monotonic() - delete_messages_start) * 1000)
commit_ms = 0
stats["total_deleted"] += delete_result["messages_deleted"]
batch_deleted_messages = delete_result["messages_deleted"]
stats["total_deleted"] += messages_deleted
batch_deleted_messages = messages_deleted
logger.info(
"clean_messages (batch %s): processed %s candidate messages, deleted %s messages in %s chunks",
stats["batches"],
len(messages),
delete_result["messages_deleted"],
delete_result["chunks"],
)
logger.info(
"clean_messages (batch %s): relations %sms, messages %sms, batch total %sms",
stats["batches"],
delete_result["relations_ms"],
delete_result["messages_ms"],
int((time.monotonic() - batch_start) * 1000),
)
logger.info(
"clean_messages (batch %s): processed %s messages, deleted %s messages",
stats["batches"],
len(messages),
messages_deleted,
)
logger.info(
"clean_messages (batch %s): relations %sms, messages %sms, commit %sms, batch total %sms",
stats["batches"],
delete_relations_ms,
delete_messages_ms,
commit_ms,
int((time.monotonic() - batch_start) * 1000),
)
# Random sleep between batches to avoid overwhelming the database
sleep_ms = random.uniform(0, max_batch_interval_ms) # noqa: S311
logger.info("clean_messages (batch %s): sleeping for %.2fms", stats["batches"], sleep_ms)
time.sleep(sleep_ms / 1000)
self._sleep_after_batch(stats["batches"], batch_deleted_messages)
else:
# Log random sample of message IDs that would be deleted (up to 10)
sample_size = min(10, len(message_ids_to_delete))
@@ -544,6 +698,267 @@ class MessagesCleanService:
return stats
def _clean_messages_by_eligible_apps(self) -> MessagesCleanStatsDict:
"""
Clean messages by first discovering eligible apps, then scanning messages with per-app cursors.
The policy may cache plans during app discovery. Before each delete chunk, the policy must revalidate
the selected message ids against fresh tenant plans.
"""
if not isinstance(self._policy, EligibleAppMessagesCleanPolicy):
raise ValueError("eligible app cleanup requires an eligible-app cleanup policy")
stats: MessagesCleanStatsDict = {
"batches": 0,
"total_messages": 0,
"filtered_messages": 0,
"total_deleted": 0,
}
session_factory = sessionmaker(bind=db.engine, expire_on_commit=False)
scanner = EligibleAppRoundRobinScanner(
policy=self._policy,
start_from=self._start_from,
end_before=self._end_before,
app_page_size=self._app_page_size,
per_app_batch_size=self._per_app_batch_size,
)
logger.info(
"clean_messages: start eligible-app cleanup (dry_run=%s), start_from=%s, end_before=%s, "
"app_page_size=%s, per_app_batch_size=%s, delete_batch_size=%s",
self._dry_run,
self._start_from,
self._end_before,
self._app_page_size,
self._per_app_batch_size,
self._delete_batch_size,
)
while True:
stats["batches"] += 1
batch_start = time.monotonic()
fetch_start = time.monotonic()
scan_batch = scanner.fetch_batch(session_factory, target_message_count=self._delete_batch_size)
batch_scanned_messages = len(scan_batch.messages)
stats["total_messages"] += batch_scanned_messages
logger.info(
"clean_messages (batch %s, eligible_apps): fetched %s messages from %s app turns in %sms "
"(apps_scanned=%s, eligible_apps=%s, empty_apps=%s, exhausted_apps=%s)",
stats["batches"],
batch_scanned_messages,
scan_batch.app_fetches,
int((time.monotonic() - fetch_start) * 1000),
scanner.scanned_apps,
scanner.eligible_apps,
scanner.empty_apps,
scan_batch.exhausted_apps,
)
if not scan_batch.messages:
logger.info("clean_messages (batch %s, eligible_apps): no more messages to process", stats["batches"])
self._metrics.record_batch(
scanned_messages=batch_scanned_messages,
filtered_messages=0,
deleted_messages=0,
batch_duration_seconds=time.monotonic() - batch_start,
)
break
revalidate_start = time.monotonic()
message_ids_to_delete = list(
self._policy.revalidate_message_ids(scan_batch.messages, scan_batch.app_to_tenant)
)
batch_filtered_messages = len(message_ids_to_delete)
stats["filtered_messages"] += batch_filtered_messages
revalidated_dropped = batch_scanned_messages - batch_filtered_messages
logger.info(
"clean_messages (batch %s, eligible_apps): revalidated %s/%s messages in %sms (dropped=%s)",
stats["batches"],
batch_filtered_messages,
batch_scanned_messages,
int((time.monotonic() - revalidate_start) * 1000),
revalidated_dropped,
)
batch_deleted_messages = self._handle_delete_candidates(
session_factory=session_factory,
message_ids_to_delete=message_ids_to_delete,
batch_index=stats["batches"],
scanned_messages=batch_scanned_messages,
)
stats["total_deleted"] += batch_deleted_messages
self._metrics.record_batch(
scanned_messages=batch_scanned_messages,
filtered_messages=batch_filtered_messages,
deleted_messages=batch_deleted_messages,
batch_duration_seconds=time.monotonic() - batch_start,
)
logger.info(
"clean_messages completed: total batches: %s, total messages: %s, filtered messages: %s, total deleted: %s",
stats["batches"],
stats["total_messages"],
stats["filtered_messages"],
stats["total_deleted"],
)
logger.info(
"clean_messages eligible-app discovery completed: apps_scanned=%s, eligible_apps=%s, empty_apps=%s",
scanner.scanned_apps,
scanner.eligible_apps,
scanner.empty_apps,
)
return stats
def _handle_delete_candidates(
self,
*,
session_factory: sessionmaker[Session],
message_ids_to_delete: Sequence[str],
batch_index: int,
scanned_messages: int,
) -> int:
if not message_ids_to_delete:
logger.info("clean_messages (batch %s): no messages to delete, skip", batch_index)
return 0
if self._dry_run:
sample_size = min(10, len(message_ids_to_delete))
sampled_ids = random.sample(list(message_ids_to_delete), sample_size)
logger.info(
"clean_messages (batch %s, dry_run): would delete %s messages, sampling %s ids:",
batch_index,
len(message_ids_to_delete),
sample_size,
)
for msg_id in sampled_ids:
logger.info("clean_messages (batch %s, dry_run) sample: message_id=%s", batch_index, msg_id)
return 0
delete_result = self._delete_messages_in_chunks(
session_factory=session_factory,
message_ids=message_ids_to_delete,
)
logger.info(
"clean_messages (batch %s): processed %s candidate messages, deleted %s messages in %s chunks",
batch_index,
scanned_messages,
delete_result["messages_deleted"],
delete_result["chunks"],
)
logger.info(
"clean_messages (batch %s): relations %sms, messages %sms",
batch_index,
delete_result["relations_ms"],
delete_result["messages_ms"],
)
self._sleep_after_batch(batch_index, delete_result["messages_deleted"])
return delete_result["messages_deleted"]
@staticmethod
def _load_app_to_tenant_mapping(
*,
session: Session,
app_ids: Sequence[str],
app_to_tenant_cache: dict[str, str | None],
) -> tuple[dict[str, str], int, int]:
unique_app_ids = sorted(set(app_ids))
missing_app_ids = [app_id for app_id in unique_app_ids if app_id not in app_to_tenant_cache]
found_apps = 0
if missing_app_ids:
app_stmt = select(App.id, App.tenant_id).where(App.id.in_(missing_app_ids))
apps = list(session.execute(app_stmt).all())
found_app_ids: set[str] = set()
for app_id, tenant_id in apps:
app_to_tenant_cache[app_id] = tenant_id
found_app_ids.add(app_id)
found_apps = len(found_app_ids)
for app_id in set(missing_app_ids) - found_app_ids:
app_to_tenant_cache[app_id] = None
app_to_tenant: dict[str, str] = {}
for app_id in unique_app_ids:
tenant_id = app_to_tenant_cache.get(app_id)
if tenant_id is not None:
app_to_tenant[app_id] = tenant_id
return app_to_tenant, len(missing_app_ids), found_apps
def _adjust_candidate_batch_size(
self,
*,
smoothed_hit_rate: float | None,
candidate_count: int,
eligible_count: int,
) -> tuple[float, int]:
if candidate_count <= 0:
next_smoothed_hit_rate = smoothed_hit_rate or 0.0
return next_smoothed_hit_rate, self._candidate_batch_size_for_hit_rate(next_smoothed_hit_rate)
hit_rate = eligible_count / candidate_count
if smoothed_hit_rate is None:
next_smoothed_hit_rate = hit_rate
else:
next_smoothed_hit_rate = (_HIT_RATE_EMA_ALPHA * hit_rate) + ((1 - _HIT_RATE_EMA_ALPHA) * smoothed_hit_rate)
return next_smoothed_hit_rate, self._candidate_batch_size_for_hit_rate(next_smoothed_hit_rate)
def _candidate_batch_size_for_hit_rate(self, hit_rate: float) -> int:
effective_hit_rate = max(hit_rate, _MIN_ELIGIBLE_HIT_RATE)
desired_batch_size = math.ceil(self._delete_batch_size / effective_hit_rate)
lower_bound = min(self._delete_batch_size, self._max_candidate_batch_size)
return min(max(desired_batch_size, lower_bound), self._max_candidate_batch_size)
def _delete_messages_in_chunks(
self,
*,
session_factory: sessionmaker[Session],
message_ids: Sequence[str],
) -> MessageDeleteResultDict:
result: MessageDeleteResultDict = {
"messages_deleted": 0,
"chunks": 0,
"relations_ms": 0,
"messages_ms": 0,
}
for message_id_chunk in self._iter_message_id_chunks(message_ids, self._delete_batch_size):
result["chunks"] += 1
with session_factory.begin() as session:
delete_relations_start = time.monotonic()
self._batch_delete_message_relations(session, message_id_chunk)
result["relations_ms"] += int((time.monotonic() - delete_relations_start) * 1000)
delete_messages_start = time.monotonic()
delete_stmt = delete(Message).where(Message.id.in_(message_id_chunk))
delete_result = cast(CursorResult, session.execute(delete_stmt))
result["messages_deleted"] += delete_result.rowcount
result["messages_ms"] += int((time.monotonic() - delete_messages_start) * 1000)
return result
@staticmethod
def _iter_message_id_chunks(message_ids: Sequence[str], chunk_size: int) -> Iterator[Sequence[str]]:
for start_index in range(0, len(message_ids), chunk_size):
yield message_ids[start_index : start_index + chunk_size]
def _sleep_after_batch(self, batch_index: int, deleted_messages: int) -> None:
if deleted_messages <= 0:
return
max_batch_interval_ms = dify_config.SANDBOX_EXPIRED_RECORDS_CLEAN_BATCH_MAX_INTERVAL
sleep_ratio = min(1.0, deleted_messages / self._delete_batch_size)
sleep_ms = random.uniform(0, max_batch_interval_ms * sleep_ratio) # noqa: S311
logger.info("clean_messages (batch %s): sleeping for %.2fms", batch_index, sleep_ms)
time.sleep(sleep_ms / 1000)
@staticmethod
def _batch_delete_message_relations(session: Session, message_ids: Sequence[str]) -> None:
"""
@@ -1,11 +1,14 @@
import datetime
import math
import uuid
from unittest.mock import MagicMock, patch
import pytest
from sqlalchemy import delete, func, select
from sqlalchemy.orm import Session
from core.db.session_factory import session_factory
from enums.cloud_plan import CloudPlan
from models import Tenant
from models.enums import FeedbackFromSource, FeedbackRating
from models.model import (
@@ -15,7 +18,7 @@ from models.model import (
MessageAnnotation,
MessageFeedback,
)
from services.retention.conversation.messages_clean_policy import BillingDisabledPolicy
from services.retention.conversation.messages_clean_policy import BillingDisabledPolicy, BillingSandboxPolicy
from services.retention.conversation.messages_clean_service import MessagesCleanService
_NOW = datetime.datetime(2026, 1, 15, 12, 0, 0, tzinfo=datetime.UTC)
@@ -91,6 +94,35 @@ def _make_message(app_id: str, conversation_id: str, created_at: datetime.dateti
)
def _create_tenant_app_conversation(session: Session, name_suffix: str) -> tuple[str, str, str]:
tenant = Tenant(name=f"retention_it_tenant_{name_suffix}")
session.add(tenant)
session.flush()
app = App(
tenant_id=tenant.id,
name=f"Retention IT App {name_suffix}",
mode="chat",
enable_site=True,
enable_api=True,
)
session.add(app)
session.flush()
conv = Conversation(
app_id=app.id,
mode="chat",
name=f"test_conv_{name_suffix}",
status="normal",
from_source="console",
_inputs={},
)
session.add(conv)
session.flush()
return tenant.id, app.id, conv.id
class TestMessagesCleanServiceIntegration:
@pytest.fixture
def seed_messages(self, tenant_and_app):
@@ -285,6 +317,118 @@ class TestMessagesCleanServiceIntegration:
remaining = session.scalar(select(func.count()).select_from(Message).where(Message.id.in_(msg_ids)))
assert remaining == 0
def test_candidate_cursor_advances_when_first_batch_has_no_eligible_messages(self, flask_req_ctx):
"""A paid-only candidate batch must not prevent later sandbox messages from being cleaned."""
del flask_req_ctx
with session_factory.create_session() as session:
paid_tenant_id, paid_app_id, paid_conv_id = _create_tenant_app_conversation(session, "paid")
free_tenant_id, free_app_id, free_conv_id = _create_tenant_app_conversation(session, "free")
paid_msg_1 = _make_message(paid_app_id, paid_conv_id, _OLD)
paid_msg_2 = _make_message(paid_app_id, paid_conv_id, _OLD + datetime.timedelta(seconds=1))
free_msg = _make_message(free_app_id, free_conv_id, _OLD + datetime.timedelta(seconds=2))
session.add_all([paid_msg_1, paid_msg_2, free_msg])
session.commit()
paid_message_ids = [paid_msg_1.id, paid_msg_2.id]
free_message_id = free_msg.id
app_ids = [paid_app_id, free_app_id]
conversation_ids = [paid_conv_id, free_conv_id]
tenant_ids = [paid_tenant_id, free_tenant_id]
plan_map = {
paid_tenant_id: {"plan": CloudPlan.PROFESSIONAL, "expiration_date": -1},
free_tenant_id: {"plan": CloudPlan.SANDBOX, "expiration_date": -1},
}
plan_provider = MagicMock(
side_effect=lambda tenant_ids: {tenant_id: plan_map[tenant_id] for tenant_id in tenant_ids}
)
policy = BillingSandboxPolicy(plan_provider=plan_provider, graceful_period_days=21)
try:
with patch("services.retention.conversation.messages_clean_service.time.sleep"):
svc = MessagesCleanService.from_time_range(
policy=policy,
start_from=_OLD - datetime.timedelta(seconds=1),
end_before=_OLD + datetime.timedelta(seconds=3),
batch_size=2,
max_candidate_batch_size=2,
delete_batch_size=1,
scan_strategy="global",
)
stats = svc.run()
assert stats["total_messages"] == 3
assert stats["filtered_messages"] == 1
assert stats["total_deleted"] == 1
assert plan_provider.call_count == 2
assert set(plan_provider.call_args_list[0].args[0]) == {paid_tenant_id}
assert set(plan_provider.call_args_list[1].args[0]) == {free_tenant_id}
with session_factory.create_session() as session:
remaining_paid = session.scalar(
select(func.count()).select_from(Message).where(Message.id.in_(paid_message_ids))
)
remaining_free = session.scalar(
select(func.count()).select_from(Message).where(Message.id == free_message_id)
)
assert remaining_paid == 2
assert remaining_free == 0
finally:
with session_factory.create_session() as session:
session.execute(delete(Message).where(Message.id.in_([*paid_message_ids, free_message_id])))
session.execute(delete(Conversation).where(Conversation.id.in_(conversation_ids)))
session.execute(delete(App).where(App.id.in_(app_ids)))
session.execute(delete(Tenant).where(Tenant.id.in_(tenant_ids)))
session.commit()
def test_delete_batch_size_chunks_eligible_message_deletes(self, tenant_and_app):
"""Candidate scans can be larger than the delete transaction chunk size."""
data = tenant_and_app
app_id = data["app_id"]
conv_id = data["conversation_id"]
msg_ids: list[str] = []
with session_factory.create_session() as session:
for index in range(5):
msg = _make_message(app_id, conv_id, _OLD + datetime.timedelta(seconds=index))
session.add(msg)
session.flush()
msg_ids.append(msg.id)
session.commit()
try:
with (
patch.object(
MessagesCleanService,
"_batch_delete_message_relations",
wraps=MessagesCleanService._batch_delete_message_relations,
) as delete_relations,
patch("services.retention.conversation.messages_clean_service.time.sleep"),
):
svc = MessagesCleanService.from_time_range(
policy=BillingDisabledPolicy(),
start_from=_OLD - datetime.timedelta(seconds=1),
end_before=_OLD + datetime.timedelta(seconds=6),
batch_size=5,
max_candidate_batch_size=5,
delete_batch_size=2,
)
stats = svc.run()
assert stats["total_deleted"] == 5
assert delete_relations.call_count == 3
with session_factory.create_session() as session:
remaining = session.scalar(select(func.count()).select_from(Message).where(Message.id.in_(msg_ids)))
assert remaining == 0
finally:
with session_factory.create_session() as session:
session.execute(delete(Message).where(Message.id.in_(msg_ids)))
session.commit()
def test_no_messages_in_range_returns_empty_stats(self, seed_messages):
"""A window entirely in the future must yield zero matches."""
far_future = _NOW + datetime.timedelta(days=365)
@@ -32,6 +32,12 @@ def test_absolute_mode_calls_from_time_range():
):
clean_expired_messages.callback(
batch_size=200,
candidate_batch_size=None,
max_candidate_batch_size=None,
delete_batch_size=None,
per_app_batch_size=None,
app_page_size=500,
scan_strategy="auto",
graceful_period=21,
start_from=start_from,
end_before=end_before,
@@ -45,6 +51,11 @@ def test_absolute_mode_calls_from_time_range():
start_from=start_from,
end_before=end_before,
batch_size=200,
max_candidate_batch_size=None,
delete_batch_size=200,
per_app_batch_size=200,
app_page_size=500,
scan_strategy="auto",
dry_run=True,
task_label="custom",
)
@@ -62,6 +73,12 @@ def test_relative_mode_before_days_only_calls_from_days():
):
clean_expired_messages.callback(
batch_size=500,
candidate_batch_size=None,
max_candidate_batch_size=None,
delete_batch_size=None,
per_app_batch_size=None,
app_page_size=500,
scan_strategy="auto",
graceful_period=14,
start_from=None,
end_before=None,
@@ -74,6 +91,11 @@ def test_relative_mode_before_days_only_calls_from_days():
policy=policy,
days=30,
batch_size=500,
max_candidate_batch_size=None,
delete_batch_size=500,
per_app_batch_size=500,
app_page_size=500,
scan_strategy="auto",
dry_run=False,
task_label="before-30",
)
@@ -93,6 +115,12 @@ def test_relative_mode_with_from_days_ago_calls_from_time_range():
):
clean_expired_messages.callback(
batch_size=1000,
candidate_batch_size=None,
max_candidate_batch_size=None,
delete_batch_size=None,
per_app_batch_size=None,
app_page_size=500,
scan_strategy="auto",
graceful_period=21,
start_from=None,
end_before=None,
@@ -106,12 +134,60 @@ def test_relative_mode_with_from_days_ago_calls_from_time_range():
start_from=fixed_now - datetime.timedelta(days=60),
end_before=fixed_now - datetime.timedelta(days=30),
batch_size=1000,
max_candidate_batch_size=None,
delete_batch_size=1000,
per_app_batch_size=1000,
app_page_size=500,
scan_strategy="auto",
dry_run=False,
task_label="60to30",
)
mock_from_days.assert_not_called()
def test_explicit_batch_knobs_are_passed_to_service():
policy = object()
service = _mock_service()
start_from = datetime.datetime(2024, 1, 1, 0, 0, 0)
end_before = datetime.datetime(2024, 2, 1, 0, 0, 0)
with (
patch("commands.retention.create_message_clean_policy", return_value=policy),
patch("commands.retention.MessagesCleanService.from_time_range", return_value=service) as mock_from_time_range,
patch("commands.retention.MessagesCleanService.from_days") as mock_from_days,
):
clean_expired_messages.callback(
batch_size=1000,
candidate_batch_size=10000,
max_candidate_batch_size=50000,
delete_batch_size=500,
per_app_batch_size=250,
app_page_size=50,
scan_strategy="global",
graceful_period=21,
start_from=start_from,
end_before=end_before,
from_days_ago=None,
before_days=None,
dry_run=True,
)
mock_from_time_range.assert_called_once_with(
policy=policy,
start_from=start_from,
end_before=end_before,
batch_size=10000,
max_candidate_batch_size=50000,
delete_batch_size=500,
per_app_batch_size=250,
app_page_size=50,
scan_strategy="global",
dry_run=True,
task_label="custom",
)
mock_from_days.assert_not_called()
@pytest.mark.parametrize(
("kwargs", "message"),
[
@@ -175,6 +251,12 @@ def test_invalid_inputs_raise_usage_error(kwargs: dict, message: str):
with pytest.raises(click.UsageError, match=re.escape(message)):
clean_expired_messages.callback(
batch_size=1000,
candidate_batch_size=None,
max_candidate_batch_size=None,
delete_batch_size=None,
per_app_batch_size=None,
app_page_size=500,
scan_strategy="auto",
graceful_period=21,
start_from=kwargs["start_from"],
end_before=kwargs["end_before"],
@@ -5,6 +5,7 @@ from unittest.mock import MagicMock, patch
import pytest
from enums.cloud_plan import CloudPlan
from services.retention.conversation.messages_clean_app_scan import EligibleAppRoundRobinScanner, EligibleAppScanBatch
from services.retention.conversation.messages_clean_policy import (
BillingDisabledPolicy,
BillingSandboxPolicy,
@@ -296,6 +297,130 @@ class TestBillingSandboxPolicyFilterMessageIds:
called_tenant_ids = set(plan_provider.call_args[0][0])
assert called_tenant_ids == {"tenant1", "tenant2"}
def test_plan_provider_reuses_job_level_cache(self):
"""Test that repeated tenants are not fetched from billing more than once."""
# Arrange
now = self.CURRENT_TIMESTAMP
tenant_plans = {
"tenant1": {"plan": CloudPlan.SANDBOX, "expiration_date": -1},
"tenant2": {"plan": CloudPlan.SANDBOX, "expiration_date": -1},
}
plan_provider = MagicMock(
side_effect=lambda tenant_ids: {tenant_id: tenant_plans[tenant_id] for tenant_id in tenant_ids}
)
policy = BillingSandboxPolicy(
plan_provider=plan_provider,
graceful_period_days=self.GRACEFUL_PERIOD_DAYS,
current_timestamp=now,
)
# Act
first_result = policy.filter_message_ids(
[make_simple_message("msg1", "app1")],
{"app1": "tenant1"},
)
second_result = policy.filter_message_ids(
[
make_simple_message("msg2", "app1"),
make_simple_message("msg3", "app2"),
],
{"app1": "tenant1", "app2": "tenant2"},
)
# Assert
assert set(first_result) == {"msg1"}
assert set(second_result) == {"msg2", "msg3"}
assert plan_provider.call_count == 2
assert set(plan_provider.call_args_list[0].args[0]) == {"tenant1"}
assert set(plan_provider.call_args_list[1].args[0]) == {"tenant2"}
def test_whitelisted_tenants_are_not_fetched_from_plan_provider(self):
"""Test that whitelisted tenants are skipped before billing plan lookup."""
# Arrange
plan_provider = make_plan_provider({})
policy = BillingSandboxPolicy(
plan_provider=plan_provider,
graceful_period_days=self.GRACEFUL_PERIOD_DAYS,
tenant_whitelist=["tenant1"],
current_timestamp=self.CURRENT_TIMESTAMP,
)
# Act
result = policy.filter_message_ids(
[make_simple_message("msg1", "app1")],
{"app1": "tenant1"},
)
# Assert
assert list(result) == []
plan_provider.assert_not_called()
def test_filter_app_to_tenant_returns_only_eligible_apps(self):
"""Test eligible-app discovery filters out paid, grace-period, and whitelisted tenants."""
now = self.CURRENT_TIMESTAMP
expired_old = now - (15 * 24 * 60 * 60)
expired_recent = now - (3 * 24 * 60 * 60)
plan_provider = make_plan_provider(
{
"tenant1": {"plan": CloudPlan.SANDBOX, "expiration_date": -1},
"tenant2": {"plan": CloudPlan.SANDBOX, "expiration_date": expired_old},
"tenant3": {"plan": CloudPlan.SANDBOX, "expiration_date": expired_recent},
"tenant4": {"plan": CloudPlan.PROFESSIONAL, "expiration_date": -1},
"tenant5": {"plan": CloudPlan.SANDBOX, "expiration_date": -1},
}
)
policy = BillingSandboxPolicy(
plan_provider=plan_provider,
graceful_period_days=self.GRACEFUL_PERIOD_DAYS,
tenant_whitelist=["tenant5"],
current_timestamp=now,
)
result = policy.filter_app_to_tenant(
{
"app1": "tenant1",
"app2": "tenant2",
"app3": "tenant3",
"app4": "tenant4",
"app5": "tenant5",
}
)
assert result == {"app1": "tenant1", "app2": "tenant2"}
def test_revalidate_message_ids_uses_fresh_provider(self):
"""Test delete-time revalidation bypasses the job-level cached provider."""
cached_plan_provider = make_plan_provider(
{
"tenant1": {"plan": CloudPlan.SANDBOX, "expiration_date": -1},
"tenant2": {"plan": CloudPlan.SANDBOX, "expiration_date": -1},
}
)
fresh_plan_provider = make_plan_provider(
{
"tenant1": {"plan": CloudPlan.SANDBOX, "expiration_date": -1},
"tenant2": {"plan": CloudPlan.PROFESSIONAL, "expiration_date": -1},
}
)
policy = BillingSandboxPolicy(
plan_provider=cached_plan_provider,
fresh_plan_provider=fresh_plan_provider,
graceful_period_days=self.GRACEFUL_PERIOD_DAYS,
current_timestamp=self.CURRENT_TIMESTAMP,
)
messages = [
make_simple_message("msg1", "app1"),
make_simple_message("msg2", "app2"),
]
app_to_tenant = {"app1": "tenant1", "app2": "tenant2"}
discovery_result = policy.filter_message_ids(messages, app_to_tenant)
revalidated_result = policy.revalidate_message_ids(messages, app_to_tenant)
assert set(discovery_result) == {"msg1", "msg2"}
assert set(revalidated_result) == {"msg1"}
fresh_plan_provider.assert_called_once()
def test_complex_mixed_scenario(self):
"""Test complex scenario with mixed plans, expirations, whitelist, and missing mappings."""
# Arrange
@@ -435,6 +560,7 @@ class TestCreateMessageCleanPolicy:
assert policy._graceful_period_days == 14
assert list(policy._tenant_whitelist) == whitelist
assert policy._plan_provider == mock_plan_provider
assert policy._fresh_plan_provider == mock_billing_service.get_plan_bulk
assert policy._current_timestamp == 1234567
@@ -497,6 +623,38 @@ class TestMessagesCleanServiceFromTimeRange:
batch_size=-100,
)
with pytest.raises(ValueError, match="max_candidate_batch_size .* must be greater than 0"):
MessagesCleanService.from_time_range(
policy=policy,
start_from=start_from,
end_before=end_before,
max_candidate_batch_size=0,
)
with pytest.raises(ValueError, match="delete_batch_size .* must be greater than 0"):
MessagesCleanService.from_time_range(
policy=policy,
start_from=start_from,
end_before=end_before,
delete_batch_size=0,
)
with pytest.raises(ValueError, match="per_app_batch_size .* must be greater than 0"):
MessagesCleanService.from_time_range(
policy=policy,
start_from=start_from,
end_before=end_before,
per_app_batch_size=0,
)
with pytest.raises(ValueError, match="app_page_size .* must be greater than 0"):
MessagesCleanService.from_time_range(
policy=policy,
start_from=start_from,
end_before=end_before,
app_page_size=0,
)
def test_valid_params_creates_instance(self):
"""Test that valid parameters create a correctly configured instance."""
# Arrange
@@ -512,6 +670,11 @@ class TestMessagesCleanServiceFromTimeRange:
start_from=start_from,
end_before=end_before,
batch_size=batch_size,
max_candidate_batch_size=5000,
delete_batch_size=200,
per_app_batch_size=100,
app_page_size=50,
scan_strategy="global",
dry_run=dry_run,
)
@@ -521,6 +684,12 @@ class TestMessagesCleanServiceFromTimeRange:
assert service._start_from == start_from
assert service._end_before == end_before
assert service._batch_size == batch_size
assert service._candidate_batch_size == batch_size
assert service._max_candidate_batch_size == 5000
assert service._delete_batch_size == 200
assert service._per_app_batch_size == 100
assert service._app_page_size == 50
assert service._scan_strategy == "global"
assert service._dry_run == dry_run
def test_default_params(self):
@@ -539,6 +708,12 @@ class TestMessagesCleanServiceFromTimeRange:
# Assert
assert service._batch_size == 1000 # default
assert service._candidate_batch_size == 1000 # default
assert service._max_candidate_batch_size == 1000 # default
assert service._delete_batch_size == 1000 # default
assert service._per_app_batch_size == 1000 # default
assert service._app_page_size == 500 # default
assert service._scan_strategy == "auto" # default
assert service._dry_run is False # default
def test_explicit_task_label(self):
@@ -590,6 +765,18 @@ class TestMessagesCleanServiceFromDays:
with pytest.raises(ValueError, match="batch_size .* must be greater than 0"):
MessagesCleanService.from_days(policy=policy, days=30, batch_size=-500)
with pytest.raises(ValueError, match="max_candidate_batch_size .* must be greater than 0"):
MessagesCleanService.from_days(policy=policy, days=30, max_candidate_batch_size=0)
with pytest.raises(ValueError, match="delete_batch_size .* must be greater than 0"):
MessagesCleanService.from_days(policy=policy, days=30, delete_batch_size=0)
with pytest.raises(ValueError, match="per_app_batch_size .* must be greater than 0"):
MessagesCleanService.from_days(policy=policy, days=30, per_app_batch_size=0)
with pytest.raises(ValueError, match="app_page_size .* must be greater than 0"):
MessagesCleanService.from_days(policy=policy, days=30, app_page_size=0)
def test_valid_params_creates_instance(self):
"""Test that valid parameters create a correctly configured instance."""
# Arrange
@@ -606,6 +793,11 @@ class TestMessagesCleanServiceFromDays:
policy=policy,
days=days,
batch_size=batch_size,
max_candidate_batch_size=2500,
delete_batch_size=250,
per_app_batch_size=125,
app_page_size=25,
scan_strategy="global",
dry_run=dry_run,
)
@@ -616,6 +808,12 @@ class TestMessagesCleanServiceFromDays:
assert service._start_from is None
assert service._end_before == expected_end_before
assert service._batch_size == batch_size
assert service._candidate_batch_size == batch_size
assert service._max_candidate_batch_size == 2500
assert service._delete_batch_size == 250
assert service._per_app_batch_size == 125
assert service._app_page_size == 25
assert service._scan_strategy == "global"
assert service._dry_run == dry_run
def test_default_params(self):
@@ -637,6 +835,152 @@ class TestMessagesCleanServiceFromDays:
assert service._metrics._base_attributes["task_label"] == "custom"
class TestMessagesCleanServiceBatchHelpers:
"""Unit tests for cache and adaptive batch helpers."""
def test_load_app_to_tenant_mapping_reuses_cache(self):
class ExecuteResult:
def __init__(self, rows: list[tuple[str, str]]) -> None:
self._rows = rows
def all(self) -> list[tuple[str, str]]:
return self._rows
class FakeSession:
execute_calls: int
def __init__(self) -> None:
self.execute_calls = 0
def execute(self, _stmt: object) -> ExecuteResult:
self.execute_calls += 1
return ExecuteResult([("app1", "tenant1")])
session = FakeSession()
cache: dict[str, str | None] = {}
first_mapping, first_cache_misses, first_found = MessagesCleanService._load_app_to_tenant_mapping(
session=session, # type: ignore[arg-type]
app_ids=["app1"],
app_to_tenant_cache=cache,
)
second_mapping, second_cache_misses, second_found = MessagesCleanService._load_app_to_tenant_mapping(
session=session, # type: ignore[arg-type]
app_ids=["app1"],
app_to_tenant_cache=cache,
)
assert first_mapping == {"app1": "tenant1"}
assert second_mapping == {"app1": "tenant1"}
assert first_cache_misses == 1
assert first_found == 1
assert second_cache_misses == 0
assert second_found == 0
assert session.execute_calls == 1
def test_candidate_batch_size_grows_for_low_hit_rate(self):
service = MessagesCleanService(
policy=BillingDisabledPolicy(),
start_from=datetime.datetime(2024, 1, 1),
end_before=datetime.datetime(2024, 1, 2),
batch_size=1000,
max_candidate_batch_size=50000,
delete_batch_size=1000,
)
smoothed_hit_rate, next_batch_size = service._adjust_candidate_batch_size(
smoothed_hit_rate=None,
candidate_count=10000,
eligible_count=55,
)
assert smoothed_hit_rate == 0.0055
assert next_batch_size == 50000
def test_candidate_batch_size_shrinks_when_hit_rate_is_high(self):
service = MessagesCleanService(
policy=BillingDisabledPolicy(),
start_from=datetime.datetime(2024, 1, 1),
end_before=datetime.datetime(2024, 1, 2),
batch_size=10000,
max_candidate_batch_size=50000,
delete_batch_size=1000,
)
_smoothed_hit_rate, next_batch_size = service._adjust_candidate_batch_size(
smoothed_hit_rate=None,
candidate_count=10000,
eligible_count=10000,
)
assert next_batch_size == 1000
def test_iter_message_id_chunks_uses_delete_batch_size(self):
chunks = list(MessagesCleanService._iter_message_id_chunks(["msg1", "msg2", "msg3", "msg4", "msg5"], 2))
assert chunks == [["msg1", "msg2"], ["msg3", "msg4"], ["msg5"]]
def test_eligible_app_scanner_filters_apps_and_round_robins(self):
class ExecuteResult:
def __init__(self, rows: list[tuple[str, str] | tuple[str, str, datetime.datetime]]) -> None:
self._rows = rows
def all(self) -> list[tuple[str, str] | tuple[str, str, datetime.datetime]]:
return self._rows
class FakeSession:
def __init__(self) -> None:
self.rows_by_call: list[list[tuple[str, str] | tuple[str, str, datetime.datetime]]] = [
[
("app1", "tenant1"),
("app2", "tenant2"),
("app3", "tenant3"),
],
[("msg1", "app1", datetime.datetime(2024, 1, 1, 0, 0, 0))],
[("msg2", "app2", datetime.datetime(2024, 1, 1, 0, 0, 1))],
]
def execute(self, _stmt: object) -> ExecuteResult:
return ExecuteResult(self.rows_by_call.pop(0))
def __enter__(self) -> "FakeSession":
return self
def __exit__(self, exc_type: object, exc_value: object, traceback: object) -> None:
return None
class FakeSessionFactory:
def __init__(self) -> None:
self.session = FakeSession()
def __call__(self) -> FakeSession:
return self.session
plan_provider = make_plan_provider(
{
"tenant1": {"plan": CloudPlan.SANDBOX, "expiration_date": -1},
"tenant2": {"plan": CloudPlan.SANDBOX, "expiration_date": -1},
"tenant3": {"plan": CloudPlan.PROFESSIONAL, "expiration_date": -1},
}
)
policy = BillingSandboxPolicy(plan_provider=plan_provider, current_timestamp=1000)
scanner = EligibleAppRoundRobinScanner(
policy=policy,
start_from=datetime.datetime(2023, 12, 1),
end_before=datetime.datetime(2024, 2, 1),
app_page_size=10,
per_app_batch_size=1,
)
batch = scanner.fetch_batch(FakeSessionFactory(), target_message_count=2) # type: ignore[arg-type]
assert [message.id for message in batch.messages] == ["msg1", "msg2"]
assert batch.app_to_tenant == {"app1": "tenant1", "app2": "tenant2"}
assert batch.app_fetches == 2
assert scanner.scanned_apps == 3
assert scanner.eligible_apps == 2
class TestMessagesCleanServiceRun:
"""Unit tests for MessagesCleanService.run instrumentation behavior."""
@@ -667,6 +1011,126 @@ class TestMessagesCleanServiceRun:
assert len(completion_calls) == 1
assert completion_calls[0]["status"] == "success"
def test_run_uses_eligible_app_strategy_for_sandbox_policy(self):
service = MessagesCleanService(
policy=BillingSandboxPolicy(plan_provider=make_plan_provider({})),
start_from=datetime.datetime(2024, 1, 1),
end_before=datetime.datetime(2024, 1, 2),
)
expected_stats = {
"batches": 1,
"total_messages": 10,
"filtered_messages": 5,
"total_deleted": 5,
}
service._clean_messages_by_eligible_apps = MagicMock(return_value=expected_stats) # type: ignore[method-assign]
service._clean_messages_by_time_range = MagicMock() # type: ignore[method-assign]
result = service.run()
assert result == expected_stats
service._clean_messages_by_eligible_apps.assert_called_once()
service._clean_messages_by_time_range.assert_not_called()
def test_global_strategy_forces_global_scan(self):
service = MessagesCleanService(
policy=BillingSandboxPolicy(plan_provider=make_plan_provider({})),
start_from=datetime.datetime(2024, 1, 1),
end_before=datetime.datetime(2024, 1, 2),
scan_strategy="global",
)
expected_stats = {
"batches": 1,
"total_messages": 10,
"filtered_messages": 5,
"total_deleted": 5,
}
service._clean_messages_by_time_range = MagicMock(return_value=expected_stats) # type: ignore[method-assign]
service._clean_messages_by_eligible_apps = MagicMock() # type: ignore[method-assign]
result = service.run()
assert result == expected_stats
service._clean_messages_by_time_range.assert_called_once()
service._clean_messages_by_eligible_apps.assert_not_called()
def test_eligible_app_strategy_requires_supported_policy(self):
service = MessagesCleanService(
policy=BillingDisabledPolicy(),
start_from=datetime.datetime(2024, 1, 1),
end_before=datetime.datetime(2024, 1, 2),
scan_strategy="eligible_apps",
)
with pytest.raises(ValueError, match="eligible-app cleanup policy"):
service.run()
def test_eligible_app_strategy_revalidates_before_dry_run_counting(self):
class FakeSessionFactory:
pass
class FakeScanner:
scanned_apps = 2
eligible_apps = 2
empty_apps = 0
def __init__(self, **_kwargs: object) -> None:
self.call_count = 0
def fetch_batch(self, _session_factory: object, *, target_message_count: int) -> EligibleAppScanBatch:
del target_message_count
self.call_count += 1
if self.call_count == 1:
return EligibleAppScanBatch(
messages=[
make_simple_message("msg1", "app1"),
make_simple_message("msg2", "app2"),
],
app_to_tenant={"app1": "tenant1", "app2": "tenant2"},
app_fetches=2,
exhausted_apps=2,
)
return EligibleAppScanBatch(messages=[], app_to_tenant={}, app_fetches=0, exhausted_apps=0)
policy = BillingSandboxPolicy(
plan_provider=make_plan_provider(
{
"tenant1": {"plan": CloudPlan.SANDBOX, "expiration_date": -1},
"tenant2": {"plan": CloudPlan.SANDBOX, "expiration_date": -1},
}
),
fresh_plan_provider=make_plan_provider(
{
"tenant1": {"plan": CloudPlan.SANDBOX, "expiration_date": -1},
"tenant2": {"plan": CloudPlan.PROFESSIONAL, "expiration_date": -1},
}
),
current_timestamp=1000,
)
service = MessagesCleanService(
policy=policy,
start_from=datetime.datetime(2024, 1, 1),
end_before=datetime.datetime(2024, 1, 2),
dry_run=True,
)
fake_db = MagicMock()
fake_db.engine = object()
with (
patch(
"services.retention.conversation.messages_clean_service.sessionmaker",
return_value=FakeSessionFactory(),
),
patch("services.retention.conversation.messages_clean_service.EligibleAppRoundRobinScanner", FakeScanner),
patch("services.retention.conversation.messages_clean_service.db", fake_db),
):
stats = service.run()
assert stats["batches"] == 2
assert stats["total_messages"] == 2
assert stats["filtered_messages"] == 1
assert stats["total_deleted"] == 0
def test_run_records_completion_metrics_on_failure(self):
# Arrange
service = MessagesCleanService(