Source code for codegen_database.ext.refresh.query

"""Query builders for incremental refresh patterns.

:func:`construct_unincorporated_since_query` returns a SELECT for rows in the
source table that have not yet been rolled into an aggregate, based on
a watermark timestamp.

Usage::

    # Find orders not yet incorporated into the daily aggregate.
    q = construct_unincorporated_since_query(
        orders,
        watermark_col="created_at",
        since="2024-06-01T00:00:00Z",
    )

    # Without a fixed timestamp: use a subquery as the checkpoint.
    q = construct_unincorporated_since_query(
        orders,
        watermark_col="created_at",
        since_subquery="SELECT MIN(day::timestamptz) FROM app.orders_daily_agg",
    )

"""

from __future__ import annotations

from typing import TYPE_CHECKING

from sqlalchemy import select, text

from codegen_database.errors import CodegenDatabaseValidationError

if TYPE_CHECKING:
    from sqlalchemy import Select

    from codegen_database.factory.context import ContextSource




def construct_unincorporated_since_query(
    source: ContextSource,
    *,
    watermark_col: str = "created_at",
    since: str | None = None,
    since_subquery: str | None = None,
    table_key: str = "raw_table",
) -> Select:
    """Return rows in *source* not yet incorporated into an aggregate.

    Generates::

        SELECT * FROM <source_table>
        WHERE <watermark_col> >= <since>

    where ``<since>`` is either a literal timestamp or the result of a
    subquery.

    Args:
        source: A factory instance or declarative class with a
            populated context.
        watermark_col: Column in the source table to compare against
            the cutoff (default ``"created_at"``).
        since: ISO 8601 timestamp string used as the lower bound.
            Exactly one of *since* or *since_subquery* must be given.
        since_subquery: SQL expression (a scalar subquery) that returns
            a ``TIMESTAMPTZ`` cutoff.  Exactly one of *since* or
            *since_subquery* must be given.
        table_key: Key in ``ctx`` for the source table
            (default ``"raw_table"``).

    Returns:
        A SQLAlchemy :class:`~sqlalchemy.Select`.

    Raises:
        CodegenDatabaseValidationError: If neither or both of *since* and
            *since_subquery* are provided.

    """
    if (since is None) == (since_subquery is None):
        msg = "Exactly one of 'since' or 'since_subquery' must be provided."
        raise CodegenDatabaseValidationError(msg)

    table = source.ctx[table_key]

    if since is not None:
        cutoff = text(f"'{since}'::timestamptz")

    else:
        cutoff = text(f"({since_subquery})")

    return select(table).where(table.c[watermark_col] >= cutoff)