activity-core/src/activity_core/sync_schedules.py

"""Bootstrap script: sync Temporal Schedules with the ActivityDefinition DB.

T23: On startup, ensures every enabled cron ActivityDefinition has a live
Temporal Schedule, and removes orphaned schedules that have no matching DB row.

Run directly:
    ACTCORE_DB_URL=... uv run python -m activity_core.sync_schedules

Also called from worker.py before the worker enters its run loop.
"""

from __future__ import annotations

import asyncio
import logging
import os
import uuid
from dataclasses import dataclass
from typing import Sequence

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
from temporalio.client import Client

from activity_core.models import ActivityDefinition, CronTriggerConfig, ScheduledTriggerConfig
from activity_core.orm import ActivityDefinition as ActivityDefinitionRow
from activity_core.schedule_manager import (
    cancel_scheduled,
    delete_schedule,
    list_schedules,
    upsert_schedule,
)

logger = logging.getLogger(__name__)

TEMPORAL_HOST = os.environ.get("TEMPORAL_HOST", "localhost:7233")
TEMPORAL_NAMESPACE = os.environ.get("TEMPORAL_NAMESPACE", "default")


@dataclass
class ScheduleSyncResult:
    upserted: int = 0
    paused: int = 0
    deleted_orphans: int = 0
    errors: int = 0
    error_details: list[str] | None = None

    def to_dict(self) -> dict[str, int]:
        return {
            "upserted": self.upserted,
            "paused": self.paused,
            "deleted_orphans": self.deleted_orphans,
            "errors": self.errors,
        }


def _row_to_domain(row: ActivityDefinitionRow) -> ActivityDefinition:
    """Convert an ORM row to a domain ActivityDefinition for schedule_manager."""
    return ActivityDefinition.model_validate(
        {
            "id": row.id,
            "name": row.name,
            "enabled": row.enabled,
            "trigger_config": row.trigger_config,
            "context_sources": row.context_sources,
            "task_templates": row.task_templates,
            "dedupe_key_strategy": row.dedupe_key_strategy,
            "version": row.version,
        }
    )


def _valid_schedule_activity_id(defn: ActivityDefinition) -> str:
    if isinstance(defn.trigger_config, ScheduledTriggerConfig):
        return f"{defn.id}-once"
    return str(defn.id)


async def _load_schedule_rows(
    session_factory: async_sessionmaker[AsyncSession],
) -> Sequence[ActivityDefinitionRow]:
    async with session_factory() as session:
        return (
            await session.scalars(
                select(ActivityDefinitionRow).where(
                    ActivityDefinitionRow.trigger_type.in_(["cron", "scheduled"])
                )
            )
        ).all()


async def sync_schedule_rows(
    client: Client,
    rows: Sequence[ActivityDefinitionRow],
) -> ScheduleSyncResult:
    """Reconcile Temporal Schedules against already-loaded definition rows.

    ACTIVITY-WP-0021-T06: one failing upsert (e.g. ScheduleAlreadyRunningError
    on pause/unpause) must not abort the remaining rows.
    """
    valid_schedule_activity_ids: set[str] = set()
    result = ScheduleSyncResult(error_details=[])

    for row in rows:
        defn = _row_to_domain(row)
        if not isinstance(
            defn.trigger_config,
            (CronTriggerConfig, ScheduledTriggerConfig),
        ):
            continue

        valid_schedule_activity_ids.add(_valid_schedule_activity_id(defn))

        try:
            if isinstance(defn.trigger_config, ScheduledTriggerConfig) and not defn.enabled:
                # A disabled one-shot has no future recurrence to preserve. A
                # completed Temporal schedule can reject in-place updates with
                # ScheduleAlreadyRunningError, leaving every later sync noisy.
                # Delete it idempotently instead; re-enabling recreates it.
                await cancel_scheduled(client, defn.id)
            else:
                await upsert_schedule(client, defn)
        except Exception as exc:  # noqa: BLE001 — continue reconcile for other rows
            result.errors += 1
            detail = f"{defn.id} ({defn.name}): {type(exc).__name__}: {exc}"
            result.error_details.append(detail)
            logger.error("upsert_schedule failed for activity %s — continuing: %s", defn.id, exc)
            continue

        if defn.enabled:
            result.upserted += 1
            logger.info("upserted schedule for activity %s (%s)", defn.id, defn.name)
        else:
            result.paused += 1
            if isinstance(defn.trigger_config, ScheduledTriggerConfig):
                logger.info("removed schedule for disabled one-shot activity %s", defn.id)
            else:
                logger.info("upserted paused schedule for disabled activity %s", defn.id)

    # Tombstone cleanup: remove Temporal Schedules with no matching DB row.
    try:
        existing_schedules = await list_schedules(client)
    except Exception as exc:  # noqa: BLE001
        result.errors += 1
        detail = f"list_schedules: {type(exc).__name__}: {exc}"
        result.error_details.append(detail)
        logger.error("list_schedules failed — skipping orphan cleanup: %s", exc)
        existing_schedules = []

    for entry in existing_schedules:
        if entry["activity_id"] not in valid_schedule_activity_ids:
            try:
                await delete_schedule(client, entry["activity_id"])
                result.deleted_orphans += 1
                logger.info("deleted orphaned schedule %s", entry["schedule_id"])
            except Exception as exc:  # noqa: BLE001
                result.errors += 1
                detail = f"delete {entry['schedule_id']}: {type(exc).__name__}: {exc}"
                result.error_details.append(detail)
                logger.error("delete_schedule failed for %s — continuing: %s", entry["schedule_id"], exc)

    logger.info(
        "sync_schedules complete — upserted=%d paused=%d deleted_orphans=%d errors=%d",
        result.upserted,
        result.paused,
        result.deleted_orphans,
        result.errors,
    )
    return result


async def sync_with_session_factory(
    client: Client,
    session_factory: async_sessionmaker[AsyncSession],
) -> ScheduleSyncResult:
    """Reconcile Temporal Schedules using an existing DB session factory."""
    return await sync_schedule_rows(client, await _load_schedule_rows(session_factory))


async def sync(client: Client, db_url: str) -> ScheduleSyncResult:
    """Reconcile Temporal Schedules against the ActivityDefinition table.

    Steps:
      1. Load all cron/scheduled ActivityDefinitions from Postgres.
      2. Upsert a Temporal Schedule for each one, paused when disabled.
      3. Delete Temporal Schedules whose activity_id has no matching DB row
         (tombstone cleanup for deleted or trigger-type-changed definitions).
    """
    engine = create_async_engine(db_url)
    session_factory = async_sessionmaker(engine, expire_on_commit=False)

    try:
        return await sync_with_session_factory(client, session_factory)
    finally:
        await engine.dispose()


async def main() -> None:
    logging.basicConfig(level=logging.INFO)
    db_url = os.environ.get("ACTCORE_DB_URL")
    if not db_url:
        raise RuntimeError("ACTCORE_DB_URL is required")

    client = await Client.connect(TEMPORAL_HOST, namespace=TEMPORAL_NAMESPACE)
    result = await sync(client, db_url)
    print(
        "Synced schedules: "
        f"upserted={result.upserted} "
        f"paused={result.paused} "
        f"deleted_orphans={result.deleted_orphans} "
        f"errors={result.errors}"
    )
    if result.error_details:
        for detail in result.error_details:
            print(f"  error: {detail}")


if __name__ == "__main__":
    asyncio.run(main())