Source code for codegen_database.plugins.fk

"""Foreign key plugin for codegen_database dimensions.

:class:`TableFKPlugin` converts inline
:class:`~codegen_database.fk.CodegenDatabaseForeignKey` column markers into real
SQLAlchemy ``ForeignKeyConstraint`` objects on a table.  Two-part
references (``"dimension.column"``) are resolved via the dimension
registry in ``metadata.info``.
"""

from __future__ import annotations

from typing import TYPE_CHECKING

from codegen_database.errors import CodegenDatabaseValidationError
from codegen_database.fk import (
    _INLINE_FK_INFO_KEY,
    CodegenDatabaseForeignKey,
    defer_fk_resolution,
    materialize_fk,
    resolve_fk_reference,
)
from codegen_database.plugin import Dynamic, Plugin, requires

if TYPE_CHECKING:
    from codegen_database.factory.context import FactoryContext

#: Number of parts in a fully-qualified ``schema.table.column`` ref.
_FULLY_QUALIFIED_PARTS = 3

#: Number of parts in a ``dimension.column`` ref.
_DIMENSION_PARTS = 2




@requires(Dynamic("table_key"))
class TableFKPlugin(Plugin):
    """Materialize FK declarations as ``ForeignKeyConstraint`` objects.

    Handles inline
    :class:`~codegen_database.fk.CodegenDatabaseForeignKey` markers
    attached to ``Column`` constructors (single-column).

    Two-part ``"dimension.column"`` references are resolved via the
    dimension registry.  Factories run at class-creation time, so a
    reference may name a dimension whose model hasn't been imported
    yet -- those are deferred and materialize when the dimension
    registers (see :func:`~codegen_database.fk.register_dimension`);
    import order never constrains who may reference whom.
    Three-part ``"schema.table.column"`` references are passed
    through directly.

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

    """

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



    def run(self, ctx: FactoryContext) -> None:
        """Collect and create foreign key constraints from inline markers."""
        table = ctx[self.table_key]
        registry: dict[str, object] = ctx.metadata.info.get(
            "codegen_database_dimensions", {}
        )

        for col in table.columns:
            inline_fks: list[CodegenDatabaseForeignKey] = col.info.get(
                _INLINE_FK_INFO_KEY, []
            )

            for inline_fk in inline_fks:
                ref = inline_fk.reference
                parts = ref.split(".")

                if len(parts) == _FULLY_QUALIFIED_PARTS:
                    resolved_ref = ref

                elif len(parts) != _DIMENSION_PARTS:
                    msg = (
                        f"FK reference {ref!r} must be "
                        f"'dimension.column' or "
                        f"'schema.table.column' format."
                    )
                    raise CodegenDatabaseValidationError(msg)

                elif parts[0] not in registry:
                    # The dimension's model may simply not be
                    # imported yet; register_dimension materializes
                    # this entry when it arrives.
                    defer_fk_resolution(
                        ctx.metadata,
                        parts[0],
                        table,
                        col.name,
                        inline_fk,
                    )
                    continue

                else:
                    resolved_ref = resolve_fk_reference(ctx.metadata, ref)

                materialize_fk(
                    ctx.metadata,
                    table,
                    col.name,
                    inline_fk,
                    resolved_ref,
                )