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 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 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}."
            )

@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()