import datetime import logging import random import time from collections.abc import Sequence from typing import TYPE_CHECKING, TypedDict, cast import sqlalchemy as sa from sqlalchemy import delete, select, tuple_ from sqlalchemy.engine import CursorResult from sqlalchemy.orm import Session, sessionmaker from configs import dify_config from extensions.ext_database import db from libs.datetime_utils import naive_utc_now from models.model import ( App, AppAnnotationHitHistory, DatasetRetrieverResource, Message, MessageAgentThought, MessageAnnotation, MessageChain, MessageFeedback, MessageFile, ) from models.web import SavedMessage from services.retention.conversation.messages_clean_policy import ( MessagesCleanPolicy, SimpleMessage, ) logger = logging.getLogger(__name__) if TYPE_CHECKING: from opentelemetry.metrics import Counter, Histogram class MessagesCleanupMetrics: """ Records low-cardinality OpenTelemetry metrics for expired message cleanup jobs. We keep labels stable (dry_run/window_mode/task_label/status) so these metrics remain dashboard-friendly for long-running CronJob executions. """ _job_runs_total: "Counter | None" _batches_total: "Counter | None" _messages_scanned_total: "Counter | None" _messages_filtered_total: "Counter | None" _messages_deleted_total: "Counter | None" _job_duration_seconds: "Histogram | None" _batch_duration_seconds: "Histogram | None" _base_attributes: dict[str, str] def __init__(self, *, dry_run: bool, has_window: bool, task_label: str) -> None: self._job_runs_total = None self._batches_total = None self._messages_scanned_total = None self._messages_filtered_total = None self._messages_deleted_total = None self._job_duration_seconds = None self._batch_duration_seconds = None self._base_attributes = { "job_name": "messages_cleanup", "dry_run": str(dry_run).lower(), "window_mode": "between" if has_window else "before_cutoff", "task_label": task_label, } self._init_instruments() def _init_instruments(self) -> None: if not dify_config.ENABLE_OTEL: return try: from opentelemetry.metrics import get_meter meter = get_meter("messages_cleanup", version=dify_config.project.version) self._job_runs_total = meter.create_counter( "messages_cleanup_jobs_total", description="Total number of expired message cleanup jobs by status.", unit="{job}", ) self._batches_total = meter.create_counter( "messages_cleanup_batches_total", description="Total number of message cleanup batches processed.", unit="{batch}", ) self._messages_scanned_total = meter.create_counter( "messages_cleanup_scanned_messages_total", description="Total messages scanned by cleanup jobs.", unit="{message}", ) self._messages_filtered_total = meter.create_counter( "messages_cleanup_filtered_messages_total", description="Total messages selected by cleanup policy.", unit="{message}", ) self._messages_deleted_total = meter.create_counter( "messages_cleanup_deleted_messages_total", description="Total messages deleted by cleanup jobs.", unit="{message}", ) self._job_duration_seconds = meter.create_histogram( "messages_cleanup_job_duration_seconds", description="Duration of expired message cleanup jobs in seconds.", unit="s", ) self._batch_duration_seconds = meter.create_histogram( "messages_cleanup_batch_duration_seconds", description="Duration of expired message cleanup batch processing in seconds.", unit="s", ) except Exception: logger.exception("messages_cleanup_metrics: failed to initialize instruments") def _attrs(self, **extra: str) -> dict[str, str]: return {**self._base_attributes, **extra} @staticmethod def _add(counter: "Counter | None", value: int, attributes: dict[str, str]) -> None: if not counter or value <= 0: return try: counter.add(value, attributes) except Exception: logger.exception("messages_cleanup_metrics: failed to add counter value") @staticmethod def _record(histogram: "Histogram | None", value: float, attributes: dict[str, str]) -> None: if not histogram: return try: histogram.record(value, attributes) except Exception: logger.exception("messages_cleanup_metrics: failed to record histogram value") def record_batch( self, *, scanned_messages: int, filtered_messages: int, deleted_messages: int, batch_duration_seconds: float, ) -> None: attributes = self._attrs() self._add(self._batches_total, 1, attributes) self._add(self._messages_scanned_total, scanned_messages, attributes) self._add(self._messages_filtered_total, filtered_messages, attributes) self._add(self._messages_deleted_total, deleted_messages, attributes) self._record(self._batch_duration_seconds, batch_duration_seconds, attributes) def record_completion(self, *, status: str, job_duration_seconds: float) -> None: attributes = self._attrs(status=status) self._add(self._job_runs_total, 1, attributes) self._record(self._job_duration_seconds, job_duration_seconds, attributes) class MessagesCleanStatsDict(TypedDict): batches: int total_messages: int filtered_messages: int total_deleted: 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). """ def __init__( self, policy: MessagesCleanPolicy, end_before: datetime.datetime, start_from: datetime.datetime | None = None, batch_size: int = 1000, dry_run: bool = False, task_label: str = "custom", ) -> None: """ Initialize the service with cleanup parameters. Args: 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 dry_run: Whether to perform a dry run (no actual deletion) task_label: Optional task label for retention metrics """ self._policy = policy self._end_before = end_before self._start_from = start_from self._batch_size = batch_size self._dry_run = dry_run self._metrics = MessagesCleanupMetrics( dry_run=dry_run, has_window=bool(start_from), task_label=task_label, ) @classmethod def from_time_range( cls, policy: MessagesCleanPolicy, start_from: datetime.datetime, end_before: datetime.datetime, batch_size: int = 1000, dry_run: bool = False, task_label: str = "custom", ) -> "MessagesCleanService": """ Create a service instance for cleaning messages within a specific time range. Time range is [start_from, end_before). Args: 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 dry_run: Whether to perform a dry run (no actual deletion) task_label: Optional task label for retention metrics Returns: MessagesCleanService instance Raises: ValueError: If start_from >= end_before or invalid parameters """ if start_from >= end_before: raise ValueError(f"start_from ({start_from}) must be less than end_before ({end_before})") if batch_size <= 0: raise ValueError(f"batch_size ({batch_size}) must be greater than 0") logger.info( "clean_messages: start_from=%s, end_before=%s, batch_size=%s, policy=%s", start_from, end_before, batch_size, policy.__class__.__name__, ) return cls( policy=policy, end_before=end_before, start_from=start_from, batch_size=batch_size, dry_run=dry_run, task_label=task_label, ) @classmethod def from_days( cls, policy: MessagesCleanPolicy, days: int = 30, batch_size: int = 1000, dry_run: bool = False, task_label: str = "custom", ) -> "MessagesCleanService": """ Create a service instance for cleaning messages older than specified days. 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 dry_run: Whether to perform a dry run (no actual deletion) task_label: Optional task label for retention metrics Returns: MessagesCleanService instance Raises: ValueError: If invalid parameters """ if days < 0: raise ValueError(f"days ({days}) must be greater than or equal to 0") if batch_size <= 0: raise ValueError(f"batch_size ({batch_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", days, end_before, batch_size, policy.__class__.__name__, ) return cls( policy=policy, end_before=end_before, start_from=None, batch_size=batch_size, dry_run=dry_run, task_label=task_label, ) def run(self) -> MessagesCleanStatsDict: """ Execute the message cleanup operation. Returns: Dict with statistics: batches, filtered_messages, total_deleted """ status = "success" run_start = time.monotonic() try: return self._clean_messages_by_time_range() except Exception: status = "failed" raise finally: self._metrics.record_completion( status=status, job_duration_seconds=time.monotonic() - run_start, ) def _clean_messages_by_time_range(self) -> MessagesCleanStatsDict: """ Clean messages within a time range using cursor-based pagination. 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 3. Delegate to policy to determine which messages to delete 4. Batch delete messages and their relations Returns: Dict with statistics: batches, filtered_messages, total_deleted """ stats: MessagesCleanStatsDict = { "batches": 0, "total_messages": 0, "filtered_messages": 0, "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 logger.info( "clean_messages: start cleaning messages (dry_run=%s), start_from=%s, end_before=%s", self._dry_run, self._start_from, self._end_before, ) max_batch_interval_ms = dify_config.SANDBOX_EXPIRED_RECORDS_CLEAN_BATCH_MAX_INTERVAL while True: stats["batches"] += 1 batch_start = time.monotonic() batch_scanned_messages = 0 batch_filtered_messages = 0 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: 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) ) 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: 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), ) ) raw_messages = list(session.execute(msg_stmt).all()) messages = [ SimpleMessage(id=msg_id, app_id=app_id, created_at=msg_created_at) for msg_id, app_id, msg_created_at in raw_messages ] logger.info( "clean_messages (batch %s): fetched %s messages in %sms", stats["batches"], len(messages), int((time.monotonic() - fetch_messages_start) * 1000), ) # Track total messages fetched across all batches stats["total_messages"] += len(messages) batch_scanned_messages = len(messages) if not messages: logger.info("clean_messages (batch %s): no more messages to process", stats["batches"]) 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, ) break # Update cursor to the last message's (created_at, id) _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"]) 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()) logger.info( "clean_messages (batch %s): fetched %s apps for %s app_ids in %sms", stats["batches"], len(apps), len(app_ids), int((time.monotonic() - fetch_apps_start) * 1000), ) if not apps: logger.info("clean_messages (batch %s): no apps found, skip", stats["batches"]) 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 # 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) logger.info( "clean_messages (batch %s): policy selected %s/%s messages in %sms", stats["batches"], len(message_ids_to_delete), len(messages), int((time.monotonic() - policy_start) * 1000), ) if not message_ids_to_delete: logger.info("clean_messages (batch %s): no messages to delete, skip", stats["batches"]) 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 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 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"] += messages_deleted batch_deleted_messages = messages_deleted 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) else: # Log random sample of message IDs that would be deleted (up to 10) 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:", stats["batches"], 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", stats["batches"], msg_id) 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"], ) return stats @staticmethod def _batch_delete_message_relations(session: Session, message_ids: Sequence[str]) -> None: """ Batch delete all related records for given message IDs. Args: session: Database session message_ids: List of message IDs to delete relations for """ if not message_ids: return # Delete all related records in batch session.execute(delete(MessageFeedback).where(MessageFeedback.message_id.in_(message_ids))) session.execute(delete(MessageAnnotation).where(MessageAnnotation.message_id.in_(message_ids))) session.execute(delete(MessageChain).where(MessageChain.message_id.in_(message_ids))) session.execute(delete(MessageAgentThought).where(MessageAgentThought.message_id.in_(message_ids))) session.execute(delete(MessageFile).where(MessageFile.message_id.in_(message_ids))) session.execute(delete(SavedMessage).where(SavedMessage.message_id.in_(message_ids))) session.execute(delete(AppAnnotationHitHistory).where(AppAnnotationHitHistory.message_id.in_(message_ids))) session.execute(delete(DatasetRetrieverResource).where(DatasetRetrieverResource.message_id.in_(message_ids)))