deer-flow/backend/app/infra/schedule/run_sql.py

"""SQL adapter for the schedule context's execution-record repository.

Secondary adapter implementing `ScheduledRunRepository`. Queries are migrated
unchanged from the legacy repository.

The load-bearing piece is the `IntegrityError` translation in `add`: the
partial unique index `uq_scheduled_task_run_active` is the atomic arbiter of
"at most one active run per task", and turning its rejection into
`ActiveRunConflictError` is what lets the service collapse a lost race into
exactly the same outcome as its own non-atomic fast path.
"""

from __future__ import annotations

from datetime import UTC, datetime

from sqlalchemy import func, select
from sqlalchemy.exc import IntegrityError
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker

from deerflow.domain.schedule.model import (
    ACTIVE_RUN_STATUSES,
    TERMINAL_RUN_STATUSES,
    ActiveRunConflictError,
    RunStatus,
    ScheduledRun,
    TriggerKind,
)
from deerflow.domain.schedule.ports import ScheduledRunRepository

# Transitional: the ORM row stays in the harness until engine, models, and
# migrations move into app/infra together.
from deerflow.persistence.scheduled_task_runs.model import ScheduledTaskRunRow

_ACTIVE_STATUS_VALUES = tuple(str(status) for status in ACTIVE_RUN_STATUSES)


def _tz_aware(value: datetime | None) -> datetime | None:
    """SQLite drops tzinfo on read; stored values are always UTC."""
    return value if value is None or value.tzinfo is not None else value.replace(tzinfo=UTC)


class SqlScheduledRunRepository(ScheduledRunRepository):
    """SQL implementation of the `ScheduledRunRepository` port."""

    def __init__(self, session_factory: async_sessionmaker[AsyncSession]) -> None:
        self._sf = session_factory

    @staticmethod
    def _to_domain(row: ScheduledTaskRunRow) -> ScheduledRun:
        return ScheduledRun(
            record_id=row.id,
            task_id=row.task_id,
            thread_id=row.thread_id,
            scheduled_for=_tz_aware(row.scheduled_for),
            trigger=TriggerKind(row.trigger),
            status=RunStatus(row.status),
            run_id=row.run_id,
            error=row.error,
            started_at=_tz_aware(row.started_at),
            finished_at=_tz_aware(row.finished_at),
            created_at=_tz_aware(row.created_at),
        )

    async def add(self, run: ScheduledRun) -> ScheduledRun:
        row = ScheduledTaskRunRow(
            id=run.record_id,
            task_id=run.task_id,
            thread_id=run.thread_id,
            run_id=run.run_id,
            scheduled_for=run.scheduled_for,
            trigger=str(run.trigger),
            status=str(run.status),
            error=run.error,
            started_at=run.started_at,
            finished_at=run.finished_at,
            created_at=run.created_at,
        )
        async with self._sf() as session:
            session.add(row)
            try:
                await session.commit()
            except IntegrityError:
                await session.rollback()
                # Only an active-status insert can trip the partial unique
                # index; a terminal row (e.g. a skip tombstone) is outside its
                # predicate and cannot conflict, so an IntegrityError there is
                # a genuine fault and is re-raised untranslated.
                if run.is_active:
                    raise ActiveRunConflictError(f"scheduled task {run.task_id!r} already has an active run") from None
                raise
            await session.refresh(row)
            return self._to_domain(row)

    async def list_by_task(self, task_id: str, *, limit: int = 50, offset: int = 0) -> list[ScheduledRun]:
        stmt = (
            select(ScheduledTaskRunRow)
            .where(ScheduledTaskRunRow.task_id == task_id)
            .order_by(
                ScheduledTaskRunRow.created_at.desc(),
                ScheduledTaskRunRow.id.desc(),
            )
            .limit(limit)
            .offset(offset)
        )
        async with self._sf() as session:
            result = await session.execute(stmt)
            return [self._to_domain(row) for row in result.scalars()]

    async def count_active(self) -> int:
        """Global count of active rows, used to bound cross-task concurrency."""
        stmt = select(func.count()).select_from(ScheduledTaskRunRow).where(ScheduledTaskRunRow.status.in_(_ACTIVE_STATUS_VALUES))
        async with self._sf() as session:
            result = await session.execute(stmt)
            return int(result.scalar() or 0)

    async def has_active(self, task_id: str) -> bool:
        stmt = (
            select(ScheduledTaskRunRow.id)
            .where(
                ScheduledTaskRunRow.task_id == task_id,
                ScheduledTaskRunRow.status.in_(_ACTIVE_STATUS_VALUES),
            )
            .limit(1)
        )
        async with self._sf() as session:
            result = await session.execute(stmt)
            return result.scalars().first() is not None

    async def update_status(
        self,
        record_id: str,
        *,
        status: RunStatus,
        run_id: str | None = None,
        error: str | None = None,
        started_at: datetime | None = None,
        finished_at: datetime | None = None,
        protect_terminal: bool = False,
    ) -> None:
        async with self._sf() as session:
            row = await session.get(ScheduledTaskRunRow, record_id)
            if row is None:
                return
            if protect_terminal and RunStatus(row.status) in TERMINAL_RUN_STATUSES:
                # The launch-path "running" write lost the race against the
                # completion hook; keep the terminal status/error and only
                # backfill bookkeeping the completion write could not know.
                if row.run_id is None and run_id is not None:
                    row.run_id = run_id
                if row.started_at is None and started_at is not None:
                    row.started_at = started_at
                await session.commit()
                return
            row.status = str(status)
            row.run_id = run_id
            row.error = error
            if started_at is not None:
                row.started_at = started_at
            if finished_at is not None:
                row.finished_at = finished_at
            await session.commit()

    async def mark_stale_active(self, *, error: str) -> int:
        """Fail-fast bookkeeping for runs orphaned by a process crash.

        Agent runs execute in-process, so any active row found at scheduler
        startup belongs to a run whose process is gone. Only valid under the
        single-scheduler-instance assumption.
        """
        stmt = select(ScheduledTaskRunRow).where(ScheduledTaskRunRow.status.in_(_ACTIVE_STATUS_VALUES))
        now = datetime.now(UTC)
        async with self._sf() as session:
            result = await session.execute(stmt)
            rows = list(result.scalars())
            for row in rows:
                row.status = str(RunStatus.INTERRUPTED)
                row.error = error
                row.finished_at = now
            await session.commit()
            return len(rows)