Source code for codegen_database.factory.dimension.append_only

"""Append-only dimension resource factory."""

from __future__ import annotations

from pathlib import Path
from typing import TYPE_CHECKING, ClassVar

from sqlalchemy import (
    BigInteger,
    Column,
    DateTime,
    ForeignKey,
    Table,
    func,
    select,
)

from codegen_database.ext.audit import ActivityViewPlugin
from codegen_database.factory.base import ResourceFactory
from codegen_database.factory.dimension._server_default import (
    coalesced_value,
    literal_defaults,
)
from codegen_database.plugin import (
    Dynamic,
    Plugin,
    PluginOrCollection,
    produces,
    requires,
    singleton,
)
from codegen_database.plugins.column_name import construct_column_name_plugin
from codegen_database.plugins.fk import TableFKPlugin
from codegen_database.plugins.index import TableIndexPlugin
from codegen_database.plugins.protect import RawTableProtectionPlugin
from codegen_database.plugins.soft_delete import soft_delete_columns
from codegen_database.plugins.trigger import (
    InsteadOfTriggerPlugin,
    TriggerOp,
    render_delete_op,
)
from codegen_database.plugins.view import ViewPlugin
from codegen_database.utils.naming import resolve_name
from codegen_database.utils.query import compile_query
from codegen_database.utils.template import load_template
from codegen_database.utils.trigger import register_view_triggers

if TYPE_CHECKING:
    from collections.abc import Callable

    from codegen_database.factory.context import FactoryContext

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

_CHECK_TEMPLATES = (
    Path(__file__).resolve().parents[2] / "plugins" / "templates" / "check"
)

_NAMING_DEFAULTS = {
    "append_only_root": "%(table_name)s_root",
    "append_only_attributes": "%(table_name)s_attributes",
    "append_only_function": ("%(schema)s_%(table_name)s_%(op)s"),
    "append_only_trigger": ("%(schema)s_%(table_name)s_%(op)s"),
}

# `_check_` prefix makes Postgres fire these triggers alphabetically
# before the main `%(schema)s_%(table_name)s_%(op)s` triggers, so a
# uniqueness violation is caught before the append rewrites the row.
_CHECK_NAMING_DEFAULTS = {
    "check_function": "_check_%(schema)s_%(table_name)s_%(op)s",
    "check_trigger": "_check_%(schema)s_%(table_name)s_%(op)s",
}


def _resolve_root_name(ctx: FactoryContext) -> str:
    return resolve_name(
        ctx.metadata,
        "append_only_root",
        {
            "table_name": ctx.tablename,
            "schema": ctx.schemaname,
        },
        _NAMING_DEFAULTS,
    )


def _resolve_attributes_name(
    ctx: FactoryContext,
) -> str:
    return resolve_name(
        ctx.metadata,
        "append_only_attributes",
        {
            "table_name": ctx.tablename,
            "schema": ctx.schemaname,
        },
        _NAMING_DEFAULTS,
    )


# -- builder factories ------------------------------------------------


def _make_query_builder(
    root_key: str,
    attributes_key: str,
) -> Callable[[FactoryContext], str]:
    """Return a query builder for an append-only join view."""

    def build(ctx: FactoryContext) -> str:
        pk_col_name = ctx.pk_column_name
        created_at_col = ctx["created_at_column"]
        updated_at_col = ctx["updated_at_column"]
        root_table = ctx[root_key]
        attribute_table = ctx[attributes_key]

        view_query = (
            select(
                root_table.c[pk_col_name].label(pk_col_name),
                root_table.c[created_at_col].label(created_at_col),
                attribute_table.c["created_at"].label(updated_at_col),
                *[col.label(col.key) for col in ctx.columns],
            )
            .select_from(root_table)
            .join(
                attribute_table,
                attribute_table.c[pk_col_name]
                == root_table.c[f"{attribute_table.name}_id"],
            )
        )
        return compile_query(view_query)

    return build


def _build_proxy(ctx: FactoryContext) -> list[Column]:
    """Build proxy columns for an append-only join view."""
    pk_col_name = ctx.pk_column_name
    created_at_col = ctx["created_at_column"]
    updated_at_col = ctx["updated_at_column"]
    return [
        ctx["pk_columns"].make_derived(name=pk_col_name),
        Column(created_at_col, DateTime(timezone=True)),
        Column(updated_at_col, DateTime(timezone=True)),
        *[Column(col.key, col.type) for col in ctx.columns],
    ]


def _make_ops_builder(
    root_key: str,
    attributes_key: str,
) -> Callable[[FactoryContext], list[TriggerOp]]:
    """Return an ops builder for append-only dimensions."""

    def build(ctx: FactoryContext) -> list[TriggerOp]:
        root_table = ctx[root_key]
        attribute_table = ctx[attributes_key]
        root_fullname = f"{ctx.schemaname}.{root_table.name}"
        attr_fullname = f"{ctx.schemaname}.{attribute_table.name}"
        dim_cols = [
            col
            for col in ctx.dim_column_names
            if col not in ctx.get("exclude_from_mutations", set())
        ]

        pk_col_name = ctx.pk_column_name
        entity_col = pk_col_name + "_entity"
        # Defer to a column's server_default when the view insert omits
        # it: emit COALESCE(NEW.<col>, <default>) so an omitted NOT NULL
        # defaulted column doesn't write NULL into the attribute table.
        col_default = literal_defaults(attribute_table.columns)
        template_vars = {
            "attr_table": attr_fullname,
            "root_table": root_fullname,
            "attr_cols": ", ".join([*dim_cols, entity_col]),
            "new_cols": ", ".join(
                coalesced_value(c, col_default) for c in dim_cols
            )
            + f", NEW.{pk_col_name}",
            "old_cols": ", ".join(f"OLD.{c}" for c in dim_cols),
            "attr_fk_col": f"{attribute_table.name}_id",
            "pk_col": pk_col_name,
            "entity_col": entity_col,
        }

        return [
            TriggerOp(
                "insert",
                load_template(_TEMPLATES / "insert.plpgsql.mako").render(
                    **template_vars
                ),
            ),
            TriggerOp(
                "update",
                load_template(_TEMPLATES / "update.plpgsql.mako").render(
                    **template_vars
                ),
            ),
            render_delete_op(ctx, _TEMPLATES, template_vars),
        ]

    return build


# -- plugins -----------------------------------------------------------


[docs] @requires("pk_columns") @singleton("__append_only_unique_check__") class UniqueColumnCheckPlugin(Plugin): """Lift column-level ``unique=True`` into INSTEAD OF trigger checks. SQLAlchemy column-level uniqueness propagates onto the underlying attributes log. An append-only UPDATE rewrites into an INSERT that carries every column value forward, so the new revision collides with its own predecessor on any unique column even when the value is unchanged. For each column declared with ``unique=True`` this plugin: 1. Strips the ``unique`` flag *before* :class:`AppendOnlyTablePlugin` attaches the column to the attributes table -- once attached, SQLAlchemy materialises a ``UniqueConstraint`` eagerly and clearing the flag no longer suppresses it. 2. Registers an INSTEAD OF trigger on the public view that rejects duplicates among the *current* revisions only. The trigger is best-effort: it does not take a row lock, so two concurrent inserts can both pass the check before either commits. Use ``SERIALIZABLE`` isolation or an explicit advisory lock if airtight uniqueness is required. """
[docs] def run(self, ctx: FactoryContext) -> None: """Strip ``unique=True`` and register trigger-based checks.""" unique_cols = [c for c in ctx.columns if c.unique] if not unique_cols: return pk_col = ctx.pk_column_name view_fullname = f"{ctx.schemaname}.{ctx.tablename}" insert_checks: list[tuple[str, str]] = [] update_checks: list[tuple[str, str]] = [] for col in unique_cols: # Must clear before the column reaches the Table -- once # attached, the UniqueConstraint is already materialised # and clearing this flag has no effect. col.unique = None col_ref = col.key check_name = f"{ctx.tablename}_{col_ref}_unique" insert_checks.append( ( f"NOT EXISTS (SELECT 1 FROM {view_fullname} " # noqa: S608 f"WHERE {col_ref} = NEW.{col_ref})", check_name, ) ) # On INSERT, OLD is unassigned in PL/pgSQL -- referencing # OLD.<pk> raises "record 'old' is not assigned yet" -- # so the two ops need different bodies. update_checks.append( ( f"NOT EXISTS (SELECT 1 FROM {view_fullname} " # noqa: S608 f"WHERE {col_ref} = NEW.{col_ref} " f"AND {pk_col} <> OLD.{pk_col})", check_name, ) ) template = load_template(_CHECK_TEMPLATES / "validate.plpgsql.mako") register_view_triggers( metadata=ctx.metadata, view_schema=ctx.schemaname, view_fullname=view_fullname, tablename=ctx.tablename, ops=[ ("insert", template.render(checks=insert_checks)), ("update", template.render(checks=update_checks)), ], naming_defaults=_CHECK_NAMING_DEFAULTS, function_key="check_function", trigger_key="check_trigger", )
[docs] @produces(Dynamic("root_key"), Dynamic("attributes_key")) @requires( "pk_columns", "created_at_column", "updated_at_column", "exclude_from_mutations", "deleted_at_column", ) @singleton("__table__") class AppendOnlyTablePlugin(Plugin): """Create the root and attributes tables for an append-only dim. Args: root_key: Key in ``ctx`` for the entity root table (default ``"root_table"``). attributes_key: Key in ``ctx`` for the append-only attributes log (default ``"attributes"``). """ def __init__( self, root_key: str = "root_table", attributes_key: str = "attributes", ) -> None: """Store the context keys.""" self.root_key = root_key self.attributes_key = attributes_key
[docs] def run(self, ctx: FactoryContext) -> None: """Create root and attributes tables.""" pk_col_name = ctx.pk_column_name pk_columns = ctx["pk_columns"] created_at_col = ctx["created_at_column"] attr_name = _resolve_attributes_name(ctx) # The attributes log is internal storage; users only see it # via the join view. Use a BigInteger surrogate so we don't # pay UUID's storage and B-tree fragmentation costs on a # high-insert log table. The user's PK plugin still governs # the public-facing root and view. attributes_table = Table( attr_name, ctx.metadata, Column(pk_col_name, BigInteger, primary_key=True), # Stamp every version with the entity (root) PK so audit # queries can link historical versions to their entity -- # the root table's FK only points to the *current* version. Column( pk_col_name + "_entity", ctx["pk_columns"].first.type, nullable=False, index=True, ), Column( "created_at", DateTime(timezone=True), server_default=func.now(), ), *ctx.table_items, *soft_delete_columns(ctx), schema=ctx.schemaname, ) ctx[self.attributes_key] = attributes_table root_name = _resolve_root_name(ctx) root_fk = f"{ctx.schemaname}.{attr_name}.{pk_col_name}" root_table = Table( root_name, ctx.metadata, pk_columns.make_derived(name=pk_col_name), Column( created_at_col, DateTime(timezone=True), server_default=func.now(), ), Column( f"{attr_name}_id", ForeignKey(root_fk), ), schema=ctx.schemaname, ) ctx[self.root_key] = root_table
# -- factory functions -------------------------------------------------
[docs] def AppendOnlyViewPlugin( # noqa: N802 root_key: str = "root_table", attributes_key: str = "attributes", primary_key: str = "primary", ) -> ViewPlugin: """Create a configured ViewPlugin for append-only dimensions. Args: root_key: Key in ``ctx`` for the entity root table (default ``"root_table"``). attributes_key: Key in ``ctx`` for the attributes log (default ``"attributes"``). primary_key: Key in ``ctx`` to store the view proxy under (default ``"primary"``). Returns: A :class:`~codegen_database.plugins.view.ViewPlugin` configured for append-only join views. """ return ViewPlugin( query_builder=_make_query_builder(root_key, attributes_key), proxy_builder=_build_proxy, primary_key=primary_key, extra_requires=[ root_key, attributes_key, "pk_columns", "created_at_column", "updated_at_column", ], )
[docs] class CodegenDatabaseAppendOnly(ResourceFactory): """Create an append-only (SCD Type 2) dimension. Internal plugins (always present), in order: The column-name plugin :func:`~codegen_database.plugins.column_name.construct_column_name_plugin` sets the ``created_at`` column name. Then: 1. :class:`UniqueColumnCheckPlugin` -- lifts column-level ``unique=True`` into INSTEAD OF trigger checks (must run before the table is built so the unique flag can be stripped in time). 2. :class:`AppendOnlyTablePlugin` -- root + attributes tables. 3. :class:`AppendOnlyViewPlugin` -- join view proxy. 4. :class:`~codegen_database.plugins.index.TableIndexPlugin` -- indices on the attributes table. 5. :class:`~codegen_database.plugins.fk.TableFKPlugin` -- foreign keys on the attributes table. 6. :class:`~codegen_database.plugins.trigger.InsteadOfTriggerPlugin` -- INSTEAD OF triggers (activates when a view plugin produces ``"primary"``). :class:`~codegen_database.plugins.check.TableCheckPlugin` is auto-added by the base factory when not already present. A :class:`~codegen_database.plugins.pk.SerialPKPlugin` is auto-added when no user plugin produces ``pk_columns``. """ _FK_TARGET_KEY: ClassVar[str] = "root_table" _INTERNAL_PLUGINS: ClassVar[list[PluginOrCollection]] = [ construct_column_name_plugin("created_at_column", "created_at"), construct_column_name_plugin("updated_at_column", "updated_at"), # Must precede AppendOnlyTablePlugin: see UniqueColumnCheckPlugin. UniqueColumnCheckPlugin(), AppendOnlyTablePlugin(), AppendOnlyViewPlugin(), TableIndexPlugin(table_key="attributes"), TableFKPlugin(table_key="attributes"), RawTableProtectionPlugin("root_table", "attributes"), ActivityViewPlugin(), InsteadOfTriggerPlugin( ops_builder=_make_ops_builder("root_table", "attributes"), naming_defaults=_NAMING_DEFAULTS, function_key="append_only_function", trigger_key="append_only_trigger", view_key="primary", extra_requires=[ "root_table", "attributes", ], ), ]