"""Rolling and period aggregation query builders for ledger factories.
Provides pure-Python helpers that return SQLAlchemy selects without
registering anything on metadata. Pass results to
:class:`~codegen_database.views.view.CodegenDatabasePlainView` or
:class:`~codegen_database.views.view.CodegenDatabaseMaterializedView`.
"""
from __future__ import annotations
from typing import TYPE_CHECKING, Literal, get_args
from sqlalchemy import (
DateTime,
and_,
cast,
func,
literal,
select,
text,
)
if TYPE_CHECKING:
from sqlalchemy import Select
from sqlalchemy.sql import ColumnElement
from codegen_database.factory.context import ContextSource
from codegen_database.errors import CodegenDatabaseValidationError
Period = Literal[
"microsecond",
"millisecond",
"second",
"minute",
"hour",
"day",
"week",
"month",
"quarter",
"year",
"decade",
"century",
"millennium",
]
_VALID_PERIODS = frozenset(get_args(Period))
def _check_period(period: str) -> None:
"""Raise if *period* is not a recognized ``date_trunc`` field."""
if period not in _VALID_PERIODS:
msg = (
f"Unknown period {period!r}. "
f"Must be one of: {sorted(_VALID_PERIODS)}"
)
raise CodegenDatabaseValidationError(msg)
[docs]
def construct_ledger_running_balance_query(
source: ContextSource,
dimensions: list[str],
) -> Select:
"""Build a running balance query using a window function.
Generates::
SELECT dim_col1, ..., created_at, value,
SUM(value) OVER (
PARTITION BY dim_col1, ...
ORDER BY created_at
ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW
) AS running_balance
FROM <primary_table>
ORDER BY dim_col1, ..., created_at
Args:
source: A
:class:`~codegen_database.factory.ledger.CodegenDatabaseLedger`
instance or a
:class:`~codegen_database.declarative.CodegenDatabaseBase`
subclass using ``CodegenDatabaseLedger`` as its factory.
dimensions: Column names to partition by. Must be a
non-empty list.
Returns:
A SQLAlchemy :class:`~sqlalchemy.Select`.
Raises:
CodegenDatabaseValidationError: If *dimensions* is empty.
"""
if not dimensions:
msg = "dimensions must be a non-empty list"
raise CodegenDatabaseValidationError(msg)
table = source.ctx["primary"]
created_at_col = source.ctx["created_at_column"]
dim_cols = [table.c[d] for d in dimensions]
order_col = table.c[created_at_col]
running_balance = (
func.sum(table.c["value"])
.over(
partition_by=dim_cols,
order_by=order_col,
rows=(None, 0),
)
.label("running_balance")
)
return (
select(
*[c.label(c.key) for c in dim_cols],
order_col.label(created_at_col),
table.c["value"].label("value"),
running_balance,
)
.select_from(table)
.order_by(*dim_cols, order_col)
)
[docs]
def construct_ledger_period_rollup_query(
source: ContextSource,
dimensions: list[str],
period: Period = "day",
) -> Select:
"""Build a period rollup query using ``date_trunc``.
Generates::
SELECT date_trunc('day', created_at) AS period,
dim_col1, ...,
SUM(value) AS balance
FROM <primary_table>
GROUP BY 1, dim_col1, ...
ORDER BY 1, dim_col1, ...
Args:
source: A
:class:`~codegen_database.factory.ledger.CodegenDatabaseLedger`
instance or a
:class:`~codegen_database.declarative.CodegenDatabaseBase`
subclass using ``CodegenDatabaseLedger`` as its factory.
dimensions: Column names to group by. Must be a
non-empty list.
period: ``date_trunc`` precision string, e.g. ``"day"``,
``"week"``, ``"month"``, ``"year"``.
Defaults to ``"day"``.
Returns:
A SQLAlchemy :class:`~sqlalchemy.Select`.
Raises:
CodegenDatabaseValidationError: If *dimensions* is empty or *period*
is not a valid ``date_trunc`` precision.
"""
if not dimensions:
msg = "dimensions must be a non-empty list"
raise CodegenDatabaseValidationError(msg)
_check_period(period)
table = source.ctx["primary"]
created_at_col = source.ctx["created_at_column"]
dim_cols = [table.c[d] for d in dimensions]
period_expr = func.date_trunc(
text(f"'{period}'"),
table.c[created_at_col],
).label("period")
return (
select(
period_expr,
*[c.label(c.key) for c in dim_cols],
func.sum(table.c["value"]).label("balance"),
)
.select_from(table)
.group_by(text("1"), *dim_cols)
.order_by(text("1"), *dim_cols)
)
def _validate_split_by(
split_by: tuple[str, list[str]] | None,
table: object,
) -> 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)
col_names = {c.key for c in table.columns} # type: ignore[attr-defined]
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_sum_columns(
table: object,
split_by: tuple[str, list[str]] | None,
) -> list:
"""Build FILTER-based SUM columns for each split value."""
if split_by is None:
return []
split_col, split_values = split_by
return [
func.coalesce(
func.sum(table.c["value"]).filter( # type: ignore[attr-defined]
table.c[split_col] == v # type: ignore[attr-defined]
),
0,
).label(v)
for v in split_values
]
[docs]
def construct_ledger_period_running_balance_query(
source: ContextSource,
dimensions: list[str],
period: Period = "day",
split_by: tuple[str, list[str]] | None = None,
) -> Select:
"""Build a per-period running balance query.
Generates::
SELECT date_trunc('<period>', created_at) AS period,
dim_col1, ...,
COALESCE(SUM(SUM(value)) OVER (
PARTITION BY dim_col1, ...
ORDER BY date_trunc('<period>', created_at)
ROWS BETWEEN UNBOUNDED PRECEDING AND 1 PRECEDING
), 0) AS starting_balance,
SUM(value) AS delta,
SUM(SUM(value)) OVER (
PARTITION BY dim_col1, ...
ORDER BY date_trunc('<period>', created_at)
) AS ending_balance
FROM <primary_table>
GROUP BY 1, dim_col1, ...
ORDER BY 1, dim_col1, ...
Columns:
* ``starting_balance`` -- cumulative balance just before this
period (i.e. sum of all prior-period deltas). ``0`` for the
first period in each partition.
* ``delta`` -- net change within this period.
* ``ending_balance`` -- cumulative balance at the close of
this period (``starting_balance + delta``).
Passing *split_by* adds one ``SUM(value) FILTER (WHERE col = v)``
column per listed value, interleaved between ``delta`` and
``ending_balance``. Typical use: ``split_by=("direction",
["debit", "credit"])`` on a double-entry ledger.
Args:
source: A
:class:`~codegen_database.factory.ledger.CodegenDatabaseLedger`
instance or a
:class:`~codegen_database.declarative.CodegenDatabaseBase`
subclass using ``CodegenDatabaseLedger`` as its factory.
dimensions: Column names to partition by. Must be a
non-empty list.
period: ``date_trunc`` precision string. Defaults to
``"day"``.
split_by: Optional ``(column_name, values)`` pair. Each
value becomes a named ``SUM`` column filtered on equality.
Returns:
A SQLAlchemy :class:`~sqlalchemy.Select`.
Raises:
CodegenDatabaseValidationError: On empty *dimensions*, unknown
*period*, empty or unknown *split_by*.
"""
if not dimensions:
msg = "dimensions must be a non-empty list"
raise CodegenDatabaseValidationError(msg)
_check_period(period)
table = source.ctx["primary"]
created_at_col = source.ctx["created_at_column"]
_validate_split_by(split_by, table)
dim_cols = [table.c[d] for d in dimensions]
period_expr = func.date_trunc(
text(f"'{period}'"),
table.c[created_at_col],
)
delta = func.sum(table.c["value"])
ending_balance = func.sum(delta).over(
partition_by=dim_cols,
order_by=period_expr,
)
starting_balance = func.coalesce(
func.sum(delta).over(
partition_by=dim_cols,
order_by=period_expr,
rows=(None, -1),
),
0,
)
return (
select(
period_expr.label("period"),
*[c.label(c.key) for c in dim_cols],
starting_balance.label("starting_balance"),
delta.label("delta"),
*_split_sum_columns(table, split_by),
ending_balance.label("ending_balance"),
)
.select_from(table)
.group_by(period_expr, *dim_cols)
.order_by(period_expr, *dim_cols)
)
def _bound_expr(
value: str | ColumnElement | None,
table: object,
created_at_col: str,
*,
kind: Literal["min", "max"],
) -> ColumnElement:
"""Return a ``timestamptz`` SQL expression for a bound."""
if value is None:
agg = func.min if kind == "min" else func.max
return select(agg(table.c[created_at_col])).scalar_subquery() # type: ignore[attr-defined]
if isinstance(value, str):
return cast(literal(value), DateTime(timezone=True))
return value
[docs]
def construct_ledger_gap_filled_period_rollup_query(
source: ContextSource,
dimensions: list[str],
*,
start: str | ColumnElement | None = None,
end: str | ColumnElement | None = None,
period: Period = "day",
) -> Select:
"""Build a gap-filled period rollup query.
Generates the Cartesian product of every period in the range
``[start, end]`` with every dimension combination observed in
the ledger, then LEFT JOINs deltas back in so empty periods are
emitted with ``delta = 0``::
SELECT p.period, d.dim_col1, ...,
COALESCE(SUM(l.value), 0) AS delta
FROM generate_series(<start>, <end>,
INTERVAL '1 <period>') p(period)
CROSS JOIN (
SELECT DISTINCT dim_col1, ...
FROM <primary_table>
) d
LEFT JOIN <primary_table> l
ON date_trunc('<period>', l.created_at) = p.period
AND l.dim_col1 = d.dim_col1 AND ...
GROUP BY p.period, d.dim_col1, ...
ORDER BY p.period, d.dim_col1, ...
When *start* or *end* is ``None``, a ``MIN(created_at)``
/ ``MAX(created_at)`` subquery is substituted so the range
auto-covers the observed data.
Args:
source: Ledger source.
dimensions: Column names to group by. Must be a non-empty
list.
start: Lower bound of the period axis, or ``None`` to
auto-detect. Accepts a timestamp literal string or a
SQLAlchemy expression.
end: Upper bound of the period axis, or ``None``.
period: ``date_trunc`` precision string. Defaults to
``"day"``.
Returns:
A SQLAlchemy :class:`~sqlalchemy.Select`.
Raises:
CodegenDatabaseValidationError: If *dimensions* is empty or *period*
is not a valid ``date_trunc`` precision.
"""
if not dimensions:
msg = "dimensions must be a non-empty list"
raise CodegenDatabaseValidationError(msg)
_check_period(period)
table = source.ctx["primary"]
created_at_col = source.ctx["created_at_column"]
start_expr = _bound_expr(start, table, created_at_col, kind="min")
end_expr = _bound_expr(end, table, created_at_col, kind="max")
period_series = func.generate_series(
func.date_trunc(text(f"'{period}'"), start_expr),
func.date_trunc(text(f"'{period}'"), end_expr),
text(f"INTERVAL '1 {period}'"),
).column_valued("period")
period_subq = select(period_series).subquery("p")
dim_cols = [table.c[d] for d in dimensions]
dim_subq = select(*dim_cols).distinct().subquery("d")
period_on_l = func.date_trunc(
text(f"'{period}'"),
table.c[created_at_col],
)
join_conditions = [
period_on_l == period_subq.c.period,
*[table.c[d] == dim_subq.c[d] for d in dimensions],
]
return (
select(
period_subq.c.period.label("period"),
*[dim_subq.c[d].label(d) for d in dimensions],
func.coalesce(func.sum(table.c["value"]), 0).label("delta"),
)
.select_from(
period_subq.join(dim_subq, text("true")).join(
table,
and_(*join_conditions),
isouter=True,
)
)
.group_by(
period_subq.c.period,
*[dim_subq.c[d] for d in dimensions],
)
.order_by(
period_subq.c.period,
*[dim_subq.c[d] for d in dimensions],
)
)
[docs]
def construct_ledger_pivoted_period_query(
source: ContextSource,
dimensions: list[str],
periods: list[str],
period: Period = "day",
) -> Select:
"""Build a pivoted period query with one column per listed period.
Rows are dimension combinations; columns are the listed period
start timestamps. Each cell holds the summed ``value`` for that
period x dimension combination::
SELECT dim_col1, ...,
COALESCE(SUM(value) FILTER (
WHERE date_trunc('<period>', created_at)
= '<periods[0]>'::timestamptz
), 0) AS "<periods[0]>",
...
FROM <primary_table>
GROUP BY dim_col1, ...
ORDER BY dim_col1, ...
Postgres requires the column set to be known at plan time, so
*periods* must be passed in as an explicit list. This mirrors
what the running-balance view provides but with periods on the
column axis instead of rows, which is often the shape a
spreadsheet or stacked bar chart wants.
Args:
source: Ledger source.
dimensions: Column names to group by. Must be a non-empty
list.
periods: Period start timestamps as ISO strings, e.g.
``["2024-01-01", "2024-01-02"]``. Each becomes a named
output column. Must be a non-empty list.
period: ``date_trunc`` precision string. Defaults to
``"day"``. Values in *periods* are matched after both
sides are truncated to this precision.
Returns:
A SQLAlchemy :class:`~sqlalchemy.Select`.
Raises:
CodegenDatabaseValidationError: If *dimensions* or *periods* is
empty, or *period* is invalid.
"""
if not dimensions:
msg = "dimensions must be a non-empty list"
raise CodegenDatabaseValidationError(msg)
if not periods:
msg = "periods must be a non-empty list"
raise CodegenDatabaseValidationError(msg)
_check_period(period)
table = source.ctx["primary"]
created_at_col = source.ctx["created_at_column"]
dim_cols = [table.c[d] for d in dimensions]
period_expr = func.date_trunc(
text(f"'{period}'"),
table.c[created_at_col],
)
pivot_cols = [
func.coalesce(
func.sum(table.c["value"]).filter(
period_expr
== func.date_trunc(
text(f"'{period}'"),
cast(literal(p), DateTime(timezone=True)),
),
),
0,
).label(p)
for p in periods
]
return (
select(
*[c.label(c.key) for c in dim_cols],
*pivot_cols,
)
.select_from(table)
.group_by(*dim_cols)
.order_by(*dim_cols)
)
[docs]
def construct_ledger_rolling_window_query(
source: ContextSource,
dimensions: list[str],
period: Period = "day",
window_size: int = 7,
) -> Select:
"""Build a rolling-window sum query over per-period deltas.
Generates::
SELECT date_trunc('<period>', created_at) AS period,
dim_col1, ...,
SUM(value) AS delta,
SUM(SUM(value)) OVER (
PARTITION BY dim_col1, ...
ORDER BY date_trunc('<period>', created_at)
ROWS BETWEEN <window_size-1> PRECEDING
AND CURRENT ROW
) AS rolling_sum
FROM <primary_table>
GROUP BY 1, dim_col1, ...
ORDER BY 1, dim_col1, ...
The rolling sum is over *rows*, not *calendar periods*: missing
periods are skipped rather than counted as zero. Use
:func:`construct_ledger_gap_filled_period_rollup_query` as a CTE
first if you need calendar-aligned rolling windows.
Args:
source: Ledger source.
dimensions: Column names to partition by. Must be a
non-empty list.
period: ``date_trunc`` precision string. Defaults to
``"day"``.
window_size: Number of rows in the trailing window,
inclusive of the current row. Must be ``>= 1``.
Returns:
A SQLAlchemy :class:`~sqlalchemy.Select`.
Raises:
CodegenDatabaseValidationError: If *dimensions* is empty, *period*
is invalid, or *window_size* is less than 1.
"""
if not dimensions:
msg = "dimensions must be a non-empty list"
raise CodegenDatabaseValidationError(msg)
_check_period(period)
if window_size < 1:
msg = f"window_size must be >= 1, got {window_size}"
raise CodegenDatabaseValidationError(msg)
table = source.ctx["primary"]
created_at_col = source.ctx["created_at_column"]
dim_cols = [table.c[d] for d in dimensions]
period_expr = func.date_trunc(
text(f"'{period}'"),
table.c[created_at_col],
)
delta = func.sum(table.c["value"])
rolling = func.sum(delta).over(
partition_by=dim_cols,
order_by=period_expr,
rows=(-(window_size - 1), 0),
)
return (
select(
period_expr.label("period"),
*[c.label(c.key) for c in dim_cols],
delta.label("delta"),
rolling.label("rolling_sum"),
)
.select_from(table)
.group_by(period_expr, *dim_cols)
.order_by(period_expr, *dim_cols)
)