Source code for codegen_database.ext.ledger.events

"""LedgerEvent configuration and helpers.

A :class:`LedgerEvent` declares a named operation on a ledger.  The
user provides lambdas that produce SQLAlchemy selects; the plugin
compiles them into a single PostgreSQL function per event.

Two modes are supported:

- **Simple mode** (``input`` only): the input select is inserted
  directly into the ledger view.
- **Diff mode** (``input`` + ``desired`` + ``existing``): the desired
  state is diffed against the existing state and only the correcting
  deltas are inserted.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import TYPE_CHECKING

from sqlalchemy import (
    column,
    func,
    literal_column,
    select,
    table,
)
from sqlalchemy.dialects import postgresql as pg_dialect
from sqlalchemy.sql.expression import tuple_
from sqlalchemy_declarative_extensions.dialects.postgresql import FunctionParam

if TYPE_CHECKING:
    from collections.abc import Callable

    from sqlalchemy import FromClause, Table
    from sqlalchemy.sql.elements import ColumnClause, ColumnElement
    from sqlalchemy.sql.expression import Select
    from sqlalchemy.sql.selectable import SelectBase
    from sqlalchemy.sql.sqltypes import TypeEngine


[docs] class ParamCollector: """Collect SQL function parameters during lambda evaluation. Usage inside an ``input`` lambda:: lambda p: select( p("warehouse", String).label("warehouse"), p("sku", String).label("sku"), ) Each call to ``p(name, sa_type)`` records the parameter and returns a :func:`~sqlalchemy.sql.expression.literal_column` reference (``p_name``) suitable for embedding in a select. Args: None. """ def __init__(self) -> None: """Initialise with empty parameter list.""" self._params: list[tuple[str, TypeEngine]] = [] def __call__( self, name: str, sa_type: type[TypeEngine] | TypeEngine, ) -> ColumnClause[str]: """Register a parameter and return a column reference. Args: name: Parameter name (without ``p_`` prefix). sa_type: SQLAlchemy type (class or instance). Returns: A ``literal_column("p_name")`` for use in selects. """ if isinstance(sa_type, type): sa_type = sa_type() self._params.append((name, sa_type)) return literal_column(f"p_{name}") @property def function_params(self) -> list[FunctionParam]: """Build the ``FunctionParam`` list for function registration. Returns: List of ``FunctionParam.input(...)`` entries. """ return [ FunctionParam.input( f"p_{name}", sa_type.compile(dialect=pg_dialect.dialect()), ) for name, sa_type in self._params ]
[docs] @dataclass class LedgerEvent: """Declare a named ledger operation. Each event compiles into a single PostgreSQL function that inserts rows into the ledger view and returns the inserted rows via ``RETURNING *``. **Simple mode** — provide only ``input``. The input select's columns are inserted directly. **Diff mode** — provide ``input``, ``desired``, ``existing``, and ``diff_keys``. The desired and existing selects are unioned and only non-zero deltas are inserted. Args: name: Unique event name within the ledger. input: Lambda ``(p) -> Select`` that builds the input CTE. ``p`` is a :class:`ParamCollector`. desired: Lambda ``(pginput) -> SelectBase`` that builds the desired-state CTE. ``pginput`` is a synthetic table reference to the ``input`` CTE -- read function-param values as ``pginput.c.<name>``, no ``p_`` prefix. May return a ``union_all`` or other compound select. existing: Lambda ``(table, desired, pginput) -> Select`` that builds the existing-state CTE. ``pginput`` is the same input-CTE reference passed to ``desired``, so you can reach function params the same way -- ``pginput.c.<name>`` -- in the predicate or join. Use :func:`construct_ledger_balances_query` (filter by diff-key tuples in ``desired``) or :func:`construct_ledger_scoped_balances_query` (filter by an owner predicate) for the common patterns. diff_keys: Column names used for grouping in diff mode. Required when ``desired`` is set. """ name: str input: Callable[[ParamCollector], Select] desired: Callable[[FromClause], SelectBase] | None = None existing: Callable[[Table, FromClause, FromClause], Select] | None = None diff_keys: list[str] = field(default_factory=list)
[docs] def construct_ledger_balances_query( *keys: str, ) -> Callable[[Table, FromClause, FromClause], Select]: """Return an ``existing`` callable for common balance lookup. Produces a select that negates the current balances for each diff-key group present in the desired CTE:: SELECT key1, key2, SUM(value) * -1 AS value FROM ledger_table WHERE (key1, key2) IN (SELECT key1, key2 FROM desired) GROUP BY key1, key2 Use this when every key the event might want to unwind is also present in ``desired``. If a status transition or cancellation can *retire* keys -- so they disappear from ``desired`` -- use :func:`construct_ledger_scoped_balances_query` instead, which slices by an owner predicate rather than by the diff-key tuple. Args: *keys: Dimension column names to group by. Returns: A callable suitable for ``LedgerEvent(existing=...)``. """ def _existing( root: Table, desired: FromClause, pginput: FromClause, ) -> Select: _ = pginput key_cols = [root.c[k] for k in keys] desired_key_cols = [desired.c[k] for k in keys] return ( select( *key_cols, (func.sum(root.c.value) * -1).label("value"), ) .where(tuple_(*key_cols).in_(select(*desired_key_cols))) .group_by(*key_cols) ) return _existing
[docs] def construct_ledger_scoped_balances_query( *keys: str, where: Callable[[Table, FromClause], ColumnElement[bool]], ) -> Callable[[Table, FromClause, FromClause], Select]: """Return an ``existing`` callable scoped by an owner predicate. Unlike :func:`construct_ledger_balances_query`, which filters existing rows to diff-key tuples that appear in ``desired``, this filters by an arbitrary predicate over the ledger root table:: SELECT key1, key2, SUM(value) * -1 AS value FROM ledger_table WHERE <where(root, pginput)> GROUP BY key1, key2 The scoped form is the right tool whenever an event reconciles every ledger row belonging to some owner entity (a purchase order, an invoice, a work order) and a state transition might *retire* a diff key -- i.e. a row that was posted previously should no longer be posted after the event runs. The default ``construct_ledger_balances_query`` would miss such retired keys because they no longer appear in ``desired``; the scoped form sees them because it slices by the owner, so the diff naturally unwinds them. Args: *keys: Dimension column names to group by (the diff keys). where: Callable ``(ledger_root, pginput) -> predicate`` that returns a boolean restricting rows to the owner scope. ``pginput`` is a table reference to the generated ``input`` CTE, so function-param values are reached as ``pginput.c.<name>`` -- the same accessor ``desired`` uses. Returns: A callable suitable for ``LedgerEvent(existing=...)``. Example: When the owner id is denormalized on the ledger:: def owned_by_input_po(ledger, pginput): return ( ledger.c.purchase_order_id == pginput.c.purchase_order_id ) existing=construct_ledger_scoped_balances_query( "lot_id", "class_id", "account", "direction", where=owned_by_input_po, ) When the owner is reached via a join, use a subquery:: def owned_by_input_po(ledger, pginput): return ledger.c.lot_id.in_( select(PurchaseLot.id) .join(LineItem, LineItem.id == PurchaseLot.line_item_id) .where( LineItem.purchase_order_id == pginput.c.purchase_order_id, ), ) existing=construct_ledger_scoped_balances_query( "lot_id", "class_id", "account", "direction", where=owned_by_input_po, ) """ def _existing( root: Table, desired: FromClause, pginput: FromClause, ) -> Select: _ = desired key_cols = [root.c[k] for k in keys] return ( select( *key_cols, (func.sum(root.c.value) * -1).label("value"), ) .where(where(root, pginput)) .group_by(*key_cols) ) return _existing
def _input_table_ref( input_select: Select, ) -> FromClause: """Build a synthetic ``table("input", ...)`` from the input select. Args: input_select: The compiled input select. Returns: A :func:`~sqlalchemy.table` reference with matching columns. """ cols: list[ColumnClause] = [ column(col.name) for col in input_select.selected_columns ] return table("input", *cols) def _desired_table_ref( desired_select: SelectBase, ) -> FromClause: """Build a synthetic ``table("desired", ...)`` from desired select. Args: desired_select: The compiled desired select. Returns: A :func:`~sqlalchemy.table` reference with matching columns. """ cols: list[ColumnClause] = [ column(col.name) for col in desired_select.selected_columns ] return table("desired", *cols)