Source code for codegen_database.plugins.check

"""Check constraint plugins for codegen_database dimensions.

:class:`TableCheckPlugin` converts
:class:`~codegen_database.check.CodegenDatabaseCheck`
items into real SQLAlchemy ``CheckConstraint`` objects on a table
(for simple and append-only dimensions).
"""

from __future__ import annotations

from typing import TYPE_CHECKING

from sqlalchemy import CheckConstraint

from codegen_database.check import CodegenDatabaseCheck, collect_checks
from codegen_database.plugin import Dynamic, Plugin, requires
from codegen_database.validation import validate_column_references

if TYPE_CHECKING:
    from codegen_database.factory.context import FactoryContext


@requires(Dynamic("table_key"))
class _CheckPlugin(Plugin):
    """Base class for check-constraint enforcement plugins.

    Handles the shared logic of collecting
    :class:`~codegen_database.check.CodegenDatabaseCheck` items,
    performing an early exit when there are none, and validating that
    each check only
    references known columns.  Concrete subclasses supply the column
    name set and the application logic.

    """

    def __init__(self, table_key: str = "primary") -> None:
        """Store the context key."""
        self.table_key = table_key

    def _column_names(self, ctx: FactoryContext) -> set[str]:
        """Return the set of column names available for validation.

        Args:
            ctx: The active factory context.

        Returns:
            Set of known column names for this plugin's target.

        """
        raise NotImplementedError

    def _apply(
        self,
        ctx: FactoryContext,
        checks: list[CodegenDatabaseCheck],
    ) -> None:
        """Apply validated checks to the target.

        Args:
            ctx: The active factory context.
            checks: Validated
                :class:`~codegen_database.check.CodegenDatabaseCheck`
                items to apply.

        """
        raise NotImplementedError

    def run(self, ctx: FactoryContext) -> None:
        """Collect, validate, and apply checks."""
        checks = collect_checks(ctx.schema_items)

        if not checks:
            return

        col_names = self._column_names(ctx)

        for cave_check in checks:
            validate_column_references(
                f"CodegenDatabaseCheck {cave_check.name!r}",
                cave_check.column_names(),
                col_names,
            )

        self._apply(ctx, checks)




class TableCheckPlugin(_CheckPlugin):
    """Materialize check items as table constraints.

    Turns :class:`~codegen_database.check.CodegenDatabaseCheck` items
    into SQLAlchemy ``CheckConstraint`` objects.

    Reads ``CodegenDatabaseCheck`` items from ``ctx.schema_items``, resolves
    ``{col}`` markers to plain column names (identity), and appends
    real ``CheckConstraint`` objects to the target table.

    Args:
        table_key: Key in ``ctx`` for the target table
            (default ``"primary"``).

    """

    def _column_names(self, ctx: FactoryContext) -> set[str]:
        """Return column names from the physical table."""
        return {c.name for c in ctx[self.table_key].columns}

    def _apply(
        self,
        ctx: FactoryContext,
        checks: list[CodegenDatabaseCheck],
    ) -> None:
        """Append SQLAlchemy CheckConstraints to the table."""
        table = ctx[self.table_key]

        for cave_check in checks:
            expr = cave_check.resolve(lambda c: c)
            constraint = CheckConstraint(expr, name=cave_check.name)
            table.append_constraint(constraint)