Source code for codegen_database.ext.cron.base

"""pg_cron job dataclasses and metadata registration."""

from __future__ import annotations

from dataclasses import dataclass, field

from sqlalchemy import MetaData




@dataclass(frozen=True)
class CronJob:
    """Describe a pg_cron job to register.

    Args:
        name: Unique job name used as the identity key in ``cron.job``.
        schedule: Cron expression, e.g. ``"*/5 * * * *"``.
        command: SQL to execute, e.g. ``"SELECT app.app_orders_refresh()"``.
        database: Database to run the command in.  When ``None``, the
            current database is used.

    """

    name: str
    schedule: str
    command: str
    database: str | None = None



    def to_sql_schedule(self) -> str:
        """Render ``SELECT cron.schedule(...)`` DDL.

        Returns:
            A complete ``SELECT cron.schedule(...)`` SQL string.

        """
        db_clause = (
            f", database := '{self.database}'"
            if self.database is not None
            else ""
        )
        return (
            f"SELECT cron.schedule("
            f"'{self.name}', "
            f"'{self.schedule}', "
            f"$${self.command}$${db_clause}"
            f")"
        )




    def to_sql_unschedule(self) -> str:
        """Render ``SELECT cron.unschedule(...)`` DDL.

        Returns:
            A complete ``SELECT cron.unschedule(...)`` SQL string.

        """
        return f"SELECT cron.unschedule('{self.name}')"






@dataclass
class CronJobs:
    """Container for all cron jobs registered on a MetaData.

    Stored under ``metadata.info["cron_jobs"]`` by
    :func:`register_cron_job`.
    """

    jobs: list[CronJob] = field(default_factory=list)



    @classmethod
    def extract(cls, metadata: object) -> CronJobs | None:
        """Return registered cron jobs or ``None`` if none exist.

        Args:
            metadata: SQLAlchemy :class:`~sqlalchemy.MetaData`, or
                ``None``.

        Returns:
            The :class:`CronJobs` holder or ``None``.

        """
        if metadata is None or not isinstance(metadata, MetaData):
            return None

        return metadata.info.get("cron_jobs")






def register_cron_job(
    metadata: MetaData,
    job: CronJob,
) -> None:
    """Store *job* in *metadata* for Alembic autogenerate.

    Args:
        metadata: SQLAlchemy :class:`~sqlalchemy.MetaData` to
            register on.
        job: The cron job to register.

    """
    holder: CronJobs | None = metadata.info.get("cron_jobs")

    if holder is None:
        holder = CronJobs()
        metadata.info["cron_jobs"] = holder

    holder.jobs.append(job)