"""Foreign key support for codegen_database dimensions.
Provides :class:`CodegenDatabaseForeignKey` — an inline single-column FK
passed directly to a ``Column(...)`` constructor, analogous to
SQLAlchemy's ``ForeignKey``. Accepts either a two-part
``"dimension.column"`` reference (resolved via the registry) or a
fully-qualified ``"schema.table.column"`` reference.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any
from sqlalchemy import Column, ForeignKeyConstraint
from sqlalchemy.sql.schema import SchemaItem
from codegen_database.errors import CodegenDatabaseValidationError
if TYPE_CHECKING:
from sqlalchemy import MetaData, Table
from sqlalchemy.sql.base import SchemaEventTarget
[docs]
@dataclass(frozen=True)
class DimensionRef:
"""Registry entry for a dimension's FK-targetable table.
Stored in ``metadata.info["codegen_database_dimensions"]`` keyed by
dimension name (``tablename``).
Args:
schema: PostgreSQL schema name.
table: Physical table name for FK targets.
"""
schema: str
table: str
[docs]
def register_dimension(
metadata: MetaData,
name: str,
ref: DimensionRef,
) -> None:
"""Register a dimension for FK resolution.
Also materializes any FK declarations that were deferred
because they referenced *name* before it was registered --
model import order doesn't constrain who may reference whom.
Args:
metadata: SQLAlchemy MetaData instance.
name: Dimension name (``tablename``).
ref: The dimension's FK target info.
"""
registry: dict[str, DimensionRef] = metadata.info.setdefault(
"codegen_database_dimensions", {}
)
registry[name] = ref
pending: dict[str, list[_PendingFK]] = metadata.info.get(
_PENDING_FK_INFO_KEY, {}
)
for entry in pending.pop(name, []):
_, col_name = entry.fk.reference.split(".")
materialize_fk(
metadata,
entry.table,
entry.column_name,
entry.fk,
f"{ref.schema}.{ref.table}.{col_name}",
)
[docs]
def resolve_fk_reference(
metadata: MetaData,
reference: str,
) -> str:
"""Resolve a ``"dimension.column"`` reference.
Looks up the dimension name in the registry and expands it
to ``"schema.table.column"``.
Args:
metadata: SQLAlchemy MetaData for registry lookup.
reference: Two-part ``"dimension.column"`` string.
Returns:
Fully qualified ``"schema.table.column"`` string.
Raises:
CodegenDatabaseValidationError: If the reference does not
contain exactly one dot, or names an unknown
dimension.
"""
parts = reference.split(".")
if len(parts) != 2: # noqa: PLR2004
msg = (
f"FK reference {reference!r} must be "
f"'dimension.column' format. "
f"Use CodegenDatabaseForeignKey with a three-part "
f"'schema.table.column' reference instead."
)
raise CodegenDatabaseValidationError(msg)
dim_name, col_name = parts
registry: dict[str, DimensionRef] = metadata.info.get(
"codegen_database_dimensions", {}
)
if dim_name not in registry:
known = sorted(registry)
msg = (
f"FK reference {reference!r} names unknown "
f"dimension {dim_name!r}. "
f"Known dimensions: {known}. "
f"Use a three-part 'schema.table.column' reference "
f"for tables outside codegen_database."
)
raise CodegenDatabaseValidationError(msg)
ref = registry[dim_name]
return f"{ref.schema}.{ref.table}.{col_name}"
_INLINE_FK_INFO_KEY = "codegen_database_fks"
#: ``metadata.info`` key holding FK declarations whose dimension was
#: not registered yet at factory time, keyed by dimension name.
_PENDING_FK_INFO_KEY = "codegen_database_pending_fks"
@dataclass(frozen=True)
class _PendingFK:
"""A deferred FK declaration awaiting its dimension.
Args:
table: The table the constraint will be appended to.
column_name: The referencing column on *table*.
fk: The original inline FK marker (carries the two-part
reference and the referential actions).
"""
table: Table
column_name: str
fk: CodegenDatabaseForeignKey
[docs]
def defer_fk_resolution(
metadata: MetaData,
dimension: str,
table: Table,
column_name: str,
fk: CodegenDatabaseForeignKey,
) -> None:
"""Park an FK declaration until *dimension* is registered.
Factories run at class-creation time, so a two-part reference
can name a dimension whose model simply hasn't been imported
yet. :func:`register_dimension` materializes parked entries
the moment the dimension arrives; :func:`validate_fks_resolved`
reports any that never do.
Args:
metadata: SQLAlchemy MetaData instance.
dimension: The (not yet registered) dimension name.
table: The table the constraint will be appended to.
column_name: The referencing column on *table*.
fk: The original inline FK marker.
"""
pending: dict[str, list[_PendingFK]] = metadata.info.setdefault(
_PENDING_FK_INFO_KEY, {}
)
pending.setdefault(dimension, []).append(
_PendingFK(table=table, column_name=column_name, fk=fk)
)
[docs]
def validate_fks_resolved(metadata: MetaData) -> None:
"""Fail if any deferred FK never found its dimension.
Call after every model module is imported (the alembic
``configure_metadata`` hook does) -- a leftover entry means the
reference names a dimension that doesn't exist, and silently
omitting the constraint would be far worse than failing here.
Args:
metadata: SQLAlchemy MetaData instance.
Raises:
CodegenDatabaseValidationError: A two-part FK reference
names a dimension that was never registered.
"""
pending: dict[str, list[_PendingFK]] = metadata.info.get(
_PENDING_FK_INFO_KEY, {}
)
if not pending:
return
registry: dict[str, DimensionRef] = metadata.info.get(
"codegen_database_dimensions", {}
)
dangling = ", ".join(
f"{entry.fk.reference!r} (on {entry.table.fullname}."
f"{entry.column_name})"
for entries in pending.values()
for entry in entries
)
msg = (
f"FK reference(s) name dimensions that were never "
f"registered: {dangling}. "
f"Known dimensions: {sorted(registry)}. "
f"Use a three-part 'schema.table.column' reference for "
f"tables outside codegen_database."
)
raise CodegenDatabaseValidationError(msg)
[docs]
def validate_fk_target(resolved_ref: str, metadata: MetaData) -> None:
"""Validate a fully-qualified FK target against the metadata.
Only validates when the target table is present in *metadata*
(i.e. it is a codegen_database-managed table). External tables are
silently skipped.
Args:
resolved_ref: ``"schema.table.column"`` string.
metadata: SQLAlchemy MetaData instance.
Raises:
CodegenDatabaseValidationError: If the target column does not exist
on a known table.
"""
schema, tname, col_name = resolved_ref.rsplit(".", 2)
table_key = f"{schema}.{tname}"
target = metadata.tables.get(table_key)
if target is None:
return
if col_name not in target.c:
known = sorted(c.name for c in target.columns)
msg = (
f"FK target {resolved_ref!r}: column {col_name!r} "
f"does not exist on {table_key!r}. "
f"Known columns: {known}."
)
raise CodegenDatabaseValidationError(msg)
[docs]
def materialize_fk(
metadata: MetaData,
table: Table,
column_name: str,
fk: CodegenDatabaseForeignKey,
resolved_ref: str,
) -> None:
"""Append the ``ForeignKeyConstraint`` for an inline FK marker.
Args:
metadata: SQLAlchemy MetaData instance (for target
validation).
table: The table to append the constraint to.
column_name: The referencing column on *table*.
fk: The inline FK marker carrying referential actions.
resolved_ref: Fully qualified ``"schema.table.column"``
target.
"""
validate_fk_target(resolved_ref, metadata)
table.append_constraint(
ForeignKeyConstraint(
[column_name],
[resolved_ref],
ondelete=fk.ondelete,
onupdate=fk.onupdate,
)
)
[docs]
class CodegenDatabaseForeignKey(SchemaItem):
"""Inline single-column FK with codegen_database dimension resolution.
Pass directly to a ``Column`` constructor like SQLAlchemy's
``ForeignKey``. The *reference* string accepts two formats:
- ``"dimension.column"`` — resolved via the dimension registry
at factory time (two dot-separated parts)::
Column("user_id", Integer, CodegenDatabaseForeignKey("users.id"))
- ``"schema.table.column"`` — passed through directly::
Column("user_id", Integer,
CodegenDatabaseForeignKey("public.users_raw.id"))
Args:
reference: Target reference string.
ondelete: ON DELETE action (e.g. ``"CASCADE"``).
onupdate: ON UPDATE action (e.g. ``"CASCADE"``).
"""
inherit_cache = False
__visit_name__ = "codegen_database_foreign_key"
def __init__(
self,
reference: str,
*,
ondelete: str | None = None,
onupdate: str | None = None,
) -> None:
"""Store the reference and optional referential actions."""
self.reference = reference
self.ondelete = ondelete
self.onupdate = onupdate
def _set_parent(
self,
parent: SchemaEventTarget,
**_kwargs: Any, # noqa: ANN401
) -> None:
"""Attach this FK marker to the parent column's info dict."""
# codegen_database only uses this marker on Column objects.
assert isinstance(parent, Column) # noqa: S101
parent.info.setdefault(_INLINE_FK_INFO_KEY, []).append(self)
def __repr__(self) -> str:
"""Return a readable representation."""
return f"CodegenDatabaseForeignKey({self.reference!r})"