Source code for codegen_database.ext.cron.compare

"""Alembic comparator ops and diff logic for pg_cron jobs."""

from __future__ import annotations

from dataclasses import dataclass
from typing import TYPE_CHECKING

from alembic.operations.ops import MigrateOperation
from sqlalchemy import text

from codegen_database.ext.cron.base import CronJob

if TYPE_CHECKING:
    from sqlalchemy import Connection

    from codegen_database.ext.cron.base import CronJobs

_FETCH_JOBS_SQL = "SELECT jobname, schedule, command FROM cron.job"




@dataclass
class ScheduleCronJobOp(MigrateOperation):
    """Alembic operation: create or update a pg_cron job."""

    job: CronJob



    def to_sql(self) -> list[str]:
        """Return the SQL statement for this operation."""
        return [self.job.to_sql_schedule()]




    def reverse(self) -> ScheduleCronJobOp:
        """Return a no-op downgrade: jobs are never dropped automatically.

        ``cron.schedule`` is an upsert so re-running it on downgrade
        is always safe.  Users who want to remove a job on downgrade
        should add an explicit :class:`UnscheduleCronJobOp` call.
        """
        return self






@dataclass
class UnscheduleCronJobOp(MigrateOperation):
    """Alembic operation: remove a pg_cron job.

    This op is available for manual migration authoring.
    The autogenerate comparator never emits it automatically to
    avoid accidentally removing externally-managed jobs.
    """

    name: str



    def to_sql(self) -> list[str]:
        """Return the SQL statement for this operation."""
        return [
            CronJob(name=self.name, schedule="", command="").to_sql_unschedule()
        ]




def _pg_cron_installed(connection: Connection) -> bool:
    """Return ``True`` if pg_cron's schema exists in the database.

    Args:
        connection: Active database connection.

    Returns:
        ``True`` when the ``cron`` schema is present.

    """
    row = connection.execute(
        text("SELECT 1 FROM pg_namespace WHERE nspname = 'cron'")
    ).scalar()
    return row is not None


def _fetch_current_jobs(connection: Connection) -> dict[str, dict[str, str]]:
    """Return current pg_cron jobs keyed by job name.

    Args:
        connection: Active database connection.

    Returns:
        Mapping from ``jobname`` to a dict with ``schedule`` and
        ``command`` keys.

    """
    rows = connection.execute(text(_FETCH_JOBS_SQL)).fetchall()
    return {
        row.jobname: {
            "schedule": row.schedule,
            "command": row.command,
        }
        for row in rows
    }




def compare_cron_jobs(
    connection: Connection,
    desired: CronJobs,
) -> list[ScheduleCronJobOp]:
    """Diff desired cron jobs against the current database state.

    Only emits :class:`ScheduleCronJobOp` (create or update) — jobs
    are never dropped automatically.  Use :class:`UnscheduleCronJobOp`
    in a hand-written migration to remove a job.

    Args:
        connection: Active database connection.
        desired: The :class:`~codegen_database.ext.cron.base.CronJobs`
            holder read from metadata.

    Returns:
        List of :class:`ScheduleCronJobOp` operations to bring
        declared jobs into sync with the database.

    """
    if not _pg_cron_installed(connection):
        return []

    current = _fetch_current_jobs(connection)
    ops: list[ScheduleCronJobOp] = []

    for job in desired.jobs:
        cur = current.get(job.name)

        if (
            cur is None
            or cur["schedule"] != job.schedule
            or cur["command"] != job.command
        ):
            ops.append(ScheduleCronJobOp(job=job))

    return ops