MoviePilot/app/db/maintenance.py

"""
数据维护用例的 SQLAlchemy 适配器。
"""

from typing import Any, Callable, ContextManager

from sqlalchemy import delete, exists, select

from app.db.base import execute_dml
from app.db.models.agentchat import AgentChat
from app.db.models.agentinvocation import AgentInvocation
from app.db.models.agenttask import AgentTask
from app.db.models.agenttaskrun import AgentTaskRun
from app.db.models.downloadfailure import DownloadFailure
from app.db.models.downloadhistory import DownloadFiles, DownloadHistory
from app.db.models.message import Message
from app.db.models.outbox import OutboxMessage
from app.db.models.siteuserdata import SiteUserData
from app.db.models.subscribehistory import SubscribeHistory
from app.db.models.transferhistory import TransferHistory
from app.db.uow import SqlAlchemyUnitOfWork


class DatabaseCleanupRepository:
    """
    把应用层清理端口映射到现有模型批量删除方法。
    """

    def __init__(self, *, session_factory: Callable[[], ContextManager[Any]]) -> None:
        """
        保存会话工厂，使测试和不同数据库后端可以显式注入。
        """
        self._session_factory = session_factory

    def session(self) -> ContextManager[Any]:
        """
        创建一次维护运行共用的数据库会话。
        """
        return self._session_factory()

    @staticmethod
    def unit_of_work(db: Any) -> SqlAlchemyUnitOfWork:
        """
        把当前维护 Session 适配成显式批次事务边界。
        """
        return SqlAlchemyUnitOfWork(db)

    @staticmethod
    def delete_messages(db: Any, cutoff: str, limit: int) -> int:
        """
        删除早于截止时间的消息。
        """
        return Message.delete_before(db=db, before_time=cutoff, limit=limit)

    @staticmethod
    def delete_download_history(db: Any, cutoff: str, limit: int) -> int:
        """
        删除早于截止时间的下载历史。
        """
        return DownloadHistory.delete_before(
            db=db,
            before_time=cutoff,
            limit=limit,
        )

    @staticmethod
    def delete_download_orphans(db: Any, limit: int) -> int:
        """
        删除已经失去父下载历史的文件记录。
        """
        return DownloadFiles.delete_orphans(db=db, limit=limit)

    @staticmethod
    def delete_site_userdata(db: Any, cutoff: str, limit: int) -> int:
        """
        删除早于截止日期的站点用户数据快照。
        """
        return SiteUserData.delete_before(db=db, before_day=cutoff, limit=limit)

    @staticmethod
    def delete_transfer_history(db: Any, cutoff: str, limit: int) -> int:
        """
        删除早于截止时间的整理历史。
        """
        return TransferHistory.delete_before(
            db=db,
            before_time=cutoff,
            limit=limit,
        )

    @staticmethod
    def delete_download_failures(db: Any, cutoff: str, limit: int) -> int:
        """
        删除已经过期的下载失败冷却记录。
        """
        return DownloadFailure.delete_expired(
            db=db,
            before_time=cutoff,
            limit=limit,
        )

    @staticmethod
    def delete_subscribe_history(db: Any, cutoff: str, limit: int) -> int:
        """
        分批删除超过用户保留期的已完成订阅快照。
        """
        return DatabaseCleanupRepository._delete_selected_ids(
            db=db,
            model=SubscribeHistory,
            condition=SubscribeHistory.date < cutoff,
            limit=limit,
        )

    @staticmethod
    def delete_agent_chats(db: Any, cutoff: str, limit: int) -> int:
        """
        回收旧会话和已确认回执，保留任务引用和投递结果不确定的上下文。
        """
        task_reference = exists(
            select(AgentTask.id).where(AgentTask.session_id == AgentChat.session_id)
        )
        invocation_reference = exists(
            select(AgentInvocation.id).where(
                AgentInvocation.principal_id == AgentChat.user_id,
                AgentInvocation.session_id == AgentChat.session_id,
                AgentInvocation.status.in_(("running", "unknown")),
            )
        )
        ids = list(db.execute(
            select(AgentChat.id).where(
                AgentChat.updated_at < cutoff, ~task_reference, ~invocation_reference,
            ).order_by(AgentChat.id).limit(limit)
        ).scalars())
        if not ids:
            return 0
        execute_dml(
            db,
            delete(AgentInvocation).where(
                AgentInvocation.status.in_(("succeeded", "failed", "pending")),
                exists(select(AgentChat.id).where(
                    AgentChat.id.in_(ids),
                    AgentChat.user_id == AgentInvocation.principal_id,
                    AgentChat.session_id == AgentInvocation.session_id,
                )),
            ),
        )
        return execute_dml(db, delete(AgentChat).where(AgentChat.id.in_(ids)))

    @staticmethod
    def delete_agent_invocations(db: Any, cutoff: str, limit: int) -> int:
        """
        清理无聊天行的旧提交回执和终态历史，不回收投递结果不确定的恢复状态。
        """
        chat_reference = exists(select(AgentChat.id).where(
            AgentChat.user_id == AgentInvocation.principal_id,
            AgentChat.session_id == AgentInvocation.session_id,
        ))
        return DatabaseCleanupRepository._delete_selected_ids(
            db=db,
            model=AgentInvocation,
            condition=(
                AgentInvocation.status.in_(("succeeded", "failed", "pending"))
                & (AgentInvocation.updated_at < cutoff)
                & ~chat_reference
            ),
            limit=limit,
        )

    @staticmethod
    def delete_agent_task_runs(db: Any, cutoff: str, limit: int) -> int:
        """
        清理旧终态运行，但保留运行中和父任务最后一次运行。
        """
        latest_run_reference = exists(
            select(AgentTask.id).where(AgentTask.last_run_id == AgentTaskRun.run_id)
        )
        return DatabaseCleanupRepository._delete_selected_ids(
            db=db,
            model=AgentTaskRun,
            condition=(
                (AgentTaskRun.started_at < cutoff)
                & (AgentTaskRun.status != "running")
                & ~latest_run_reference
            ),
            limit=limit,
        )

    @staticmethod
    def delete_outbox_completed(db: Any, cutoff: str, limit: int) -> int:
        """
        删除过期 completed 记录，事务提交由维护用例统一控制。
        """
        return DatabaseCleanupRepository._delete_outbox_status(
            db=db,
            status="completed",
            timestamp_column=OutboxMessage.completed_at,
            cutoff=cutoff,
            limit=limit,
        )

    @staticmethod
    def delete_outbox_dead(db: Any, cutoff: str, limit: int) -> int:
        """
        删除过期 dead-letter 记录，保留 pending/processing 恢复语义。
        """
        return DatabaseCleanupRepository._delete_outbox_status(
            db=db,
            status="dead",
            timestamp_column=OutboxMessage.next_retry_at,
            cutoff=cutoff,
            limit=limit,
        )

    @staticmethod
    def _delete_outbox_status(
        *,
        db: Any,
        status: str,
        timestamp_column: Any,
        cutoff: str,
        limit: int,
    ) -> int:
        """
        先锁定有限 ID 再删除，避免一次维护事务无界膨胀。
        """
        message_ids = db.execute(
            select(OutboxMessage.id)
            .where(
                OutboxMessage.status == status,
                timestamp_column.is_not(None),
                timestamp_column < cutoff,
            )
            .order_by(OutboxMessage.id)
            .limit(limit)
        ).scalars().all()
        if not message_ids:
            return 0
        return execute_dml(
            db,
            delete(OutboxMessage).where(OutboxMessage.id.in_(message_ids)),
            execution_options={"synchronize_session": False},
        )

    @staticmethod
    def _delete_selected_ids(
        *,
        db: Any,
        model: Any,
        condition: Any,
        limit: int,
    ) -> int:
        """
        按安全谓词锁定有限主键并暂存删除。
        """
        record_ids = db.execute(
            select(model.id)
            .where(condition)
            .order_by(model.id)
            .limit(limit)
        ).scalars().all()
        if not record_ids:
            return 0
        return execute_dml(
            db,
            delete(model).where(model.id.in_(record_ids)),
            execution_options={"synchronize_session": False},
        )