Source code for codegen_database.ext.chart.date_bin

"""``codegen_database_date_bin`` polyfill: ``date_bin`` with month/year strides.

PostgreSQL's built-in ``date_bin(stride interval, ts timestamptz,
origin timestamptz)`` refuses strides that contain units of month or
larger.  The reason is implementation-driven — ``date_bin`` works
by subtracting ``ts`` and ``origin`` and flooring that interval by
``stride``, which requires uniform-length strides, and months are
28-31 days.

This module registers an SQL function ``codegen_database_date_bin`` that
bridges the gap by falling through to ``date_bin`` for sub-month
strides and switching to calendar arithmetic for whole-month strides
(including quarter and year).  Mixed strides (e.g. ``'1 month 3
days'``) are rejected by returning ``NULL``.

The helper is used by the ledger chart functions in
:mod:`codegen_database.ext.ledger.chart_functions`, which accept an ``interval``
``p_period`` parameter.  It lives under :mod:`codegen_database.ext.chart` rather
than the ledger because :class:`~codegen_database.ext.chart.ChartExtension` is
what installs it during ``configure_metadata``; the ledger merely
calls into it.
"""

from __future__ import annotations

from typing import TYPE_CHECKING

from sqlalchemy_declarative_extensions.dialects.postgresql import (
    FunctionParam,
    FunctionSecurity,
    FunctionVolatility,
)

if TYPE_CHECKING:
    from codegen_database.functions import CodegenDatabaseFunctionSpec

CODEGEN_DATABASE_DATE_BIN_NAME = "codegen_database_date_bin"
"""Canonical unqualified name of the polyfill function."""

CODEGEN_DATABASE_DATE_BIN_DEFAULT_ORIGIN = "'2000-01-03'::timestamptz"
"""Default origin used by the polyfill (a Monday at midnight local).

``'2000-01-03'`` was chosen so that week strides align to Monday
(ISO-style), day/hour strides align to midnight, and month/quarter/
year strides align to the 1st of the month (because the month path
snaps via ``date_trunc('month', origin)``).
"""

_DEFINITION = """\
SELECT
    CASE
        -- Sub-month stride: native date_bin handles it.
        WHEN stride_months = 0
            THEN date_bin(stride, ts, origin)

        -- Whole-month (incl. quarter, year) stride.  Align on the
        -- 1st-of-month of origin; preserve origin's month-of-year
        -- offset so custom fiscal alignments work (pass origin as
        -- e.g. '2000-04-01' to get quarters starting in April).
        WHEN stride_sub = INTERVAL '0'
            THEN date_trunc('month', origin) + make_interval(
                months => FLOOR(
                    ((EXTRACT(year  FROM ts)
                    - EXTRACT(year  FROM origin)) * 12
                   + (EXTRACT(month FROM ts)
                    - EXTRACT(month FROM origin)))::numeric
                    / stride_months
                )::int * stride_months
            )

        -- Mixed stride (e.g. '1 month 3 days') has no uniform
        -- bucketing — surface the error as NULL rather than a
        -- silently-wrong result.
        ELSE NULL::timestamptz
    END
FROM (VALUES (
    (EXTRACT(year FROM stride) * 12
   + EXTRACT(month FROM stride))::int,
    stride - make_interval(
        years  => EXTRACT(year  FROM stride)::int,
        months => EXTRACT(month FROM stride)::int
    )
)) AS s(stride_months, stride_sub)
"""




def construct_date_bin_function(
    *,
    name: str = CODEGEN_DATABASE_DATE_BIN_NAME,
) -> CodegenDatabaseFunctionSpec:
    """Build a ``codegen_database_date_bin`` function spec.

    Returns a :class:`~codegen_database.functions.CodegenDatabaseFunctionSpec`.
    Register the returned spec once per schema that needs it (typically
    wherever a ledger chart function lives).  The ledger chart functions
    call ``<schema>.codegen_database_date_bin(p_period, created_at)`` with the
    chart function's own schema by default.

    Usage::

        class DateBin(CodegenDatabaseFunctionMixin, Base):
            __table_args__ = {"schema": "ops"}
            __funcspec__ = construct_date_bin_function()

    Signature::

        codegen_database_date_bin(
            stride interval,
            ts     timestamptz,
            origin timestamptz DEFAULT '2000-01-03'::timestamptz
        ) RETURNS timestamptz

    Semantics:
        * ``stride`` contains only sub-month units → delegate to
          Postgres's built-in ``date_bin``.
        * ``stride`` is a whole number of months (so includes quarter
          and year) → calendar arithmetic, aligned to the 1st of the
          month.
        * ``stride`` mixes month and sub-month units → returns
          ``NULL`` (no uniform bucketing possible).

    Args:
        name: Unqualified function name.  Defaults to
            ``"codegen_database_date_bin"``.  Chart functions will look for the
            helper under this name, so override only if you know what
            you are doing.

    Returns:
        A :class:`~codegen_database.functions.CodegenDatabaseFunctionSpec`
        ready to pass to
        :class:`~codegen_database.functions.CodegenDatabaseFunction`
        or assign to ``__funcspec__`` on a
        :class:`~codegen_database.declarative.CodegenDatabaseFunctionMixin`
        subclass.

    """
    # Returned as a plain dict literal -- equivalent to constructing
    # the ``CodegenDatabaseFunctionSpec`` TypedDict, but avoids importing it
    # at runtime here.  ``CodegenDatabaseFunctionSpec`` lives in
    # :mod:`codegen_database.functions`, which is import-safe from this module,
    # but the literal keeps this file dependency-free for the
    # ``date_bin`` polyfill which is otherwise just a static SQL body.
    return {
        "name": name,
        "definition": _DEFINITION,
        "language": "sql",
        "parameters": [
            FunctionParam.input("stride", "interval"),
            FunctionParam.input("ts", "timestamptz"),
            FunctionParam.input(
                "origin",
                "timestamptz",
                default=CODEGEN_DATABASE_DATE_BIN_DEFAULT_ORIGIN,
            ),
        ],
        "returns": "timestamptz",
        "security": FunctionSecurity.invoker,
        "volatility": FunctionVolatility.STABLE,
    }