"""Temporal validity range plugin for codegen_database dimensions.
:class:`TemporalPlugin` adds ``valid_from``/``valid_to`` columns
to a dimension's backing table and provides:
- A ``construct_temporal_as_of_query()`` helper that queries the table at a
specific point in time.
- A GiST-based exclusion constraint (requiring the ``btree_gist``
PostgreSQL extension) to prevent overlapping validity ranges for
the same entity.
The plugin works with any factory type that has a user-controlled
backing table (``CodegenDatabaseSimple``, ``CodegenDatabaseAppendOnly``,
``CodegenDatabaseLedger``). The user supplies ``valid_from`` and
``valid_to`` columns in ``schema_items``; the plugin validates them
and adds the exclusion constraint.
Usage::
orders = CodegenDatabaseSimple(
tablename="orders",
schemaname="app",
metadata=metadata,
schema_items=[
Column("order_id", Integer, nullable=False),
Column("status", String, nullable=False),
Column(
"valid_from",
DateTime(timezone=True),
nullable=False,
server_default=func.now(),
),
Column("valid_to", DateTime(timezone=True), nullable=True),
],
extra_plugins=[
TemporalPlugin(
subject_columns=["order_id"],
valid_from_column="valid_from",
valid_to_column="valid_to",
),
],
)
# Point-in-time query
query = construct_temporal_as_of_query(
orders,
as_of="2024-06-15T00:00:00Z",
subject_columns=["order_id"],
)
The exclusion constraint prevents storing two overlapping validity
windows for the same ``order_id``. Requires ``btree_gist``::
CREATE EXTENSION IF NOT EXISTS btree_gist;
"""
from __future__ import annotations
from typing import TYPE_CHECKING
from sqlalchemy import DateTime, func, select, text
from sqlalchemy.dialects.postgresql import ExcludeConstraint
from codegen_database.errors import CodegenDatabaseValidationError
from codegen_database.pg_extension import PGExtension, register_pg_extension
from codegen_database.plugin import Dynamic, Plugin, produces, requires
if TYPE_CHECKING:
from sqlalchemy import Select
from codegen_database.factory.context import ContextSource, FactoryContext
[docs]
@produces("temporal_columns")
@requires(Dynamic("table_key"))
class TemporalPlugin(Plugin):
r"""Add temporal validity range support to a dimension.
Validates that *valid_from_column* and *valid_to_column* exist
in the backing table, stores their names in
``ctx["temporal_columns"]``, and appends a GiST exclusion
constraint to the table so that no two rows with the same
*subject_columns* values have overlapping ``tstzrange``\\ s.
The exclusion constraint uses::
EXCLUDE USING GIST (
col1 WITH =, ...,
tstzrange(valid_from, COALESCE(valid_to, 'infinity'), '[)')
WITH &&
)
which requires the ``btree_gist`` extension to be installed.
Args:
subject_columns: Columns that identify a unique entity across
multiple temporal versions (e.g. ``["order_id"]``).
valid_from_column: Name of the range-start column
(default ``"valid_from"``).
valid_to_column: Name of the range-end column
(default ``"valid_to"``).
table_key: Key in ``ctx`` for the backing table to add the
constraint to (default ``"raw_table"``).
Raises:
CodegenDatabaseValidationError: If *valid_from_column* or
*valid_to_column* is missing, non-nullable for
*valid_from* (it must have a value), or not a
``DateTime`` column.
"""
def __init__(
self,
subject_columns: list[str],
valid_from_column: str = "valid_from",
valid_to_column: str = "valid_to",
table_key: str = "raw_table",
) -> None:
"""Store configuration."""
if not subject_columns:
msg = "subject_columns must be a non-empty list"
raise CodegenDatabaseValidationError(msg)
self.subject_columns = list(subject_columns)
self.valid_from_column = valid_from_column
self.valid_to_column = valid_to_column
self.table_key = table_key
[docs]
def run(self, ctx: FactoryContext) -> None:
"""Validate columns and add the exclusion constraint."""
register_pg_extension(ctx.metadata, PGExtension("btree_gist"))
table = ctx[self.table_key]
col_names = {c.name for c in table.columns}
for col_name in [self.valid_from_column, self.valid_to_column]:
if col_name not in col_names:
msg = (
f"TemporalPlugin: column {col_name!r} is not present "
f"in table {table.name!r}. Add it to schema_items."
)
raise CodegenDatabaseValidationError(msg)
col = table.c[col_name]
if not isinstance(col.type, DateTime):
msg = (
f"TemporalPlugin: column {col_name!r} must be a "
f"DateTime column, got {type(col.type).__name__!r}."
)
raise CodegenDatabaseValidationError(msg)
# valid_from must not be nullable (it always has a value).
vf_col = table.c[self.valid_from_column]
if vf_col.nullable:
msg = (
f"TemporalPlugin: valid_from column "
f"{self.valid_from_column!r} must be nullable=False. "
f"Use server_default=func.now() for auto-population."
)
raise CodegenDatabaseValidationError(msg)
ctx["temporal_columns"] = {
"valid_from": self.valid_from_column,
"valid_to": self.valid_to_column,
"subject": self.subject_columns,
}
# Append a GIST exclusion constraint directly to the table so that
# SQLAlchemy includes it in CREATE TABLE DDL and Alembic picks it up
# during autogenerate. Using ExcludeConstraint (rather than a DDL
# event listener) ensures the constraint appears in op.create_table()
# as well as metadata.create_all().
constraint_name = f"{table.name}_temporal_excl"
subject_elements = [(table.c[c], "=") for c in self.subject_columns]
range_element = (
func.tstzrange(
table.c[self.valid_from_column],
func.coalesce(
table.c[self.valid_to_column],
text("'infinity'::timestamptz"),
),
"[)",
),
"&&",
)
table.append_constraint(
ExcludeConstraint(
*subject_elements,
range_element,
using="gist",
name=constraint_name,
)
)
[docs]
def construct_temporal_as_of_query(
source: ContextSource,
as_of: str,
subject_columns: list[str] | None = None,
*,
table_key: str = "raw_table",
) -> Select:
"""Build a point-in-time query for a temporal dimension.
Returns rows where::
valid_from <= :as_of AND (valid_to IS NULL OR valid_to > :as_of)
Args:
source: A factory instance or declarative class that was
built with :class:`TemporalPlugin`.
as_of: ISO 8601 timestamp string or PostgreSQL-parseable
timestamp for the point in time to query.
subject_columns: Columns identifying unique entities.
When provided, applies ``DISTINCT ON`` to return only
the most recent version per subject at *as_of*. When
``None``, returns all matching rows.
table_key: Key in ``ctx`` for the backing table
(default ``"raw_table"``).
Returns:
A SQLAlchemy :class:`~sqlalchemy.Select`.
Raises:
CodegenDatabaseValidationError: If :class:`TemporalPlugin` has not
been applied to *source*.
"""
if "temporal_columns" not in source.ctx:
msg = (
"construct_temporal_as_of_query requires TemporalPlugin "
"to have been applied to the factory. Add "
"TemporalPlugin(...) to extra_plugins."
)
raise CodegenDatabaseValidationError(msg)
temporal = source.ctx["temporal_columns"]
vf_col = temporal["valid_from"]
vt_col = temporal["valid_to"]
table = source.ctx[table_key]
# Use literal timestamptz so the query can be used as a view definition
# without bind parameters.
as_of_text = text(f"'{as_of}'::timestamptz")
q = select(table).where(
table.c[vf_col] <= as_of_text,
(table.c[vt_col].is_(None)) | (table.c[vt_col] > as_of_text),
)
if subject_columns:
subj_cols = [table.c[c] for c in subject_columns]
q = q.distinct(*subj_cols).order_by(*subj_cols, table.c[vf_col].desc())
return q