Source code for codegen_database.ext.cron.base

from __future__ import annotations

from dataclasses import dataclass, field
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from sqlalchemy import MetaData




@dataclass(frozen=True)
class CronJob:
    name: str
    schedule: str
    command: str
    database: str | None = None

    def to_sql_schedule(self) -> str:
        name = self.name.replace("'", "''")
        schedule = self.schedule.replace("'", "''")
        command = self.command
        db = self.database.replace("'", "''") if self.database else None
        db_clause = f", database := '{db}'" if db is not None else ""
        return (
            f"SELECT cron.schedule("
            f"'{name}', "
            f"'{schedule}', "
            f"$${command}$${db_clause}"
            f")"
        )

    def to_sql_unschedule(self) -> str:
        name = self.name.replace("'", "''")
        return f"SELECT cron.unschedule('{name}')"





@dataclass
class CronJobs:
    jobs: list[CronJob] = field(default_factory=list)

    @classmethod
    def extract(cls, metadata: MetaData) -> CronJobs:
        jobs = metadata.info.get("cron_jobs", CronJobs())
        metadata.info["cron_jobs"] = jobs
        return jobs



def register_cron_job(
    metadata: MetaData,
    job: CronJob,
) -> None:
    CronJobs.extract(metadata).jobs.append(job)