Source code for codegen_database.ext.ledger.functions

"""Function spec builder for ledger factories.

Provides :func:`ledger_event_function` which compiles a
:class:`~codegen_database.ext.ledger.events.LedgerEvent` into a
:class:`~codegen_database.functions.CodegenDatabaseFunctionSpec` ready to
unpack into :class:`~codegen_database.functions.CodegenDatabaseFunction`.
"""

from __future__ import annotations

from pathlib import Path
from typing import TYPE_CHECKING

from sqlalchemy.dialects import postgresql as pg_dialect
from sqlalchemy_declarative_extensions.dialects.postgresql import (
    FunctionSecurity,
)

if TYPE_CHECKING:
    from sqlalchemy import Column, Table

    from codegen_database.factory.context import ContextSource

from codegen_database.errors import CodegenDatabaseValidationError
from codegen_database.ext.ledger.events import (
    LedgerEvent,
    ParamCollector,
    _desired_table_ref,
    _input_table_ref,
)
from codegen_database.functions import CodegenDatabaseFunctionSpec
from codegen_database.utils.query import compile_query
from codegen_database.utils.template import load_template

_TEMPLATES = (
    Path(__file__).resolve().parents[2] / "plugins" / "templates" / "ledger"
)


def _pg_type_str(col: Column) -> str:
    """Return the PostgreSQL type string for *col*.

    Args:
        col: A SQLAlchemy :class:`~sqlalchemy.Column`.

    Returns:
        Compiled PostgreSQL type string.

    """
    return col.type.compile(dialect=pg_dialect.dialect())


def _validate_single_event(
    event: LedgerEvent,
    root_table: Table,
) -> None:
    """Validate one event against the root table columns.

    Args:
        event: The event to validate.
        root_table: The ledger backing table.

    Raises:
        CodegenDatabaseValidationError: If any validation rule is violated.

    """
    if event.desired is not None and not event.diff_keys:
        msg = (
            f"LedgerEvent {event.name!r}: "
            f"diff_keys is required when desired is set."
        )
        raise CodegenDatabaseValidationError(msg)

    if event.existing is not None and event.desired is None:
        msg = (
            f"LedgerEvent {event.name!r}: existing requires desired to be set."
        )
        raise CodegenDatabaseValidationError(msg)

    col_names = {c.key for c in root_table.columns}

    for key in event.diff_keys:
        if key not in col_names:
            msg = (
                f"LedgerEvent {event.name!r}: "
                f"diff_key {key!r} is not a column on "
                f"the ledger table.  Available columns: "
                f"{sorted(col_names)}"
            )
            raise CodegenDatabaseValidationError(msg)




def ledger_event_function(
    source: ContextSource,
    event: LedgerEvent,
) -> CodegenDatabaseFunctionSpec:
    """Build a function spec for a single ledger event.

    Returns a :class:`~codegen_database.functions.CodegenDatabaseFunctionSpec`.

    The spec contains the SQL function body generated from the event
    lambdas and all metadata needed to register a PostgreSQL function
    via :class:`~codegen_database.functions.CodegenDatabaseFunction`.

    Args:
        source: A
            :class:`~codegen_database.factory.ledger.CodegenDatabaseLedger`
            instance or a
            :class:`~codegen_database.declarative.CodegenDatabaseBase`
            subclass using ``CodegenDatabaseLedger`` as its factory.
        event: The :class:`~codegen_database.ext.ledger.events.LedgerEvent` to
            compile into a function.

    Returns:
        A :class:`~codegen_database.functions.CodegenDatabaseFunctionSpec`
        ready to pass to
        :class:`~codegen_database.functions.CodegenDatabaseFunction`.

    Raises:
        CodegenDatabaseValidationError: If event configuration is invalid.

    """
    ctx = source.ctx
    # raw_table: the backing table.  Used for the INSERT (bypasses the
    # INSTEAD OF trigger so all rows land in one statement, which is
    # required for the double-entry balance check).
    raw_table: Table = ctx["raw_table"]
    # view_table (__root__): the writable view over raw_table.  Used for
    # the RETURNS type so the function is created AFTER the table (views
    # are created after tables in migration ordering).
    view_table: Table = ctx["__root__"]
    schema: str = ctx.schemaname

    _validate_single_event(event, raw_table)

    backing_table = f"{schema}.{raw_table.name}"

    # Run input lambda with ParamCollector.
    collector = ParamCollector()
    input_select = event.input(collector)
    input_sql = compile_query(input_select)

    template_ctx: dict = {
        "input_sql": input_sql,
        "backing_table": backing_table,
        "diff_mode": event.desired is not None,
    }

    if event.desired is not None:
        # Build input table ref and run desired.
        input_ref = _input_table_ref(input_select)
        desired_select = event.desired(input_ref)
        desired_sql = compile_query(desired_select)

        # Build desired table ref and run existing.
        desired_ref = _desired_table_ref(desired_select)

        if event.existing is not None:
            existing_select = event.existing(raw_table, desired_ref, input_ref)
            existing_sql = compile_query(existing_select)

        else:
            # No existing: empty select with same columns.
            existing_sql = (
                "SELECT "
                + ", ".join(
                    f"NULL::{_pg_type_str(raw_table.c[k])}"
                    if k in {c.key for c in raw_table.columns}
                    else "NULL"
                    for k in [
                        col.name for col in desired_select.selected_columns
                    ]
                )
                + " WHERE false"
            )

        # Compute column lists.
        desired_col_names = [
            col.name for col in desired_select.selected_columns
        ]
        all_cols = list(desired_col_names)

        # Existing cols: diff_keys + value, with NULL
        # padding for any extra columns (e.g. reason).
        existing_col_names = (
            [col.name for col in existing_select.selected_columns]
            if event.existing is not None
            else []
        )
        existing_cols_padded: list[str] = []

        for col_name in all_cols:
            if col_name in existing_col_names:
                existing_cols_padded.append(col_name)

            else:
                existing_cols_padded.append(f"NULL AS {col_name}")

        # Passthrough cols: in desired but not diff_keys or value.
        passthrough_cols = [
            c
            for c in desired_col_names
            if c not in event.diff_keys and c != "value"
        ]

        template_ctx.update(
            desired_sql=desired_sql,
            existing_sql=existing_sql,
            all_cols=all_cols,
            desired_cols=", ".join(desired_col_names),
            existing_cols=", ".join(existing_cols_padded),
            diff_keys=list(event.diff_keys),
            passthrough_cols=passthrough_cols,
        )

    else:
        # Simple mode: columns from input select.
        all_cols = [col.name for col in input_select.selected_columns]
        template_ctx["all_cols"] = all_cols

    fn_body: str = load_template(_TEMPLATES / "event.plpgsql.mako").render(
        **template_ctx
    )

    fn_name = f"{schema}_{ctx.tablename}_{event.name}"

    return CodegenDatabaseFunctionSpec(
        name=fn_name,
        definition=fn_body,
        language="plpgsql",
        parameters=collector.function_params,
        returns=f"SETOF {schema}.{view_table.name}",
        security=FunctionSecurity.definer,
    )