arm-core/api/tests/test_migrations.py

"""Миграции 0001 (DBML v5.13) и 0002 (v5.14) на PostgreSQL 16: полный цикл, откат на шаг, пустой autogenerate,
схема = схеме ORM.

Spec data-model, требование «Миграции применяются с нуля и откатываются» (сценарии «Полный цикл» и «Автогенерация
ничего не находит»), design D-2. Alembic запускается отдельным процессом (`python -m alembic -x dburl=...`), как из
CLI: `env.py` вызывает `asyncio.run(...)`, а внутри event loop pytest-asyncio это невозможно; заодно процесс не
перенастраивает логирование тестов через `fileConfig`. Каждая проверка идёт на своей временной БД (`scratch_db_url`).
"""

import datetime
import uuid
from typing import Any

import pytest
from alembic.autogenerate import compare_metadata
from alembic.runtime.migration import MigrationContext
from sqlalchemy import Connection, make_url, text
from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine, AsyncSession, create_async_engine

from src.models import Base
from src.models.enums import ENUMS
from tests.pg_helpers import alembic_async as alembic
from tests.pg_helpers import base_url, create_scratch_db, drop_scratch_db, run_alembic, service_url
from tests.test_models_part_01_03 import _check_related_group_member_uniqueness
from tests.test_models_part_04_06 import _check_nulls_not_distinct, _make_file_with_job
from tests.test_models_part_07_09 import (
    ALL_TABLES,
    EXPECTED_ENUM_COUNT,
    EXPECTED_TABLE_COUNT,
    _check_active_lock_is_unique,
)

HEAD = "0002_dbml_v5_14"
PREVIOUS = "0001_dbml_v5_13"
VERSION_TABLE = "alembic_version"
FK_CYCLE = "fk_stored_file_current_recognition_job_id_recognition_job"


async def _scalars(conn: AsyncConnection, sql: str) -> list[Any]:
    return list((await conn.execute(text(sql))).scalars())


async def tables_in_dkb(conn: AsyncConnection) -> set[str]:
    return set(
        await _scalars(
            conn,
            "SELECT table_name FROM information_schema.tables WHERE table_schema = 'dkb'",
        )
    )


async def enums_in_dkb(conn: AsyncConnection) -> set[str]:
    return set(
        await _scalars(
            conn,
            "SELECT t.typname FROM pg_type t JOIN pg_namespace n ON n.oid = t.typnamespace "
            "WHERE n.nspname = 'dkb' AND t.typtype = 'e'",
        )
    )


async def schema_snapshot(conn: AsyncConnection) -> dict[str, Any]:
    """Каталог схемы dkb без служебной таблицы alembic: колонки, ограничения, индексы, enum'ы (с порядком значений)."""
    columns = (
        await conn.execute(
            text(
                "SELECT table_name, ordinal_position, column_name, data_type, udt_schema, udt_name, is_nullable, "
                "column_default, numeric_precision, numeric_scale FROM information_schema.columns "
                "WHERE table_schema = 'dkb' AND table_name <> :vt ORDER BY table_name, ordinal_position"
            ),
            {"vt": VERSION_TABLE},
        )
    ).all()
    constraints = (
        await conn.execute(
            text(
                "SELECT c.conname, c.conrelid::regclass::text, pg_get_constraintdef(c.oid) FROM pg_constraint c "
                "WHERE c.connamespace = 'dkb'::regnamespace AND c.conrelid::regclass::text <> 'dkb.alembic_version' "
                "AND c.contype <> 'n'"
            )
        )
    ).all()
    indexes = (
        await conn.execute(
            text("SELECT indexname, indexdef FROM pg_indexes WHERE schemaname = 'dkb' AND tablename <> :vt"),
            {"vt": VERSION_TABLE},
        )
    ).all()
    enums = (
        await conn.execute(
            text(
                "SELECT t.typname, array_agg(e.enumlabel ORDER BY e.enumsortorder) FROM pg_type t "
                "JOIN pg_namespace n ON n.oid = t.typnamespace JOIN pg_enum e ON e.enumtypid = t.oid "
                "WHERE n.nspname = 'dkb' GROUP BY t.typname"
            )
        )
    ).all()
    return {
        "columns": [tuple(row) for row in columns],
        "constraints": {row[0]: (row[1], row[2]) for row in constraints},
        "indexes": {row[0]: row[1] for row in indexes},
        "enums": {row[0]: list(row[1]) for row in enums},
    }


async def _assert_full_model(conn: AsyncConnection) -> None:
    tables = await tables_in_dkb(conn)
    assert tables - {VERSION_TABLE} == ALL_TABLES
    assert len(tables - {VERSION_TABLE}) == EXPECTED_TABLE_COUNT
    enums = await enums_in_dkb(conn)
    assert enums == set(ENUMS)
    assert len(enums) == EXPECTED_ENUM_COUNT
    assert await _scalars(conn, "SELECT version_num FROM dkb.alembic_version") == [HEAD]


async def _truncate_model_tables(engine: AsyncEngine) -> None:
    async with engine.begin() as conn:
        await conn.execute(text("TRUNCATE " + ", ".join(f"dkb.{table}" for table in sorted(ALL_TABLES)) + " CASCADE"))


@pytest.mark.integration
async def test_full_cycle_upgrade_downgrade_upgrade(scratch_db_url: str) -> None:
    """Сценарий «Полный цикл»: upgrade head → downgrade base → upgrade head; после отката схема dkb пуста."""
    engine = create_async_engine(scratch_db_url)
    try:
        await alembic(scratch_db_url, "upgrade", "head")
        async with engine.connect() as conn:
            await _assert_full_model(conn)

        await alembic(scratch_db_url, "downgrade", "base")
        async with engine.connect() as conn:
            assert await tables_in_dkb(conn) == {VERSION_TABLE}  # служебная таблица alembic остаётся, она пуста
            assert await _scalars(conn, "SELECT count(*) FROM dkb.alembic_version") == [0]
            assert await enums_in_dkb(conn) == set()
            # в схеме не осталось других объектов (последовательностей, типов, индексов) кроме alembic_version
            leftovers = await _scalars(
                conn,
                "SELECT relname FROM pg_class WHERE relnamespace = 'dkb'::regnamespace "
                "AND relname NOT IN ('alembic_version', 'alembic_version_pkc')",
            )
            assert leftovers == []
            types = await _scalars(
                conn,
                "SELECT typname FROM pg_type WHERE typnamespace = 'dkb'::regnamespace "
                "AND typname NOT IN ('alembic_version', '_alembic_version')",
            )
            assert types == []

        await alembic(scratch_db_url, "upgrade", "head")
        async with engine.connect() as conn:
            await _assert_full_model(conn)
    finally:
        await engine.dispose()


@pytest.mark.integration
async def test_downgrade_one_step(scratch_db_url: str) -> None:
    """Сценарий «Откат на один шаг»: head → 0001 возвращает схему v5.13 (без kind, старый индекс и audit_action)."""
    engine = create_async_engine(scratch_db_url)
    try:
        await alembic(scratch_db_url, "upgrade", "head")
        await alembic(scratch_db_url, "downgrade", PREVIOUS)
        async with engine.connect() as conn:
            assert await _scalars(conn, "SELECT version_num FROM dkb.alembic_version") == [PREVIOUS]
            kind_columns = await _scalars(
                conn,
                "SELECT column_name FROM information_schema.columns WHERE table_schema = 'dkb' "
                "AND table_name = 'related_group_member' AND column_name = 'kind'",
            )
            assert kind_columns == []
            indexes = set(
                await _scalars(
                    conn,
                    "SELECT indexname FROM pg_indexes WHERE schemaname = 'dkb' AND tablename = 'related_group_member'",
                )
            )
            assert "uq_related_group_member_entity_current" in indexes
            assert "uq_related_group_member_entity_kind_current" not in indexes
            [labels] = await _scalars(conn, "SELECT enum_range(NULL::dkb.audit_action)::text[]")
            assert "LEGAL_ENTITY_CHANGE" not in labels
            assert "USER_CHANGE" not in labels
            sequences = await _scalars(
                conn, "SELECT sequence_name FROM information_schema.sequences WHERE sequence_schema = 'dkb'"
            )
            assert sequences == []

        await alembic(scratch_db_url, "upgrade", "head")
        async with engine.connect() as conn:
            await _assert_full_model(conn)
            [labels] = await _scalars(conn, "SELECT enum_range(NULL::dkb.audit_action)::text[]")
            assert labels == [m.value for m in ENUMS["audit_action"]]
            sequences = set(
                await _scalars(
                    conn, "SELECT sequence_name FROM information_schema.sequences WHERE sequence_schema = 'dkb'"
                )
            )
            assert sequences == {"company_group_code_seq", "related_group_code_seq"}
    finally:
        await engine.dispose()


@pytest.mark.integration
async def test_upgrade_0002_keeps_membership_and_fills_kind(scratch_db_url: str) -> None:
    """0002 пересоздаёт `related_group_member`: строки v5.13 сохраняются, `kind` = `related_group.kind` группы."""
    engine = create_async_engine(scratch_db_url)
    now = datetime.datetime.now(datetime.UTC)
    ids = {k: uuid.uuid4() for k in ("ro", "cg", "le", "gsk", "gsz", "m1", "m2")}
    try:
        await alembic(scratch_db_url, "upgrade", PREVIOUS)
        async with engine.begin() as conn:
            await conn.execute(text("INSERT INTO dkb.regional_office (id, name) VALUES (:ro, 'ГО')"), ids)
            await conn.execute(
                text("INSERT INTO dkb.company_group (id, code, name, valid_from) VALUES (:cg, 'G1', 'ГК', :now)"),
                {**ids, "now": now},
            )
            await conn.execute(
                text("INSERT INTO dkb.legal_entity (id, inn, name, valid_from) VALUES (:le, '1', 'ЮЛ', :now)"),
                {**ids, "now": now},
            )
            for kind in ("gsk", "gsz"):
                await conn.execute(
                    text(
                        "INSERT INTO dkb.related_group (id, code, company_group_id, kind, name) "
                        "VALUES (:g, :code, :cg, :kind, 'Группа')"
                    ),
                    {"g": ids[kind], "code": kind, "cg": ids["cg"], "kind": kind},
                )
            # v5.13: одно текущее членство (в ГСЗ) и закрытое (в ГСК)
            member = text(
                "INSERT INTO dkb.related_group_member (id, related_group_id, legal_entity_id, valid_from, valid_to, "
                "is_current) VALUES (:id, :g, :le, :now, :to, :cur)"
            )
            await conn.execute(
                member, {"id": ids["m1"], "g": ids["gsz"], "le": ids["le"], "now": now, "to": None, "cur": True}
            )
            await conn.execute(
                member, {"id": ids["m2"], "g": ids["gsk"], "le": ids["le"], "now": now, "to": now, "cur": False}
            )

        await alembic(scratch_db_url, "upgrade", "head")
        async with engine.connect() as conn:
            rows = (
                await conn.execute(
                    text("SELECT id, related_group_id, kind::text, valid_to, is_current FROM dkb.related_group_member")
                )
            ).all()
            assert {tuple(r) for r in rows} == {
                (ids["m1"], ids["gsz"], "gsz", None, True),
                (ids["m2"], ids["gsk"], "gsk", now, False),
            }
            leftovers = await _scalars(
                conn, "SELECT relname FROM pg_class WHERE relnamespace = 'dkb'::regnamespace AND relname LIKE '%v5_13'"
            )
            assert leftovers == []  # временная копия удалена
    finally:
        await engine.dispose()


@pytest.mark.integration
async def test_autogenerate_finds_nothing_after_upgrade(scratch_db_url: str) -> None:
    """Сценарий «Автогенерация ничего не находит»: compare_metadata пуст, `alembic check` не видит операций."""
    await alembic(scratch_db_url, "upgrade", "head")

    def diff(sync_conn: Connection) -> list[Any]:
        context = MigrationContext.configure(
            sync_conn,
            opts={
                # те же опции сравнения, что в alembic/env.py
                "compare_type": True,
                "compare_server_default": True,
                "include_schemas": True,
                "include_name": lambda name, type_, _parents: name == "dkb" if type_ == "schema" else True,
                "version_table_schema": "dkb",
            },
        )
        return compare_metadata(context, Base.metadata)

    engine = create_async_engine(scratch_db_url)
    try:
        async with engine.connect() as conn:
            assert await conn.run_sync(diff) == []
    finally:
        await engine.dispose()

    output = await alembic(scratch_db_url, "check")
    assert "No new upgrade operations detected" in output


@pytest.mark.integration
async def test_migrated_schema_equals_orm_create_all(scratch_db_url: str) -> None:
    """Каталог после миграции = каталог `create_all` ORM: колонки, default, ограничения, индексы, enum (имена и DDL)."""
    await alembic(scratch_db_url, "upgrade", "head")
    reference_name = f"armcore_scratch_{uuid.uuid4().hex[:12]}"
    reference_url = await create_scratch_db(reference_name)
    migrated = create_async_engine(scratch_db_url)
    reference = create_async_engine(reference_url)
    try:
        async with reference.begin() as conn:
            await conn.execute(text("CREATE SCHEMA IF NOT EXISTS dkb"))
            await conn.run_sync(Base.metadata.create_all)
        async with migrated.connect() as conn:
            got = await schema_snapshot(conn)
        async with reference.connect() as conn:
            expected = await schema_snapshot(conn)
        assert got["enums"] == expected["enums"]
        assert got["columns"] == expected["columns"]
        assert got["constraints"] == expected["constraints"]
        assert got["indexes"] == expected["indexes"]
    finally:
        await migrated.dispose()
        await reference.dispose()
        await drop_scratch_db(reference_name)


@pytest.mark.integration
async def test_migrated_schema_behaviour(scratch_db_url: str) -> None:
    """После миграции: FK-цикл, partial unique, NULLS NOT DISTINCT и server_default работают на реальной БД."""
    await alembic(scratch_db_url, "upgrade", "head")
    engine = create_async_engine(scratch_db_url)
    try:
        async with engine.connect() as conn:
            fk = (
                await conn.execute(
                    text(
                        "SELECT conrelid::regclass::text, pg_get_constraintdef(oid) FROM pg_constraint "
                        "WHERE conname = :name AND connamespace = 'dkb'::regnamespace"
                    ),
                    {"name": FK_CYCLE},
                )
            ).one()
            assert fk[0] == "dkb.stored_file"
            assert fk[1] == "FOREIGN KEY (current_recognition_job_id) REFERENCES dkb.recognition_job(id)"
            rows = await conn.execute(
                text(
                    "SELECT indexname, indexdef FROM pg_indexes WHERE schemaname = 'dkb' AND indexname IN "
                    "('uq_document_file_type_account', 'uq_consolidation_lock_version_active', "
                    "'uq_related_group_member_entity_kind_current', 'uq_related_group_member_group_entity_current')"
                )
            )
            indexdefs = {row[0]: row[1] for row in rows}
            assert "NULLS NOT DISTINCT" in indexdefs["uq_document_file_type_account"]
            assert "WHERE (released_at IS NULL)" in indexdefs["uq_consolidation_lock_version_active"]
            for name in ("uq_related_group_member_entity_kind_current", "uq_related_group_member_group_entity_current"):
                assert indexdefs[name].startswith("CREATE UNIQUE INDEX")
                assert "WHERE is_current" in indexdefs[name]

        # поведенческие проверки из тестов ORM (create_all) — теперь на схеме из миграции; данные между ними
        # очищаются: помощники создают одинаковые справочные записи (ГО, группа G1)
        await _check_related_group_member_uniqueness(engine)
        await _truncate_model_tables(engine)
        async with AsyncSession(engine, expire_on_commit=False) as session:
            chain = await _make_file_with_job(session)  # замыкает FK-цикл, проверяет server_default
            await session.commit()
            await _check_nulls_not_distinct(session, chain)
        await _truncate_model_tables(engine)
        async with AsyncSession(engine, expire_on_commit=False) as session:
            await _check_active_lock_is_unique(session)
    finally:
        await engine.dispose()


@pytest.mark.integration
async def test_upgrade_into_preprovisioned_schema_without_database_create(scratch_db_url: str) -> None:
    """Схему dkb выдала платформа во владение роли миграций, права CREATE на БД у роли нет: цикл миграций проходит.

    PostgreSQL проверяет право CREATE на БД раньше, чем IF NOT EXISTS, поэтому схема создаётся только при отсутствии.
    """
    role = f"armcore_mig_{uuid.uuid4().hex[:12]}"  # роли общие для кластера — имя уникальное, удаляется в finally
    password = uuid.uuid4().hex
    scratch_url = make_url(scratch_db_url)
    role_url = scratch_url.set(username=role, password=password).render_as_string(hide_password=False)
    server = create_async_engine(service_url(), isolation_level="AUTOCOMMIT")
    scratch = create_async_engine(scratch_db_url, isolation_level="AUTOCOMMIT")
    role_created = False
    try:
        async with server.connect() as conn:
            await conn.execute(text(f"CREATE ROLE \"{role}\" LOGIN PASSWORD '{password}'"))
            role_created = True
        async with scratch.connect() as conn:
            await conn.execute(text(f'REVOKE CREATE ON DATABASE "{scratch_url.database}" FROM PUBLIC'))
            await conn.execute(text(f'CREATE SCHEMA dkb AUTHORIZATION "{role}"'))
            can_create = await conn.scalar(
                text("SELECT has_database_privilege(:role, :db, 'CREATE')"),
                {"role": role, "db": scratch_url.database},
            )
            assert can_create is False

        await alembic(role_url, "upgrade", "head")
        await alembic(role_url, "check")
        await alembic(role_url, "downgrade", "base")
        await alembic(role_url, "upgrade", "head")
        async with scratch.connect() as conn:
            await _assert_full_model(conn)
            owner = await conn.scalar(text("SELECT nspowner::regrole::text FROM pg_namespace WHERE nspname = 'dkb'"))
            assert owner == role
    finally:
        if role_created:
            async with scratch.connect() as conn:
                await conn.execute(text(f'DROP OWNED BY "{role}" CASCADE'))
            async with server.connect() as conn:
                await conn.execute(text(f'DROP ROLE "{role}"'))
        await scratch.dispose()
        await server.dispose()


def test_offline_sql_creates_schema_before_version_table() -> None:
    """Офлайн-режим (`upgrade head --sql`): условный CREATE SCHEMA dkb идёт до таблицы версий и до enum-типов."""
    output = run_alembic("postgresql+asyncpg://user:pass@localhost/offline", "upgrade", "head", "--sql")
    schema_at = output.index(
        "DO $$ BEGIN IF NOT EXISTS (SELECT 1 FROM pg_namespace WHERE nspname = 'dkb') "
        "THEN CREATE SCHEMA dkb; END IF; END $$"
    )
    assert "CREATE SCHEMA IF NOT EXISTS" not in output  # требует права CREATE на базу даже при существующей схеме
    assert schema_at < output.index("CREATE TABLE dkb.alembic_version")
    assert schema_at < output.index("CREATE TYPE dkb.entity_type AS ENUM ('ЮЛ', 'ИП')")
    assert f"INSERT INTO dkb.alembic_version (version_num) VALUES ('{PREVIOUS}')" in output
    # с двумя ревизиями версия head пишется UPDATE'ом после 0001 — офлайн-SQL доходит до head
    assert (
        f"UPDATE dkb.alembic_version SET version_num='{HEAD}' WHERE dkb.alembic_version.version_num = '{PREVIOUS}'"
        in output
    )


def test_cluster_commands_use_service_database() -> None:
    """CREATE/DROP DATABASE и роли — через служебную БД `postgres`, а не рабочую `armka` (финальное ревью M-9)."""
    assert service_url().database == "postgres"
    base = base_url()
    assert (service_url().host, service_url().port, service_url().username) == (base.host, base.port, base.username)