"""User-facing plugins for ledger (append-only value) tables."""
from __future__ import annotations
from pathlib import Path
from typing import TYPE_CHECKING
from sqlalchemy import (
CheckConstraint,
Column,
String,
)
from sqlalchemy_declarative_extensions import (
register_function,
register_trigger,
)
from sqlalchemy_declarative_extensions.dialects.postgresql import (
Function,
FunctionSecurity,
Trigger,
)
if TYPE_CHECKING:
from codegen_database.factory.context import FactoryContext
from codegen_database.errors import CodegenDatabaseValidationError
from codegen_database.plugin import (
Dynamic,
Plugin,
produces,
requires,
singleton,
)
from codegen_database.utils.naming import resolve_name
from codegen_database.utils.template import load_template
_TEMPLATES = Path(__file__).resolve().parent / "templates" / "ledger"
_NAMING_DEFAULTS = {
"ledger_function": "%(schema)s_%(table_name)s_%(op)s",
"ledger_trigger": "%(schema)s_%(table_name)s_%(op)s",
"balance_check_function": "%(schema)s_%(table_name)s_%(op)s",
"balance_check_trigger": "%(schema)s_%(table_name)s_%(op)s",
"double_entry_function": "%(schema)s_%(table_name)s_%(op)s",
"double_entry_trigger": "%(schema)s_%(table_name)s_%(op)s",
}
DIRECTION_DEBIT = "debit"
DIRECTION_CREDIT = "credit"
_VALID_DIRECTIONS = (DIRECTION_DEBIT, DIRECTION_CREDIT)
def _require_dimensions(dimensions: list[str]) -> None:
"""Raise if *dimensions* is empty.
Args:
dimensions: The list to validate.
Raises:
CodegenDatabaseValidationError: If *dimensions* is empty.
"""
if not dimensions:
msg = "dimensions must be a non-empty list"
raise CodegenDatabaseValidationError(msg)
[docs]
@requires(Dynamic("table_key"))
class LedgerBalanceCheckPlugin(Plugin):
"""Enforce a minimum balance per dimension group.
Registers an ``AFTER INSERT FOR EACH STATEMENT`` trigger
that checks ``SUM(value) >= min_balance`` for every
dimension group affected by the inserted rows. If any
group violates the constraint the entire statement is
rejected.
Uses the same ``REFERENCING NEW TABLE`` transition-table
pattern as
:class:`~codegen_database.plugins.ledger.DoubleEntryTriggerPlugin`.
Args:
dimensions: Column names that define a balance group.
Must be a non-empty list.
min_balance: The minimum allowed ``SUM(value)`` per
group (default ``0``).
table_key: Key in ``ctx`` for the backing table
(default ``"primary"``).
Raises:
CodegenDatabaseValidationError: If *dimensions* is empty.
"""
def __init__(
self,
dimensions: list[str],
min_balance: int = 0,
table_key: str = "raw_table",
) -> None:
"""Store configuration."""
_require_dimensions(dimensions)
self.dimensions = list(dimensions)
self.min_balance = min_balance
self.table_key = table_key
[docs]
def run(self, ctx: FactoryContext) -> None:
"""Register the balance-check trigger."""
table = ctx[self.table_key]
table_fullname = f"{ctx.schemaname}.{table.name}"
dim_cols = ", ".join(self.dimensions)
dim_format = ", ".join("%" for _ in self.dimensions)
dim_values = ", ".join(f"_bad.{d}" for d in self.dimensions)
template = load_template(_TEMPLATES / "balance_check.plpgsql.mako")
body = template.render(
table=table_fullname,
dim_cols=dim_cols,
dim_format=dim_format,
dim_values=dim_values,
min_balance=self.min_balance,
)
fn_name = resolve_name(
ctx.metadata,
"balance_check_function",
{
"table_name": ctx.tablename,
"schema": ctx.schemaname,
"op": "balance_check",
},
_NAMING_DEFAULTS,
)
trigger_name = resolve_name(
ctx.metadata,
"balance_check_trigger",
{
"table_name": ctx.tablename,
"schema": ctx.schemaname,
"op": "balance_check",
},
_NAMING_DEFAULTS,
)
register_function(
ctx.metadata,
Function(
fn_name,
body,
returns="trigger",
language="plpgsql",
schema=ctx.schemaname,
security=FunctionSecurity.definer,
),
)
register_trigger(
ctx.metadata,
Trigger.after(
"insert",
on=table_fullname,
execute=f"{ctx.schemaname}.{fn_name}",
name=trigger_name,
)
.for_each_statement()
.referencing_new_table_as("new_entries"),
)
[docs]
@produces("double_entry_columns")
@singleton("__double_entry__")
class DoubleEntryPlugin(Plugin):
"""Add debit/credit semantics to a ledger table.
Adds a ``direction`` column (``'debit'`` or ``'credit'``)
to the schema items so that ``LedgerTablePlugin`` includes
it in the table. Also registers an AFTER INSERT constraint
trigger that validates all rows sharing an ``entry_id``
have equal total debits and credits.
This plugin must appear **before** ``LedgerTablePlugin`` in
the plugin list so its column is included in the table
definition.
Args:
column_name: Name of the direction column
(default ``"direction"``).
"""
def __init__(
self,
column_name: str = "direction",
) -> None:
"""Store configuration."""
self._column_name = column_name
[docs]
def run(self, ctx: FactoryContext) -> None:
"""Inject the direction column and store its name."""
valid = ", ".join(f"'{v}'" for v in _VALID_DIRECTIONS)
ctx.injected_columns.append(
Column(
self._column_name,
String,
CheckConstraint(
f"{self._column_name} IN ({valid})",
),
nullable=False,
),
)
ctx["double_entry_columns"] = self._column_name
[docs]
@requires(
Dynamic("table_key"),
"double_entry_columns",
"entry_id_column",
)
class DoubleEntryTriggerPlugin(Plugin):
"""Register an AFTER INSERT trigger enforcing balanced entries.
Validates that for every ``entry_id`` in the inserted batch,
the sum of debit values equals the sum of credit values.
Raises a PostgreSQL exception if any entry is unbalanced.
Uses a statement-level trigger with a ``REFERENCING NEW TABLE``
transition table so that multi-row inserts are checked as a
whole, not row-by-row.
Must run **after** ``LedgerTablePlugin`` (needs the table)
and **after** ``DoubleEntryPlugin`` (needs column name).
Args:
table_key: Key in ``ctx`` for the backing table
(default ``"primary"``).
"""
def __init__(
self,
table_key: str = "raw_table",
) -> None:
"""Store the context key."""
self.table_key = table_key
[docs]
def run(self, ctx: FactoryContext) -> None:
"""Register the constraint trigger on the ledger table."""
table = ctx[self.table_key]
direction_col = ctx["double_entry_columns"]
entry_id_col = ctx["entry_id_column"]
table_fullname = f"{ctx.schemaname}.{table.name}"
template = load_template(_TEMPLATES / "double_entry_check.plpgsql.mako")
body = template.render(
table=table_fullname,
direction_col=direction_col,
entry_id_col=entry_id_col.name,
debit=DIRECTION_DEBIT,
credit=DIRECTION_CREDIT,
)
fn_name = resolve_name(
ctx.metadata,
"double_entry_function",
{
"table_name": ctx.tablename,
"schema": ctx.schemaname,
"op": "double_entry_check",
},
_NAMING_DEFAULTS,
)
trigger_name = resolve_name(
ctx.metadata,
"double_entry_trigger",
{
"table_name": ctx.tablename,
"schema": ctx.schemaname,
"op": "double_entry_check",
},
_NAMING_DEFAULTS,
)
register_function(
ctx.metadata,
Function(
fn_name,
body,
returns="trigger",
language="plpgsql",
schema=ctx.schemaname,
security=FunctionSecurity.definer,
),
)
register_trigger(
ctx.metadata,
Trigger.after(
"insert",
on=table_fullname,
execute=f"{ctx.schemaname}.{fn_name}",
name=trigger_name,
)
.for_each_statement()
.referencing_new_table_as("new_entries"),
)