Source code for codegen_database.alembic.register

from __future__ import annotations

from typing import TYPE_CHECKING

from sqlalchemy_declarative_extensions.alembic import (
    register_alembic_events,
)

from codegen_database.alembic.renderer import (
    register_renderers,
    render_item,
)
from codegen_database.alembic.rewriter import (
    process_revision_directives,
)
from codegen_database.alembic.schema import register_schemas
from codegen_database.fk import validate_fks_resolved
from codegen_database.patches import apply_all
from codegen_database.pg_extension import (
    register_default_pg_extensions,
    register_pg_extension_alembic_events,
)

if TYPE_CHECKING:
    from sqlalchemy import MetaData

    from codegen_database.config import CodegenDatabaseConfig

__all__ = [
    "alembic_hook",
    "configure_metadata",
    "process_revision_directives",
    "render_item",
]




def alembic_hook(
    config: CodegenDatabaseConfig | None = None,
) -> None:
    """Register codegen_database's alembic extensions.

    Call before importing models.

    Usage in ``env.py``::

        from codegen_database.alembic.register import (
            alembic_hook,
            configure_metadata,
            process_revision_directives,
        )

        alembic_hook()

        # ... import models / build metadata ...

        configure_metadata(target_metadata)

    Then pass ``process_revision_directives`` to
    ``context.configure(process_revision_directives=...)``.

    Args:
        config: Optional config providing extensions
            whose ``configure_alembic()`` hooks will be
            called.

    """
    # This is from sqlalchemy-declarative-extensions
    register_alembic_events()

    # These are codegen_database specific
    register_renderers()
    apply_all()

    # Always register the pg_extension comparator so that extensions declared
    # via register_pg_extension (e.g. btree_gist by TemporalPlugin) are tracked
    # in Alembic autogenerate regardless of which extensions are enabled.
    register_pg_extension_alembic_events()

    if config is not None:
        for ext in config._resolved_extensions():
            ext.configure_alembic()





def configure_metadata(
    metadata: MetaData,
    config: CodegenDatabaseConfig | None = None,
) -> None:
    """Register schemas and extension hooks on *metadata*.

    Args:
        metadata: The SQLAlchemy ``MetaData`` to configure.
        config: Optional config providing extensions.
            If ``None``, falls back to
            ``metadata.info["codegen_database_config"]``.

    Raises:
        CodegenDatabaseValidationError: A two-part FK reference
            names a dimension no imported model registered.

    """
    # Every model module is imported by the time the consumer's
    # env.py calls this, so any FK still parked on the deferred
    # list names a dimension that doesn't exist.
    validate_fks_resolved(metadata)
    register_schemas(metadata)
    # Default extensions (pg_trgm) before the config's extension
    # hooks run, so an extension that asserts on them sees them and
    # a manual registration of the same name stays the only entry
    # the comparator emits (it dedupes by name).
    register_default_pg_extensions(metadata)
    cfg = config or metadata.info.get("codegen_database_config")

    if cfg is not None:
        for ext in cfg._resolved_extensions():
            ext.configure_metadata(metadata)
            ext.configure_alembic()