"""Ledger snapshot plugin: maintain a running-total table.
:class:`LedgerSnapshotPlugin` creates a ``{tablename}_snapshot``
table whose rows always reflect ``SUM(value)`` per dimension group.
An ``AFTER INSERT FOR EACH STATEMENT`` trigger on the ledger raw
table applies each batch of inserts as incremental ``UPSERT``
operations, so the snapshot stays in sync without a full re-scan.
This is the performance-optimised alternative to
:func:`~codegen_database.ext.ledger.queries.construct_ledger_balance_query`
when the ledger is large and balance reads are frequent.
Usage::
inventory = CodegenDatabaseLedger(
tablename="inventory",
schemaname="private",
metadata=metadata,
schema_items=[
Column("warehouse", String, nullable=False),
Column("sku", String, nullable=False),
],
extra_plugins=[
LedgerSnapshotPlugin(dimensions=["warehouse", "sku"]),
],
)
# inventory.snapshot_table is a joinable proxy.
After construction the ``snapshot_table`` attribute on the factory
instance holds a joinable SQLAlchemy ``Table`` whose columns are the
declared dimension columns plus ``balance`` (same type as
``value``) and ``updated_at``.
Note: the snapshot does **not** enforce a minimum balance. Use
:class:`~codegen_database.plugins.ledger.LedgerBalanceCheckPlugin` on the
ledger table for that.
"""
from __future__ import annotations
from pathlib import Path
from typing import TYPE_CHECKING
from sqlalchemy import Column, DateTime, Table, func
from sqlalchemy_declarative_extensions import (
register_function,
register_trigger,
)
from sqlalchemy_declarative_extensions.dialects.postgresql import (
Function,
FunctionSecurity,
Trigger,
)
from codegen_database.errors import CodegenDatabaseValidationError
from codegen_database.ext.refresh.plugin import IncrementalRefreshPlugin
from codegen_database.plugin import Dynamic, Plugin, produces, requires
from codegen_database.types import build_value_type
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
from codegen_database.types import ValueType
_TEMPLATES = Path(__file__).resolve().parent / "templates" / "snapshot"
_NAMING_DEFAULTS = {
"snapshot_table": "%(table_name)s_snapshot",
"snapshot_function": "%(schema)s_%(table_name)s_snapshot_upsert",
"snapshot_trigger": "%(schema)s_%(table_name)s_snapshot_upsert",
}
[docs]
@produces("snapshot_table")
@requires(Dynamic("table_key"), "entry_id_column")
class LedgerSnapshotPlugin(Plugin):
"""Maintain a running-total snapshot table for a ledger.
Creates ``{tablename}_snapshot`` with one row per unique
combination of *dimensions* columns. An ``AFTER INSERT FOR EACH
STATEMENT`` trigger incrementally applies each batch of ledger
inserts via ``INSERT ... ON CONFLICT DO UPDATE``, keeping the
snapshot current without a full aggregation scan.
After the plugin runs, ``ctx["snapshot_table"]`` is a joinable
SQLAlchemy :class:`~sqlalchemy.Table` whose columns are the
declared dimensions plus ``balance`` and ``updated_at``.
Args:
dimensions: Column names to group the balance by. Must be
a non-empty list matching columns on the ledger raw table.
value_type: Type for the ``balance`` column. One of
``"integer"``, ``"numeric"``, or ``"decimal"``. Should
match the ledger's ``value_type`` (default ``"integer"``).
precision: Total number of digits when *value_type* is
``"numeric"`` / ``"decimal"``. Should match the ledger's
``precision``. ``None`` leaves the column unconstrained.
scale: Digits after the decimal point. Requires *precision*.
table_key: Key in ``ctx`` for the ledger raw table
(default ``"raw_table"``).
Raises:
CodegenDatabaseValidationError: If *dimensions* is empty,
*value_type* is unrecognised, or precision/scale are
misconfigured.
"""
def __init__(
self,
dimensions: list[str],
value_type: ValueType = "integer",
*,
precision: int | None = None,
scale: int | None = None,
table_key: str = "raw_table",
) -> None:
"""Store configuration."""
if not dimensions:
msg = "dimensions must be a non-empty list"
raise CodegenDatabaseValidationError(msg)
# Validate eagerly so a bad type / precision combo fails at
# construction; the column itself is built fresh in ``run``.
build_value_type(value_type, precision=precision, scale=scale)
self.dimensions = list(dimensions)
self.value_type = value_type
self.precision = precision
self.scale = scale
self.table_key = table_key
[docs]
def run(self, ctx: FactoryContext) -> None:
"""Create the snapshot table and the upsert trigger."""
raw_table = ctx[self.table_key]
schema = ctx.schemaname
raw_fullname = f"{schema}.{raw_table.name}"
snapshot_name = resolve_name(
ctx.metadata,
"snapshot_table",
{"table_name": ctx.tablename, "schema": schema},
_NAMING_DEFAULTS,
)
snapshot_fullname = f"{schema}.{snapshot_name}"
# Build snapshot table columns from the raw table's dim columns.
# Dimension columns form the composite primary key so that
# ON CONFLICT (dim_cols) in the upsert trigger resolves correctly.
balance_col_type = build_value_type(
self.value_type,
precision=self.precision,
scale=self.scale,
)
dim_col_objs: list[Column] = [
Column(d, raw_table.c[d].type, nullable=False, primary_key=True)
for d in self.dimensions
]
snapshot_table = Table(
snapshot_name,
ctx.metadata,
*dim_col_objs,
Column("balance", balance_col_type, nullable=False, default=0),
Column(
"updated_at",
DateTime(timezone=True),
server_default=func.now(),
nullable=False,
),
schema=schema,
)
ctx["snapshot_table"] = snapshot_table
# Render the upsert trigger function body.
dim_cols = ", ".join(self.dimensions)
new_dim_exprs = ", ".join(f"_row.{d}" for d in self.dimensions)
template = load_template(_TEMPLATES / "upsert.plpgsql.mako")
body = template.render(
snapshot_table=snapshot_fullname,
dim_cols=dim_cols,
new_dim_exprs=new_dim_exprs,
)
fn_name = resolve_name(
ctx.metadata,
"snapshot_function",
{"table_name": ctx.tablename, "schema": schema},
_NAMING_DEFAULTS,
)
trigger_name = resolve_name(
ctx.metadata,
"snapshot_trigger",
{"table_name": ctx.tablename, "schema": schema},
_NAMING_DEFAULTS,
)
register_function(
ctx.metadata,
Function(
fn_name,
body,
returns="trigger",
language="plpgsql",
schema=schema,
security=FunctionSecurity.definer,
),
)
register_trigger(
ctx.metadata,
Trigger.after(
"insert",
on=raw_fullname,
execute=f"{schema}.{fn_name}",
name=trigger_name,
)
.for_each_statement()
.referencing_new_table_as("new_entries"),
)
# Incremental refresh function built on the generic
# IncrementalRefreshPlugin. Aggregates new ledger rows
# (created_at >= _since) and upserts into the snapshot.
# Callers can schedule this via pg_cron or call manually
# after bulk loads; the real-time trigger handles live inserts.
raw_dim_cols = ", ".join(f"{raw_fullname}.{d}" for d in self.dimensions)
refresh_body = load_template(
_TEMPLATES / "refresh_body.plpgsql.mako"
).render(
snapshot_table=snapshot_fullname,
raw_table=raw_fullname,
dim_cols=dim_cols,
raw_dim_cols=raw_dim_cols,
)
# snapshot_fullname is schema-derived (not user input) -- S608 FP.
fallback_sql = (
f"SELECT MIN(updated_at) - INTERVAL '1 second'" # noqa: S608
f" FROM {snapshot_fullname}"
)
IncrementalRefreshPlugin(
body=refresh_body,
fallback_sql=fallback_sql,
table_key=self.table_key,
).run(ctx)