"""Generic incremental refresh plugin.
:class:`IncrementalRefreshPlugin` registers a checkpoint-based
``{schema}_{tablename}_refresh(p_since TIMESTAMPTZ DEFAULT NULL)``
function in PostgreSQL. Callers invoke this function from ``pg_cron``
(or manually after bulk loads) to roll forward an aggregate table.
The refresh function resolves the starting cutoff as::
_since = COALESCE(p_since, fallback_sql, '-infinity'::TIMESTAMPTZ)
where *fallback_sql* is a SQL expression (supplied by the caller) that
reads the last-processed watermark, e.g.::
SELECT MIN(updated_at) - INTERVAL '1 second' FROM my_snapshot
The user-supplied *body* is PL/pgSQL that runs after ``_since`` is
resolved. It typically aggregates rows from the source table where
``watermark_col >= _since`` and upserts results into an aggregate table.
Usage::
orders = CodegenDatabaseSimple(
"orders",
"app",
metadata,
schema_items=[
Column("customer_id", Integer, nullable=False),
Column("amount", Numeric, nullable=False),
],
extra_plugins=[
IncrementalRefreshPlugin(
body=\"\"\"
INSERT INTO app.orders_daily_agg (customer_id, day, total)
SELECT customer_id,
date_trunc('day', created_at),
SUM(amount)
FROM app.orders_raw
WHERE created_at >= _since
GROUP BY 1, 2
ON CONFLICT (customer_id, day)
DO UPDATE SET total = orders_daily_agg.total + EXCLUDED.total;
\"\"\",
fallback_sql=(
"SELECT MIN(day::timestamptz)"
" - INTERVAL '1 second'"
" FROM app.orders_daily_agg"
),
schedule="*/5 * * * *",
),
],
)
# The above auto-registers a CronJob named "app.app_orders_refresh"
# that runs every 5 minutes. Alembic autogenerate emits:
# op.execute("SELECT cron.schedule(...)")
# Requires pg_cron to be installed. Register it alongside the
# factory:
# register_pg_extension(metadata, PGExtension("pg_cron", cascade=True))
"""
from __future__ import annotations
from pathlib import Path
from typing import TYPE_CHECKING
from sqlalchemy_declarative_extensions import register_function
from sqlalchemy_declarative_extensions.dialects.postgresql import (
Function,
FunctionParam,
FunctionSecurity,
)
from codegen_database.ext.cron.base import CronJob, register_cron_job
from codegen_database.pg_extension import assert_pg_extension_declared
from codegen_database.plugin import Dynamic, Plugin, produces, requires
from codegen_database.utils.naming import resolve_name
from codegen_database.utils.template import load_template
if TYPE_CHECKING:
from codegen_database.factory.context import FactoryContext
_TEMPLATES = Path(__file__).resolve().parent / "templates"
_NAMING_DEFAULTS = {
"refresh_function": "%(schema)s_%(table_name)s_refresh",
}
[docs]
@produces("refresh_function")
@requires(Dynamic("table_key"))
class IncrementalRefreshPlugin(Plugin):
r"""Register a checkpoint-based incremental refresh function.
Generates and registers a PostgreSQL function with signature::
{schema}.{schema}_{tablename}_refresh(
p_since TIMESTAMPTZ DEFAULT NULL
) RETURNS void
The function resolves ``_since`` as the caller-supplied *p_since*,
or falls back to the expression in *fallback_sql*. The
user-supplied *body* is embedded after ``_since`` is resolved and
can reference ``_since`` directly.
Args:
body: PL/pgSQL fragment (no ``DECLARE`` or ``BEGIN``/``END``
wrapper). Must be valid PL/pgSQL and may reference
``_since TIMESTAMPTZ``.
fallback_sql: SQL expression that returns a ``TIMESTAMPTZ``
used as the checkpoint when *p_since* is ``NULL``. A
``NULL`` result falls back to ``'-infinity'``. Typically
reads ``MIN(updated_at) - INTERVAL '1 second'`` from the
aggregate table.
table_key: Key in ``ctx`` for the source table (used only to
satisfy the ``@requires`` dependency). Default
``"raw_table"``.
"""
def __init__(
self,
body: str,
fallback_sql: str,
*,
schedule: str | None = None,
table_key: str = "raw_table",
) -> None:
"""Store configuration."""
self.body = body
self.fallback_sql = fallback_sql
self.schedule = schedule
self.table_key = table_key
[docs]
def run(self, ctx: FactoryContext) -> None:
"""Render and register the incremental refresh function."""
schema = ctx.schemaname
table_name = ctx.tablename
fn_name = resolve_name(
ctx.metadata,
"refresh_function",
{"table_name": table_name, "schema": schema},
_NAMING_DEFAULTS,
)
template = load_template(_TEMPLATES / "incremental.plpgsql.mako")
body = template.render(
fallback_sql=self.fallback_sql,
body=self.body,
)
register_function(
ctx.metadata,
Function(
fn_name,
body,
returns="void",
language="plpgsql",
schema=schema,
security=FunctionSecurity.definer,
parameters=[
FunctionParam.input("p_since", "TIMESTAMPTZ DEFAULT NULL")
],
),
)
ctx["refresh_function"] = fn_name
if self.schedule is not None:
assert_pg_extension_declared(ctx.metadata, "pg_cron")
register_cron_job(
ctx.metadata,
CronJob(
name=f"{schema}.{fn_name}",
schedule=self.schedule,
command=f"SELECT {schema}.{fn_name}()",
),
)