Source code for codegen_database.ext.ledger.queries

"""Query builders for ledger factories.

Provides pure-Python helpers that return SQLAlchemy selects
without registering anything on metadata.  Callers are responsible
for passing the results to
:class:`~codegen_database.views.view.CodegenDatabasePlainView`.
"""

from __future__ import annotations

from typing import TYPE_CHECKING

from sqlalchemy import func, select

if TYPE_CHECKING:
    from sqlalchemy import Select

    from codegen_database.factory.context import ContextSource

from codegen_database.errors import CodegenDatabaseValidationError




def construct_ledger_balance_query(
    source: ContextSource,
    dimensions: list[str],
) -> Select:
    """Build a balance aggregation query for a ledger.

    Generates::

        SELECT dim_col1, ..., SUM(value) AS balance
        FROM <primary_table>
        GROUP BY 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.

    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"]
    dim_cols = [table.c[d] for d in dimensions]
    return (
        select(
            *[c.label(c.key) for c in dim_cols],
            func.sum(table.c["value"]).label("balance"),
        )
        .select_from(table)
        .group_by(*dim_cols)
    )





def construct_ledger_latest_query(
    source: ContextSource,
    dimensions: list[str],
) -> Select:
    """Build a latest-row query for a ledger using DISTINCT ON.

    Generates::

        SELECT * FROM <primary_table>
        DISTINCT ON (dim_col1, ...)
        ORDER BY dim_col1, ..., created_at DESC

    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]
    return (
        select(table)
        .distinct(*dim_cols)
        .order_by(
            *dim_cols,
            table.c[created_at_col].desc(),
        )
    )