API Reference¶

async def run_alembic_upgrade(
    engine: AsyncEngine,
    alembic_config: Config,
    target_schema: str = "public",
    revision: str = "head",
) -> None:
    """Run ``alembic upgrade`` programmatically without spawning a subprocess.

    Args:
        engine: Async SQLAlchemy engine.
        alembic_config: Alembic config object (typically from ``alembic.ini``).
        target_schema: PostgreSQL schema to run migrations in. Defaults to ``"public"``.
        revision: Target revision identifier. Defaults to ``"head"``.
    """

    def _upgrade(sync_conn: Connection) -> None:
        validate_schema_name(target_schema, sync_conn)
        alembic_config.attributes["target_schema"] = target_schema
        alembic_config.attributes["connection"] = sync_conn
        try:
            command.upgrade(alembic_config, revision)
        finally:
            alembic_config.attributes.pop("connection", None)
            alembic_config.attributes.pop("target_schema", None)

    async with engine.begin() as conn:
        await conn.run_sync(_upgrade)

async def run_alembic_downgrade(
    engine: AsyncEngine,
    alembic_config: Config,
    target_schema: str = "public",
    revision: str = "base",
) -> None:
    """Run ``alembic downgrade`` programmatically without spawning a subprocess.

    Args:
        engine: Async SQLAlchemy engine.
        alembic_config: Alembic config object.
        target_schema: PostgreSQL schema to run migrations in. Defaults to ``"public"``.
        revision: Target revision identifier. Defaults to ``"base"``.
    """

    def _downgrade(sync_conn: Connection) -> None:
        validate_schema_name(target_schema, sync_conn)
        alembic_config.attributes["target_schema"] = target_schema
        alembic_config.attributes["connection"] = sync_conn
        try:
            command.downgrade(alembic_config, revision)
        finally:
            alembic_config.attributes.pop("connection", None)
            alembic_config.attributes.pop("target_schema", None)

    async with engine.begin() as conn:
        await conn.run_sync(_downgrade)

async def get_current_revision(
    engine: AsyncEngine,
    target_schema: str = "public",
) -> str | None:
    """Return the current Alembic revision for the given schema, or ``None`` if at base.

    Args:
        engine: Async SQLAlchemy engine.
        target_schema: PostgreSQL schema to inspect. Defaults to ``"public"``.
    """

    def _get_rev(sync_conn: Connection) -> str | None:
        validate_schema_name(target_schema, sync_conn)
        ctx = MigrationContext.configure(
            sync_conn,
            opts={"version_table_schema": target_schema},
        )
        return ctx.get_current_revision()  # type: ignore[no-any-return]

    async with engine.connect() as conn:
        return await conn.run_sync(_get_rev)

def get_all_revisions(alembic_config: Config) -> list[str]:
    """Return all migration revision IDs in chronological order (base → head).

    Args:
        alembic_config: Alembic config object.
    """
    script = ScriptDirectory.from_config(alembic_config)
    # walk_revisions() yields head → base; reverse for base → head order.
    return list(reversed([rev.revision for rev in script.walk_revisions() if rev.revision]))

async def create_isolated_migration_schema(
    migration_db_url: str,
) -> AsyncGenerator[str, None]:
    """Create a unique PostgreSQL schema for one test run, yield it, then drop it.

    Uses a dedicated engine so that schema creation and deletion do not interfere
    with the connection pool used by the migration engine.

    Schema name format: ``test_mig_{8-char hex}``.

    Args:
        migration_db_url: Async DSN for the test database.

    Yields:
        The name of the freshly created schema.
    """
    schema = f"test_mig_{uuid.uuid4().hex[:8]}"
    validate_schema_name(schema)

    engine = create_async_engine(migration_db_url, echo=False, poolclass=NullPool)
    try:
        async with engine.connect() as conn:
            await conn.execute(CreateSchema(schema))
            await conn.commit()

        yield schema

    finally:
        async with engine.connect() as conn:
            await conn.execute(DropSchema(schema, cascade=True, if_exists=True))
            await conn.commit()
        await engine.dispose()

def validate_schema_name(name: str, connection: Connection | None = None) -> None:
    """Validate a PostgreSQL schema name to prevent SQL injection and ensure sanity.

    Checks format, length, and optionally reserved words (requires a live connection).
    """
    if not name:
        raise EmptySchemaNameError("Schema name cannot be empty.")
    if not _PG_IDENTIFIER_RE.match(name):
        raise InvalidSchemaNameError(
            f"Invalid schema name: {name!r}. "
            "Must start with a letter or underscore and contain only alphanumeric characters and underscores."
        )
    if len(name) > _PG_MAX_IDENTIFIER_LEN:
        raise SchemaNameTooLongError(
            f"Schema name too long: {name!r}. Maximum length is {_PG_MAX_IDENTIFIER_LEN} characters."
        )
    if connection is not None:
        reserved = get_pg_reserved_words(connection)
        if name.lower() in reserved:
            raise ReservedWordSchemaNameError(f"Schema name {name!r} is a PostgreSQL reserved word.")

def get_pg_reserved_words(connection: Connection) -> set[str]:
    """Fetch PostgreSQL reserved keywords from system catalog."""
    result = connection.execute(text("SELECT word FROM pg_get_keywords() WHERE catcode = 'R'"))
    return {row[0].lower() for row in result if row[0]}

def is_ignored_diff_item(diff_item: tuple, ignore_tables: frozenset[str]) -> bool:
    """Return ``True`` if this diff item should be excluded from schema diff checks.

    Filters out tables and their indexes that exist in the database but are absent
    from ORM metadata — for example, partition tables auto-created by PostgreSQL.

    Args:
        diff_item: A single item from Alembic's ``compare_metadata()`` result.
        ignore_tables: Set of table names to skip.
    """
    if len(diff_item) < 2:
        return False
    op, first = diff_item[0], diff_item[1]
    if op == "remove_table":
        name = getattr(first, "name", None)
        return name in ignore_tables if name else False
    if op == "remove_index":
        table = getattr(first, "table", None)
        name = getattr(table, "name", None) if table is not None else None
        return name in ignore_tables if name else False
    return False

def compare_check_constraints(
    sync_conn: Connection,
    metadata: MetaData,
    schema: str,
    ignore_tables: frozenset[str] = frozenset(),
) -> list[str]:
    """Compare the CHECK constraints ``metadata`` declares with the ones in ``schema``, by name.

    Alembic's ``compare_metadata()`` never looks at CHECK constraints. Names are resolved
    the way the DDL would render them — the metadata's naming convention applied, a
    deferred name such as ``Boolean(create_constraint=True)`` filled in, anything over 63
    characters truncated — through the same Alembic helper autogenerate uses for index
    and unique constraint names.

    A named constraint the models declare and the database lacks is always reported.
    One the database has and the models do not is reported unless the table also
    declares an unnamed check constraint, which carries whatever name PostgreSQL gave
    it. Unnamed constraints are never compared.

    Args:
        sync_conn: Synchronous SQLAlchemy connection.
        metadata: SQLAlchemy ``MetaData`` of the ORM models.
        schema: PostgreSQL schema the migrations ran in.
        ignore_tables: Table names to skip.

    Returns:
        One line per difference; empty when the constraints match.
    """
    inspector = inspect(sync_conn)
    differences: list[str] = []
    for table in metadata.sorted_tables:
        if table.name in ignore_tables:
            continue
        names = {
            _get_constraint_final_name(constraint, sync_conn.dialect)
            for constraint in table.constraints
            if isinstance(constraint, CheckConstraint)
        }
        expected = {name for name in names if name}
        actual = {name for c in inspector.get_check_constraints(table.name, schema=schema) if (name := c["name"])}
        for name in sorted(expected - actual):
            differences.append(
                f"Check constraint '{name}' on table '{table.name}' is in the models but not in the database."
            )
        if None not in names:
            for name in sorted(actual - expected):
                differences.append(
                    f"Check constraint '{name}' on table '{table.name}' is in the database but not in the models."
                )
    return differences

def compare_enums(
    sync_conn: Connection,
    metadata: MetaData,
    schema: str,
    ignore_tables: frozenset[str] = frozenset(),
) -> list[str]:
    """Compare the values of every native ``Enum`` column in ``metadata`` with its type in the database.

    Alembic's ``compare_metadata()`` never looks at enum members. Values are compared as
    ordered lists, so a value added in the wrong position counts as a difference. The type
    is looked up in ``Enum.schema`` when set, otherwise in ``schema``. Non-native and
    unnamed enums have no type in the database and are skipped; a type in the database
    that no column uses is not reported.

    Args:
        sync_conn: Synchronous SQLAlchemy connection.
        metadata: SQLAlchemy ``MetaData`` of the ORM models.
        schema: PostgreSQL schema the migrations ran in.
        ignore_tables: Table names to skip.

    Returns:
        One line per difference; empty when the enums match.
    """
    inspector = inspect(sync_conn)
    assert isinstance(inspector, PGInspector)
    expected: dict[tuple[str, str], list[str]] = {}
    for table in metadata.sorted_tables:
        if table.name in ignore_tables:
            continue
        for column in table.columns:
            enum = column.type
            if isinstance(enum, Enum) and enum.native_enum and enum.name:
                expected[(enum.schema or schema, enum.name)] = list(enum.enums)
    actual = {(e["schema"], e["name"]): e["labels"] for e in inspector.get_enums(schema="*")}

    differences: list[str] = []
    for (type_schema, name), values in sorted(expected.items()):
        labels = actual.get((type_schema, name))
        if labels is None:
            differences.append(
                f"Enum type '{name}' in schema '{type_schema}' is in the models but not in the database."
            )
        elif labels != values:
            differences.append(f"Enum type '{name}' has values {labels} in the database and {values} in the models.")
    return differences

def fetch_table_naming_results(
    sync_conn: Connection,
    schema: str,
) -> dict[str, TableNamingResults]:
    """Fetch index, foreign key, and constraint names for every table in the given schema.

    The ``alembic_version`` table is always excluded.

    Args:
        sync_conn: Synchronous SQLAlchemy connection.
        schema: PostgreSQL schema to inspect.

    Returns:
        Mapping of ``table_name`` → :class:`TableNamingResults`.
    """
    inspector = inspect(sync_conn)
    tables = inspector.get_table_names(schema=schema)

    results: dict[str, TableNamingResults] = {}
    for table in tables:
        if table == "alembic_version":
            continue
        indexes = inspector.get_indexes(table, schema=schema)
        fks = inspector.get_foreign_keys(table, schema=schema)
        check_constraints = inspector.get_check_constraints(table, schema=schema)
        unique_constraints = inspector.get_unique_constraints(table, schema=schema)
        pk_constraint = inspector.get_pk_constraint(table, schema=schema)

        pk_name: str | None = pk_constraint.get("name") if pk_constraint else None

        results[table] = {
            "indexes": {idx["name"] for idx in indexes if idx["name"]},
            "fks": [cast(ForeignKeyInfo, fk) for fk in fks],
            "check_constraints": {n for c in check_constraints if (n := c.get("name"))},
            "unique_constraints": {n for c in unique_constraints if (n := c.get("name"))},
            "pk_constraint": pk_name,
        }

    return results

def validate_naming_results(
    results: dict[str, TableNamingResults],
    allowed_index_prefixes: list[str],
    allowed_index_suffixes: list[str],
    allowed_fk_prefixes: list[str],
    allowed_fk_suffixes: list[str],
    allowed_check_prefixes: list[str],
    allowed_check_suffixes: list[str],
    allowed_uq_prefixes: list[str],
    allowed_uq_suffixes: list[str],
    allowed_pk_prefixes: list[str],
    allowed_pk_suffixes: list[str],
) -> None:
    """Assert that every index, foreign key, and constraint follows naming conventions.

    Each name must match at least one allowed prefix **or** one allowed suffix.
    An optional trailing digit is accepted on index suffixes to accommodate PostgreSQL
    partition auto-naming (e.g. ``users_pkey1``).

    Args:
        results: Output of :func:`fetch_table_naming_results`.
        allowed_index_prefixes: Prefixes an index name may start with (e.g. ``["idx_", "uq_"]``).
        allowed_index_suffixes: Suffixes an index name may end with (e.g. ``["_pkey", "_idx"]``).
        allowed_fk_prefixes: Prefixes a foreign key name may start with (e.g. ``["fk_"]``).
        allowed_fk_suffixes: Suffixes a foreign key name must end with (e.g. ``["_fkey"]``).
        allowed_check_prefixes: Prefixes a check constraint name may start with (e.g. ``["chk_"]``).
        allowed_check_suffixes: Suffixes a check constraint name may end with.
        allowed_uq_prefixes: Prefixes a unique constraint name may start with (e.g. ``["uq_"]``).
        allowed_uq_suffixes: Suffixes a unique constraint name may end with.
        allowed_pk_prefixes: Prefixes a primary key constraint name may start with (e.g. ``["pk_"]``).
        allowed_pk_suffixes: Suffixes a primary key constraint name may end with (e.g. ``["_pkey"]``).
    """
    idx_prefix_pats, idx_suffix_pats = _make_pats(allowed_index_prefixes, allowed_index_suffixes, trailing_digit=True)
    fk_prefix_pats, fk_suffix_pats = _make_pats(allowed_fk_prefixes, allowed_fk_suffixes)
    chk_prefix_pats, chk_suffix_pats = _make_pats(allowed_check_prefixes, allowed_check_suffixes)
    uq_prefix_pats, uq_suffix_pats = _make_pats(allowed_uq_prefixes, allowed_uq_suffixes)
    pk_prefix_pats, pk_suffix_pats = _make_pats(allowed_pk_prefixes, allowed_pk_suffixes)

    for table, data in results.items():
        for idx_name in data["indexes"]:
            valid = any(p.match(idx_name) for p in idx_prefix_pats) or any(p.match(idx_name) for p in idx_suffix_pats)
            assert valid, (
                f"Index '{idx_name}' on table '{table}' does not follow naming conventions. "
                f"Allowed prefixes: {allowed_index_prefixes}, suffixes: {allowed_index_suffixes} "
                "(optional trailing digit permitted)."
            )

        for fk in data["fks"]:
            fk_name = fk.get("name")
            if fk_name:
                assert _is_valid(fk_name, fk_prefix_pats, fk_suffix_pats), (
                    f"Foreign key '{fk_name}' on table '{table}' does not follow naming conventions. "
                    f"Allowed prefixes: {allowed_fk_prefixes}, suffixes: {allowed_fk_suffixes}."
                )

        for chk_name in data["check_constraints"]:
            assert _is_valid(chk_name, chk_prefix_pats, chk_suffix_pats), (
                f"Check constraint '{chk_name}' on table '{table}' does not follow naming conventions. "
                f"Allowed prefixes: {allowed_check_prefixes}, suffixes: {allowed_check_suffixes}."
            )

        for uq_name in data["unique_constraints"]:
            assert _is_valid(uq_name, uq_prefix_pats, uq_suffix_pats), (
                f"Unique constraint '{uq_name}' on table '{table}' does not follow naming conventions. "
                f"Allowed prefixes: {allowed_uq_prefixes}, suffixes: {allowed_uq_suffixes}."
            )

        pk_name = data["pk_constraint"]
        if pk_name:
            assert _is_valid(pk_name, pk_prefix_pats, pk_suffix_pats), (
                f"Primary key '{pk_name}' on table '{table}' does not follow naming conventions. "
                f"Allowed prefixes: {allowed_pk_prefixes}, suffixes: {allowed_pk_suffixes}."
            )

def rules_from_metadata(metadata: MetaData) -> NamingConventionRules:
    """Derive :class:`NamingConventionRules` from ``metadata.naming_convention``.

    Only literal prefixes and suffixes are extracted — fully dynamic templates
    (no literal prefix or suffix) are silently ignored.

    The SQLAlchemy built-in default convention (``{"ix": "ix_%(column_0_label)s"}``)
    is ignored; only user-supplied conventions are considered.

    Args:
        metadata: SQLAlchemy ``MetaData`` instance.

    Returns:
        :class:`NamingConventionRules` with lists populated from the convention.
        Fields for which no literal prefix/suffix could be extracted are empty.
    """
    convention = getattr(metadata, "naming_convention", {}) or {}
    if set(convention.keys()) == _SQLALCHEMY_DEFAULT_CONVENTION:
        return NamingConventionRules()
    rules = NamingConventionRules()

    for key, (prefix_attr, suffix_attr) in _CONVENTION_KEY_MAP.items():
        template = convention.get(key)
        if not isinstance(template, str):
            continue

        prefix = _extract_prefix(template)
        suffix = _extract_suffix(template)

        if prefix:
            getattr(rules, prefix_attr).append(prefix)
        if suffix:
            getattr(rules, suffix_attr).append(suffix)

    return rules

@pytest.fixture
def alembic_config() -> Config:
    """Create an Alembic :class:`~alembic.config.Config` from ``alembic.ini``.

    Override this fixture in your test class to point to a different ``alembic.ini``::

        @pytest.fixture
        def alembic_config(self) -> Config:
            from pathlib import Path
            from alembic.config import Config
            ini = Path(__file__).parent.parent / "alembic.ini"
            return Config(str(ini))
    """
    return _create_alembic_config()

@pytest.fixture
async def migration_engine(migration_db_url: str) -> AsyncGenerator[AsyncEngine, None]:
    """Async engine with :class:`~sqlalchemy.pool.NullPool` for migration tests.

    NullPool prevents connection reuse between tests, which is important for
    schema isolation when running tests in parallel.

    Requires the ``migration_db_url`` fixture to be provided.
    """
    engine = create_async_engine(migration_db_url, echo=False, poolclass=NullPool)
    try:
        yield engine
    finally:
        await engine.dispose()