Source code for codegen_database.ext.ledger.chart_functions

"""Chart-focused function builders for ledger factories.

These build :class:`~codegen_database.functions.CodegenDatabaseFunctionSpec`
instances that can be unpacked into a declarative
:class:`~codegen_database.declarative.CodegenDatabaseFunctionMixin`
subclass (via ``__funcspec__``) or registered imperatively.

The generated functions are ``LANGUAGE sql STABLE``: Postgres inlines
them at plan time, so predicates added by callers (``WHERE``,
``LIMIT``) push through the parameters.  PL/pgSQL is avoided on
purpose — dynamic grouping in PL/pgSQL means ``EXECUTE format(...)``
which blocks the planner from inlining.

Function bodies are assembled with SQLAlchemy core and compiled via
:func:`~codegen_database.utils.query.compile_query` so identifiers, casts, and
window frames go through the normal dialect compiler.  Function
parameters (``p_period``, ``p_start``, ``p_end``) surface as
``literal_column`` references so they render as bare identifiers in
the generated SELECT.

Bucketing uses the :func:`codegen_database_date_bin
<codegen_database.ext.chart.date_bin.construct_date_bin_function>` helper
rather than ``date_trunc`` so callers can pass arbitrary
``interval`` strides (``'15 minutes'``, ``'3 months'``, ``'1 year'``,
etc.).  The helper must be installed in the same schema as the chart
function (or in ``date_bin_schema``, if overridden).
"""

from __future__ import annotations

import enum
from functools import wraps
from typing import TYPE_CHECKING

from sqlalchemy import (
    DateTime,
    Interval,
    Numeric,
    Table,
    Text,
    and_,
    cast,
    column,
    func,
    literal_column,
    or_,
    select,
    text,
)
from sqlalchemy_declarative_extensions.dialects.postgresql import (
    FunctionParam,
    FunctionSecurity,
    FunctionVolatility,
)

from codegen_database.errors import CodegenDatabaseValidationError
from codegen_database.ext.chart._contract import (
    CHART_SCHEMA_KEY,
    assert_chart_extension_declared,
)
from codegen_database.ext.chart.date_bin import CODEGEN_DATABASE_DATE_BIN_NAME
from codegen_database.ext.ledger.functions import _pg_type_str
from codegen_database.functions import CodegenDatabaseFunctionSpec
from codegen_database.utils.query import compile_query

if TYPE_CHECKING:
    from collections.abc import Callable

    from sqlalchemy.sql import ColumnElement
    from sqlalchemy.sql.elements import ColumnClause

    from codegen_database.factory.context import ContextSource


def _requires_chart_extension[**P](
    fn: Callable[P, CodegenDatabaseFunctionSpec],
) -> Callable[P, CodegenDatabaseFunctionSpec]:
    """Guard a chart function builder with a ChartExtension check.

    The decorated builder takes a :class:`ContextSource` as its first
    positional argument; the check reads ``source.ctx.metadata`` and
    raises :class:`~codegen_database.errors.CodegenDatabaseValidationError` if
    :class:`~codegen_database.ext.chart.ChartExtension` is not
    registered on the metadata.
    """

    @wraps(fn)
    def wrapper(
        *args: P.args, **kwargs: P.kwargs
    ) -> CodegenDatabaseFunctionSpec:
        source: ContextSource = args[0]  # type: ignore[assignment]
        assert_chart_extension_declared(source.ctx.metadata)
        return fn(*args, **kwargs)

    return wrapper


[docs] class NormalSide(enum.Enum): """Which side an account's balance naturally grows on. Double-entry accounting splits accounts into two camps: - ``DEBIT``: assets and expenses — the balance grows with debit postings (e.g. cash receipts, inventory purchases). - ``CREDIT``: liabilities, equity, and revenue — the balance grows with credit postings (e.g. sales, capital injections). :func:`construct_double_entry_chart_function` reads a column carrying these string values from the accounts dimension and uses them to normalize ``delta = debits - credits`` (for debit-normal accounts) or ``delta = credits - debits`` (for credit-normal accounts), so ``delta`` is always positive when the balance moves in the account's natural direction. Pair with :class:`~codegen_database.types.TextEnum` to persist the string value:: from codegen_database import NormalSide, TextEnum class Accounts(Base): name = Column(String, unique=True, nullable=False) normal_side = Column(TextEnum(NormalSide), nullable=False) """ DEBIT = "debit" CREDIT = "credit"
def _resolve_table(t: object) -> Table: """Accept either a SQLAlchemy declarative class or a ``Table``.""" if isinstance(t, Table): return t inner = getattr(t, "__table__", None) if isinstance(inner, Table): return inner msg = ( f"expected a SQLAlchemy Table or declarative class, got " f"{type(t).__name__}" ) raise CodegenDatabaseValidationError(msg) def _dim_type_strs(source: ContextSource, dimensions: list[str]) -> list[str]: """Return PG type strings for each dimension column.""" table = source.ctx["primary"] col_names = {c.key for c in table.columns} types: list[str] = [] for d in dimensions: if d not in col_names: msg = ( f"dimension {d!r} is not a column on the ledger " f"table. Available columns: {sorted(col_names)}" ) raise CodegenDatabaseValidationError(msg) types.append(_pg_type_str(table.c[d])) return types def _validate_split_by( source: ContextSource, split_by: tuple[str, list[str]] | None, ) -> None: """Validate split_by argument shape and column existence.""" if split_by is None: return split_col, split_values = split_by if not split_values: msg = "split_by values list must be non-empty" raise CodegenDatabaseValidationError(msg) table = source.ctx["raw_table"] col_names = {c.key for c in table.columns} if split_col not in col_names: msg = ( f"split_by column {split_col!r} is not a column on the " f"ledger table. Available columns: {sorted(col_names)}" ) raise CodegenDatabaseValidationError(msg) def _split_filter_columns( value_col: ColumnElement, split_by: tuple[str, list[str]] | None, ) -> list[ColumnElement]: """Build COALESCE-wrapped SUM FILTER columns cast to numeric.""" if split_by is None: return [] split_col, split_values = split_by split_col_ref: ColumnClause = column(split_col) return [ cast( func.coalesce( func.sum(value_col).filter(split_col_ref == v), 0, ), Numeric, ).label(v) for v in split_values ] def _split_returns_cols( split_by: tuple[str, list[str]] | None, ) -> list[str]: """Build TABLE-column declarations for split columns.""" if split_by is None: return [] _, split_values = split_by return [f"{v} numeric" for v in split_values] def _range_params(period_default: str) -> list[FunctionParam]: """Build the (p_period, p_start, p_end) parameter list. ``p_period`` is an ``interval`` (``'1 day'``, ``'15 minutes'``, ``'3 months'``, …). ``p_start`` and ``p_end`` bracket a half-open date range ``[p_start, p_end)`` on ``created_at``. Both default to ``NULL``, which the body treats as "unbounded on that side". """ return [ FunctionParam.input( "p_period", "interval", default=f"'{period_default}'::interval", ), FunctionParam.input("p_start", "timestamptz", default="NULL"), FunctionParam.input("p_end", "timestamptz", default="NULL"), ] def _validate_period_default(period_default: str) -> None: """Reject obviously-bad interval literals early. Full interval validation needs a live Postgres connection, so this only catches empty or whitespace-only strings. Postgres will reject malformed interval literals at function creation time with a clearer error than anything we could raise here. """ if not period_default or not period_default.strip(): msg = "period_default must be a non-empty interval literal" raise CodegenDatabaseValidationError(msg) def _date_bin_call( schema: str, p_period: ColumnElement, created_at: ColumnElement, *, name: str, ) -> ColumnElement: """Build a ``<schema>.<name>(p_period, created_at)`` function call. Uses the documented SQLAlchemy ``func.<schema>.<name>`` form so the identifier is schema-qualified at compile time. """ return getattr(getattr(func, schema), name)(p_period, created_at) def _period_range_expr( schema: str, date_bin_name: str, created_at_col: str ) -> ColumnElement: """Build the ``tstzrange(bucket_start, bucket_end, '[)')`` SELECT expr. ``bucket_start`` is ``codegen_database_date_bin(p_period, created_at)`` and ``bucket_end`` is ``bucket_start + p_period``. Emitted as raw text because SQLAlchemy cannot infer that the function call returns ``timestamptz`` and therefore refuses to compile ``... + interval``. GROUP BY / ORDER BY still use the scalar bucket-start (via ``period_expr``) so the planner is not asked to group on a range. Mixed-stride intervals (e.g. ``'1 month 3 days'``) are rejected by ``codegen_database_date_bin`` — it returns ``NULL``. Naively wrapping that in ``tstzrange`` would yield the unbounded range ``(,)`` and silently mask the error, so we guard with an explicit ``CASE`` that preserves the ``NULL`` signal. """ bucket_start = f"{schema}.{date_bin_name}(p_period, {created_at_col})" return literal_column( "CASE WHEN " f"{bucket_start} IS NULL THEN NULL::tstzrange ELSE " f"tstzrange({bucket_start}, {bucket_start} + p_period, '[)') " "END" ) def _period_params() -> tuple[ColumnClause, ColumnClause, ColumnClause]: """Return the ``(p_period, p_start, p_end)`` function-parameter columns. These typed ``literal_column`` references are identical across every chart function; centralised so the SQL types stay in lockstep. """ return ( literal_column("p_period", type_=Interval()), literal_column("p_start", type_=DateTime(timezone=True)), literal_column("p_end", type_=DateTime(timezone=True)), ) def _range_filter( created_at: ColumnElement, p_start: ColumnElement, p_end: ColumnElement, ) -> ColumnElement: """Build the (p_start IS NULL OR ...) AND (p_end ...) predicate.""" return and_( or_(p_start.is_(None), created_at >= p_start), or_(p_end.is_(None), created_at < p_end), )
[docs] @_requires_chart_extension def construct_ledger_chart_function( # noqa: PLR0913 source: ContextSource, *, name: str, dimensions: list[str], period_default: str = "1 day", split_by: tuple[str, list[str]] | None = None, date_bin_schema: str | None = None, date_bin_name: str = CODEGEN_DATABASE_DATE_BIN_NAME, ) -> CodegenDatabaseFunctionSpec: """Build a period x dimensions chart function spec. This is a *range* function: callers pass ``p_start`` / ``p_end`` to bracket a half-open ``[p_start, p_end)`` window on ``created_at``. Both default to ``NULL`` meaning "unbounded on that side". The function returns one row per ``(codegen_database_date_bin(p_period, created_at), dimensions...)`` bucket inside the range. Signature:: <name>(p_period interval DEFAULT '<period_default>'::interval, p_start timestamptz DEFAULT NULL, p_end timestamptz DEFAULT NULL) RETURNS TABLE ( period tstzrange, <dim1> <type>, ..., starting_balance numeric, delta numeric, [<split_value_1> numeric, ...,] ending_balance numeric ) The ``period`` column is a half-open ``[bucket_start, bucket_end)`` ``tstzrange`` so callers see the full extent of each bucket, not just its start. ``bucket_start`` is the polyfilled ``codegen_database_date_bin(p_period, created_at)`` and ``bucket_end`` is ``bucket_start + p_period``. To filter:: -- All buckets contained in Feb 2024: WHERE period <@ tstzrange('2024-02-01', '2024-03-01', '[)') -- Exactly the Feb 2024 bucket (only useful when stride and -- boundary align): WHERE lower(period) = '2024-02-01' Body (compiled from a SQLAlchemy select):: SELECT CASE WHEN <schema>.codegen_database_date_bin(p_period, created_at) IS NULL THEN NULL::tstzrange ELSE tstzrange( <schema>.codegen_database_date_bin(p_period, created_at), <schema>.codegen_database_date_bin(p_period, created_at) + p_period, '[)') END AS period, <dims>, COALESCE(SUM(SUM(value)) OVER ( PARTITION BY <dims> ORDER BY <schema>.codegen_database_date_bin( p_period, created_at) ROWS BETWEEN UNBOUNDED PRECEDING AND 1 PRECEDING ), 0)::numeric AS starting_balance, SUM(value)::numeric AS delta, [COALESCE(SUM(value) FILTER ( WHERE <split_col> = '<v>'), 0)::numeric AS <v>, ...,] SUM(SUM(value)) OVER ( PARTITION BY <dims> ORDER BY <schema>.codegen_database_date_bin( p_period, created_at) )::numeric AS ending_balance FROM <schema>.<raw_table> WHERE (p_start IS NULL OR created_at >= p_start) AND (p_end IS NULL OR created_at < p_end) GROUP BY <schema>.codegen_database_date_bin( p_period, created_at), <dims> ORDER BY <schema>.codegen_database_date_bin( p_period, created_at), <dims> ``GROUP BY`` / ``ORDER BY`` use the scalar bucket-start rather than the range expression so the planner operates on ``timestamptz``, not ``tstzrange``, and window ``PARTITION BY`` stays cheap. The ``codegen_database_date_bin`` helper (see :func:`~codegen_database.ext.chart.date_bin.construct_date_bin_function`) must exist in ``date_bin_schema``. It polyfills Postgres's ``date_bin`` for month/quarter/year strides, which the native function rejects. Mixed-stride intervals (e.g. ``'1 month 3 days'``) return ``NULL`` for ``period`` — the ``CASE`` in the body preserves that signal rather than emitting the unbounded range ``(,)``. Note: ``starting_balance`` and ``ending_balance`` are cumulative flows *within the filtered window*, not absolute ledger balances. To anchor against history, widen ``p_start`` or compose with :func:`~codegen_database.ext.ledger.queries.construct_ledger_balance_query`. Args: source: Ledger source. name: Unqualified function name. dimensions: Column names to group by. Must be non-empty. period_default: Default value for ``p_period`` as a Postgres interval literal, e.g. ``"1 day"``, ``"15 minutes"``, ``"3 months"``, ``"1 year"``. split_by: Optional ``(column_name, values)`` pair. Each value becomes a named ``SUM`` column filtered on equality, e.g. ``split_by=("direction", ["debit", "credit"])`` on a double-entry ledger. date_bin_schema: Schema where ``codegen_database_date_bin`` lives. Defaults to the schema configured by :class:`~codegen_database.ext.chart.ChartExtension` (``"codegen_database"`` unless overridden). date_bin_name: Name of the polyfill function. Defaults to ``codegen_database_date_bin``. Returns: A :class:`~codegen_database.functions.CodegenDatabaseFunctionSpec`. Raises: CodegenDatabaseValidationError: If *dimensions* is empty, any dimension is unknown, *period_default* is empty, or *split_by* names a missing column / is empty. """ if not dimensions: msg = "dimensions must be a non-empty list" raise CodegenDatabaseValidationError(msg) _validate_period_default(period_default) _validate_split_by(source, split_by) ctx = source.ctx raw_table = ctx["raw_table"] created_at_col: str = ctx["created_at_column"] schema: str = date_bin_schema or ctx.metadata.info[CHART_SCHEMA_KEY] dim_types = _dim_type_strs(source, dimensions) dim_cols: list[ColumnClause] = [column(d) for d in dimensions] created_at: ColumnClause = column(created_at_col) value: ColumnClause = column("value") p_period, p_start, p_end = _period_params() period_expr = _date_bin_call( schema, p_period, created_at, name=date_bin_name ) period_range = _period_range_expr( schema, date_bin_name, created_at_col ).label("period") delta = func.sum(value) starting_balance = cast( func.coalesce( func.sum(delta).over( partition_by=dim_cols, order_by=period_expr, rows=(None, -1), ), 0, ), Numeric, ).label("starting_balance") delta_col = cast(delta, Numeric).label("delta") ending_balance = cast( func.sum(delta).over( partition_by=dim_cols, order_by=period_expr, ), Numeric, ).label("ending_balance") stmt = ( select( period_range, *dim_cols, starting_balance, delta_col, *_split_filter_columns(value, split_by), ending_balance, ) .select_from(raw_table) .where(_range_filter(created_at, p_start, p_end)) .group_by(period_expr, *dim_cols) .order_by(period_expr, *dim_cols) ) returns_cols = [ "period tstzrange", *[f"{d} {t}" for d, t in zip(dimensions, dim_types, strict=True)], "starting_balance numeric", "delta numeric", *_split_returns_cols(split_by), "ending_balance numeric", ] returns = "TABLE (" + ", ".join(returns_cols) + ")" return CodegenDatabaseFunctionSpec( name=name, definition=compile_query(stmt), language="sql", parameters=_range_params(period_default), returns=returns, security=FunctionSecurity.invoker, volatility=FunctionVolatility.STABLE, )
[docs] @_requires_chart_extension def construct_ledger_rollup_chart_function( # noqa: PLR0913 source: ContextSource, *, name: str, dimensions: list[str], period_default: str = "1 day", date_bin_schema: str | None = None, date_bin_name: str = CODEGEN_DATABASE_DATE_BIN_NAME, ) -> CodegenDatabaseFunctionSpec: """Build a period x ROLLUP(dimensions) chart function spec. This is a *range* function: callers pass ``p_start`` / ``p_end`` to bracket a half-open ``[p_start, p_end)`` window on ``created_at``. Both default to ``NULL`` meaning "unbounded on that side". Signature:: <name>(p_period interval DEFAULT '<period_default>'::interval, p_start timestamptz DEFAULT NULL, p_end timestamptz DEFAULT NULL) RETURNS TABLE ( period tstzrange, <dim1> <type>, ..., delta numeric ) The ``period`` column is the same ``[bucket_start, bucket_end)`` ``tstzrange`` as :func:`construct_ledger_chart_function` — see that function for the rationale and for filtering examples. Unlike the regular chart function, ``starting_balance`` and ``ending_balance`` are not emitted: running totals are undefined against grouping sets. The body uses ``ROLLUP`` to produce hierarchical subtotals in a single query. Rolled-up (subtotal) rows carry ``NULL`` in the rolled-up dimension columns; callers distinguish them from detail rows with ``<dim> IS NULL`` / ``<dim> IS NOT NULL``. This assumes the dimension itself is never ``NULL`` in the underlying data — which is the norm for FK / enum dimensions; if your data has genuine ``NULL`` values in a dimension, subtotal rows will be indistinguishable from NULL-data rows. :: SELECT CASE WHEN <schema>.codegen_database_date_bin(p_period, created_at) IS NULL THEN NULL::tstzrange ELSE tstzrange( <schema>.codegen_database_date_bin(p_period, created_at), <schema>.codegen_database_date_bin(p_period, created_at) + p_period, '[)') END AS period, <dims>, SUM(value)::numeric AS delta FROM <schema>.<raw_table> WHERE (p_start IS NULL OR created_at >= p_start) AND (p_end IS NULL OR created_at < p_end) GROUP BY <schema>.codegen_database_date_bin(p_period, created_at), ROLLUP(<dims>) ORDER BY <schema>.codegen_database_date_bin(p_period, created_at), <dims> NULLS LAST Args: source: Ledger source. name: Unqualified function name. dimensions: Column names forming the rollup hierarchy (left to right = outer to inner). Must be non-empty. period_default: Default value for ``p_period`` as a Postgres interval literal, e.g. ``"1 day"``, ``"3 months"``. date_bin_schema: Schema where ``codegen_database_date_bin`` lives. Defaults to the schema configured by :class:`~codegen_database.ext.chart.ChartExtension` (``"codegen_database"`` unless overridden). date_bin_name: Name of the polyfill function. Defaults to ``codegen_database_date_bin``. Returns: A :class:`~codegen_database.functions.CodegenDatabaseFunctionSpec`. Raises: CodegenDatabaseValidationError: If *dimensions* is empty, any dimension is unknown, or *period_default* is empty. """ if not dimensions: msg = "dimensions must be a non-empty list" raise CodegenDatabaseValidationError(msg) _validate_period_default(period_default) ctx = source.ctx raw_table = ctx["raw_table"] created_at_col: str = ctx["created_at_column"] schema: str = date_bin_schema or ctx.metadata.info[CHART_SCHEMA_KEY] dim_types = _dim_type_strs(source, dimensions) dim_cols: list[ColumnClause] = [column(d) for d in dimensions] created_at: ColumnClause = column(created_at_col) value: ColumnClause = column("value") p_period, p_start, p_end = _period_params() period_expr = _date_bin_call( schema, p_period, created_at, name=date_bin_name ) period_range = _period_range_expr( schema, date_bin_name, created_at_col ).label("period") # ROLLUP is a grouping-set keyword in Postgres; SQLAlchemy has no # first-class support for it, so emit the GROUP BY tail as text. dim_list = ", ".join(dimensions) rollup_clause = text(f"ROLLUP({dim_list})") stmt = ( select( period_range, *dim_cols, cast(func.sum(value), Numeric).label("delta"), ) .select_from(raw_table) .where(_range_filter(created_at, p_start, p_end)) .group_by(period_expr, rollup_clause) .order_by(period_expr, *[c.nullslast() for c in dim_cols]) ) returns_cols = [ "period tstzrange", *[f"{d} {t}" for d, t in zip(dimensions, dim_types, strict=True)], "delta numeric", ] returns = "TABLE (" + ", ".join(returns_cols) + ")" return CodegenDatabaseFunctionSpec( name=name, definition=compile_query(stmt), language="sql", parameters=_range_params(period_default), returns=returns, security=FunctionSecurity.invoker, volatility=FunctionVolatility.STABLE, )
[docs] @_requires_chart_extension def construct_double_entry_chart_function( # noqa: PLR0913 source: ContextSource, *, name: str, accounts_table: object, dimensions: list[str], account_column: str = "account", direction_column: str = "direction", accounts_key_column: str = "name", normal_side_column: str = "normal_side", period_default: str = "1 day", date_bin_schema: str | None = None, date_bin_name: str = CODEGEN_DATABASE_DATE_BIN_NAME, ) -> CodegenDatabaseFunctionSpec: """Build a double-entry period x dimensions chart function spec. This is a specialization of :func:`construct_ledger_chart_function` for double-entry ledgers: it joins each raw row to an *accounts* dimension table, splits raw ``value`` into ``debits`` and ``credits`` by direction, and normalizes the per-bucket change (``delta``) by each account's ``normal_side``. The normalization rule is the standard accounting convention: * Debit-normal account (assets, expenses): ``delta = debits - credits``. * Credit-normal account (liabilities, equity, revenue): ``delta = credits - debits``. Running balances (``starting_balance``, ``ending_balance``) are cumulative sums of the normalized ``delta`` inside the ``[p_start, p_end)`` window, partitioned by ``dimensions``. Signature:: <name>(p_period interval DEFAULT '<period_default>'::interval, p_start timestamptz DEFAULT NULL, p_end timestamptz DEFAULT NULL) RETURNS TABLE ( period tstzrange, <dim1> <type>, ..., starting_balance numeric, debits numeric, credits numeric, delta numeric, ending_balance numeric ) The ``period`` column is the same ``[bucket_start, bucket_end)`` ``tstzrange`` as :func:`construct_ledger_chart_function` — see that function for filtering examples. Body (CTE-based for readability; the CTE holds per-bucket aggregates, the outer SELECT adds running-balance windows):: WITH per_bucket AS ( SELECT <bucket_start> AS bucket_start, <dims>, a.<normal_side> AS _normal_side, COALESCE(SUM(l.value) FILTER ( WHERE l.<direction> = 'debit'), 0) AS debits, COALESCE(SUM(l.value) FILTER ( WHERE l.<direction> = 'credit'), 0) AS credits FROM <schema>.<raw_table> l INNER JOIN <a_schema>.<accounts> a ON a.<name> = l.<account> WHERE (p_start IS NULL OR l.created_at >= p_start) AND (p_end IS NULL OR l.created_at < p_end) GROUP BY bucket_start, <dims>, a.<normal_side> ) SELECT <period_range> AS period, <dims>, COALESCE(SUM(<delta>) OVER ( PARTITION BY <dims> ORDER BY bucket_start ROWS BETWEEN UNBOUNDED PRECEDING AND 1 PRECEDING ), 0)::numeric AS starting_balance, debits::numeric AS debits, credits::numeric AS credits, (<delta>)::numeric AS delta, (SUM(<delta>) OVER ( PARTITION BY <dims> ORDER BY bucket_start ))::numeric AS ending_balance FROM per_bucket ORDER BY bucket_start, <dims> where ``<delta>`` is:: CASE WHEN _normal_side = 'debit' THEN debits - credits ELSE credits - debits END Note: ``starting_balance`` and ``ending_balance`` are cumulative *within the filtered window*, not absolute historical balances. Widen ``p_start`` to anchor against history. Args: source: Ledger source. name: Unqualified function name. accounts_table: A SQLAlchemy ``Table`` or declarative class for the accounts dimension. Must have ``name`` and ``normal_side`` columns (or override via *accounts_key_column* / *normal_side_column*). dimensions: Ledger columns to group by (typically includes ``account``). Must be non-empty. account_column: Name of the account column on the raw ledger table. Defaults to ``"account"``. direction_column: Name of the direction column on the raw ledger table. Defaults to ``"direction"``. accounts_key_column: Column on *accounts_table* that matches ``account_column`` values. Defaults to ``"name"``. normal_side_column: Column on *accounts_table* holding ``'debit'`` or ``'credit'``. Defaults to ``"normal_side"``. period_default: Default for ``p_period`` as an interval literal, e.g. ``"1 day"``, ``"1 month"``. date_bin_schema: Schema where ``codegen_database_date_bin`` lives. Defaults to the schema configured by :class:`~codegen_database.ext.chart.ChartExtension` (``"codegen_database"`` unless overridden). date_bin_name: Name of the polyfill function. Defaults to ``codegen_database_date_bin``. Returns: A :class:`~codegen_database.functions.CodegenDatabaseFunctionSpec`. Raises: CodegenDatabaseValidationError: If *dimensions* is empty, any dimension / account_column / direction_column is unknown, or *period_default* is empty. """ if not dimensions: msg = "dimensions must be a non-empty list" raise CodegenDatabaseValidationError(msg) _validate_period_default(period_default) ctx = source.ctx raw_table = ctx["raw_table"] created_at_col: str = ctx["created_at_column"] schema: str = date_bin_schema or ctx.metadata.info[CHART_SCHEMA_KEY] dim_types = _dim_type_strs(source, dimensions) raw_col_names = {c.key for c in raw_table.columns} for c in (account_column, direction_column): if c not in raw_col_names: msg = ( f"column {c!r} is not a column on the ledger table. " f"Available: {sorted(raw_col_names)}" ) raise CodegenDatabaseValidationError(msg) accounts = _resolve_table(accounts_table) a_schema = accounts.schema or ctx.schemaname a_name = accounts.name fn_schema: str = ctx.schemaname raw_name = raw_table.name bucket_expr = f"{schema}.{date_bin_name}(p_period, l.{created_at_col})" debit = NormalSide.DEBIT.value credit = NormalSide.CREDIT.value debits_sum = ( "COALESCE(SUM(l.value) FILTER " f"(WHERE l.{direction_column} = '{debit}'), 0)" ) credits_sum = ( "COALESCE(SUM(l.value) FILTER " f"(WHERE l.{direction_column} = '{credit}'), 0)" ) delta_expr = ( f"CASE WHEN _normal_side = '{debit}' " "THEN debits - credits " "ELSE credits - debits END" ) dim_csv = ", ".join(dimensions) dim_l_csv = ", ".join(f"l.{d}" for d in dimensions) body = "\n".join( [ "WITH per_bucket AS (", " SELECT", f" {bucket_expr} AS bucket_start,", f" {dim_l_csv},", f" a.{normal_side_column} AS _normal_side,", f" {debits_sum} AS debits,", f" {credits_sum} AS credits", f" FROM {fn_schema}.{raw_name} AS l", f" INNER JOIN {a_schema}.{a_name} AS a", f" ON a.{accounts_key_column} = l.{account_column}", f" WHERE (p_start IS NULL OR l.{created_at_col} >= p_start)", f" AND (p_end IS NULL OR l.{created_at_col} < p_end)", f" GROUP BY 1, {dim_l_csv}, a.{normal_side_column}", ")", "SELECT", " CASE WHEN bucket_start IS NULL THEN NULL::tstzrange", " ELSE tstzrange(bucket_start, " "bucket_start + p_period, '[)')", " END AS period,", f" {dim_csv},", f" COALESCE(SUM({delta_expr}) OVER (", f" PARTITION BY {dim_csv}", " ORDER BY bucket_start", " ROWS BETWEEN UNBOUNDED PRECEDING AND 1 PRECEDING", " ), 0)::numeric AS starting_balance,", " debits::numeric AS debits,", " credits::numeric AS credits,", f" ({delta_expr})::numeric AS delta,", f" (SUM({delta_expr}) OVER (", f" PARTITION BY {dim_csv}", " ORDER BY bucket_start", " ))::numeric AS ending_balance", "FROM per_bucket", f"ORDER BY bucket_start, {dim_csv}", ] ) returns_cols = [ "period tstzrange", *[f"{d} {t}" for d, t in zip(dimensions, dim_types, strict=True)], "starting_balance numeric", "debits numeric", "credits numeric", "delta numeric", "ending_balance numeric", ] returns = "TABLE (" + ", ".join(returns_cols) + ")" return CodegenDatabaseFunctionSpec( name=name, definition=body, language="sql", parameters=_range_params(period_default), returns=returns, security=FunctionSecurity.invoker, volatility=FunctionVolatility.STABLE, )
[docs] @_requires_chart_extension def construct_ledger_pivoted_chart_function( # noqa: PLR0913 source: ContextSource, *, name: str, dimensions: list[str], period_default: str = "1 day", date_bin_schema: str | None = None, date_bin_name: str = CODEGEN_DATABASE_DATE_BIN_NAME, ) -> CodegenDatabaseFunctionSpec: """Build a JSONB-pivoted chart function spec. Returns one row per ``dimensions`` group with a single ``buckets`` ``jsonb`` column mapping UTC-normalized bucket-start keys to ``numeric`` deltas. Keeping the shape fixed (unlike ``crosstab``, which forces callers to spell every pivot column in an ``AS t(...)`` clause at every call site) makes this function usable anywhere ``SELECT * FROM`` goes: ORMs, CTEs, views. Signature:: <name>(p_period interval DEFAULT '<period_default>'::interval, p_start timestamptz DEFAULT NULL, p_end timestamptz DEFAULT NULL) RETURNS TABLE ( <dim1> <type>, ..., buckets jsonb ) Typical usage:: SELECT warehouse, sku, (buckets->>'2024-01-01 00:00:00')::numeric AS day1, (buckets->>'2024-01-02 00:00:00')::numeric AS day2 FROM <schema>.<name>( p_period => '1 day'::interval, p_start => '2024-01-01'::timestamptz, p_end => '2024-01-03'::timestamptz ); Bucket keys are the UTC-normalized text form of the bucket-start timestamp — ``to_char(bucket_start AT TIME ZONE 'UTC', 'YYYY-MM-DD HH24:MI:SS')`` — chosen over ``bucket_start::text`` so the keys are stable across session timezones. Missing buckets are absent from the JSONB (not present with value ``0``), so use ``COALESCE((buckets->>'...')::numeric, 0)`` when you need zeros. Body (compiled from a SQLAlchemy select):: SELECT <dims>, jsonb_object_agg( to_char(bucket_start AT TIME ZONE 'UTC', 'YYYY-MM-DD HH24:MI:SS'), delta ) AS buckets FROM ( SELECT <dims>, <schema>.codegen_database_date_bin(p_period, created_at) AS bucket_start, SUM(value)::numeric AS delta FROM <fn_schema>.<raw_table> WHERE (p_start IS NULL OR created_at >= p_start) AND (p_end IS NULL OR created_at < p_end) GROUP BY <dims>, bucket_start ) AS b GROUP BY <dims> ORDER BY <dims> Note: ``jsonb_object_agg`` raises on duplicate keys. The inner aggregate's ``GROUP BY <dims>, bucket_start`` guarantees one row per key within each dimension group, so duplicates only appear if ``codegen_database_date_bin`` returns ``NULL`` (mixed-stride intervals like ``'1 month 3 days'``) for multiple source rows in the same dimension group — those keys collapse to ``null`` and Postgres rejects the aggregate. Avoid mixed-stride intervals; ``codegen_database_date_bin`` documents the supported forms. Args: source: Ledger source. name: Unqualified function name. dimensions: Column names used as row keys. Must be non-empty. period_default: Default for ``p_period`` as an interval literal, e.g. ``"1 day"``, ``"1 month"``. date_bin_schema: Schema where ``codegen_database_date_bin`` lives. Defaults to the schema configured by :class:`~codegen_database.ext.chart.ChartExtension`. date_bin_name: Name of the polyfill function. Defaults to ``codegen_database_date_bin``. Returns: A :class:`~codegen_database.functions.CodegenDatabaseFunctionSpec`. Raises: CodegenDatabaseValidationError: If *dimensions* is empty, any dimension is unknown, or *period_default* is empty. """ if not dimensions: msg = "dimensions must be a non-empty list" raise CodegenDatabaseValidationError(msg) _validate_period_default(period_default) dim_types = _dim_type_strs(source, dimensions) ctx = source.ctx raw_table = ctx["raw_table"] created_at_col: str = ctx["created_at_column"] schema: str = date_bin_schema or ctx.metadata.info[CHART_SCHEMA_KEY] dim_cols: list[ColumnClause] = [column(d) for d in dimensions] created_at: ColumnClause = column(created_at_col) value: ColumnClause = column("value") p_period, p_start, p_end = _period_params() period_expr = _date_bin_call( schema, p_period, created_at, name=date_bin_name ) inner = ( select( *dim_cols, period_expr.label("bucket_start"), cast(func.sum(value), Numeric).label("delta"), ) .select_from(raw_table) .where(_range_filter(created_at, p_start, p_end)) .group_by(*dim_cols, period_expr) .subquery("b") ) bucket_key = func.to_char( func.timezone("UTC", inner.c.bucket_start), "YYYY-MM-DD HH24:MI:SS", type_=Text(), ) outer_dims = [inner.c[d] for d in dimensions] stmt = ( select( *outer_dims, func.jsonb_object_agg(bucket_key, inner.c.delta).label("buckets"), ) .group_by(*outer_dims) .order_by(*outer_dims) ) returns_cols = [ *[f"{d} {t}" for d, t in zip(dimensions, dim_types, strict=True)], "buckets jsonb", ] returns = "TABLE (" + ", ".join(returns_cols) + ")" return CodegenDatabaseFunctionSpec( name=name, definition=compile_query(stmt), language="sql", parameters=_range_params(period_default), returns=returns, security=FunctionSecurity.invoker, volatility=FunctionVolatility.STABLE, )