X-Financial/server/tests/test_alembic_migrations.py

from __future__ import annotations

import os
import re
from collections.abc import Iterator
from pathlib import Path

import pytest
from alembic.config import Config
from sqlalchemy import create_engine, inspect, text
from sqlalchemy.engine import Engine, make_url
from sqlalchemy.exc import IntegrityError
from sqlalchemy.pool import NullPool

from alembic import command
from app.core.config import get_settings
from app.db.migration_preflight import MigrationPreflightError, validate_migration_state
from app.db.schema_ownership import MIGRATION_OWNED_TABLES, create_legacy_schema

MIGRATION_TEST_DATABASE_URL = os.getenv("MIGRATION_TEST_DATABASE_URL", "").strip()
LEGACY_PROBE_TABLE = "legacy_migration_probe_records"
HEAD_REVISION = "20260714_0004"
SERVER_DIR = Path(__file__).resolve().parents[1]
ALEMBIC_INI_PATH = SERVER_DIR / "alembic.ini"


def _normalize_probe_component(value: str) -> str:
    return re.sub(r"[^a-z0-9]+", "-", value.lower()).strip("-")


def _is_disposable_probe_host(value: str) -> bool:
    markers = ("migration-probe", "disposable-probe")
    return value in markers or any(
        value.startswith(f"{marker}-") or value.startswith(f"x-financial-{marker}-")
        for marker in markers
    )


def _is_disposable_probe_database(value: str) -> bool:
    markers = ("migration-probe", "disposable-probe")
    return value in markers or any(value.startswith(f"{marker}-") for marker in markers)


def _require_disposable_probe_url(raw_url: str) -> str:
    try:
        parsed = make_url(raw_url)
    except Exception as exc:  # pragma: no cover - SQLAlchemy 提供具体解析异常
        raise RuntimeError("MIGRATION_TEST_DATABASE_URL 不是有效的数据库 URL") from exc

    if parsed.get_backend_name() != "postgresql":
        raise RuntimeError("迁移测试只允许连接 PostgreSQL 一次性数据库")

    host = _normalize_probe_component(parsed.host or "")
    database = _normalize_probe_component(parsed.database or "")
    if not _is_disposable_probe_host(host):
        raise RuntimeError(
            "迁移测试数据库主机名必须使用 migration-probe 或 disposable-probe 前缀"
        )
    if not _is_disposable_probe_database(database):
        raise RuntimeError(
            "迁移测试数据库名必须使用 migration-probe 或 disposable-probe 前缀"
        )

    return raw_url


@pytest.fixture(scope="module")
def migration_database_url() -> Iterator[str]:
    if not MIGRATION_TEST_DATABASE_URL:
        pytest.skip("仅在显式配置 MIGRATION_TEST_DATABASE_URL 时运行一次性 PostgreSQL 迁移测试")
    database_url = _require_disposable_probe_url(MIGRATION_TEST_DATABASE_URL)
    previous_database_url = os.environ.get("DATABASE_URL")
    os.environ["DATABASE_URL"] = database_url
    get_settings.cache_clear()

    try:
        resolved_url = get_settings().resolved_database_url
        if make_url(resolved_url) != make_url(database_url):
            raise RuntimeError("运行时数据库 URL 未解析到 MIGRATION_TEST_DATABASE_URL")
        yield database_url
    finally:
        if previous_database_url is None:
            os.environ.pop("DATABASE_URL", None)
        else:
            os.environ["DATABASE_URL"] = previous_database_url
        get_settings.cache_clear()


def _alembic_config(database_url: str) -> Config:
    config = Config(str(ALEMBIC_INI_PATH))
    config.set_main_option("sqlalchemy.url", database_url.replace("%", "%%"))
    return config


def _upgrade_head(database_url: str) -> None:
    get_settings.cache_clear()
    command.upgrade(_alembic_config(database_url), "head")


def _downgrade_base(database_url: str) -> None:
    get_settings.cache_clear()
    command.downgrade(_alembic_config(database_url), "base")


def _table_names(engine: Engine) -> set[str]:
    return set(inspect(engine).get_table_names(schema="public"))


def _assert_unique_constraint(
    engine: Engine,
    table_name: str,
    constraint_name: str,
    expected_columns: tuple[str, ...],
) -> None:
    constraints = {
        str(item["name"]): tuple(item["column_names"])
        for item in inspect(engine).get_unique_constraints(table_name, schema="public")
    }
    assert constraints.get(constraint_name) == expected_columns


def _assert_indexes(
    engine: Engine,
    table_name: str,
    expected_indexes: dict[str, tuple[str, ...]],
) -> None:
    indexes = {
        str(item["name"]): tuple(item["column_names"])
        for item in inspect(engine).get_indexes(table_name, schema="public")
    }
    for index_name, expected_columns in expected_indexes.items():
        assert indexes.get(index_name) == expected_columns


def _assert_cascade_foreign_key(engine: Engine, table_name: str) -> None:
    foreign_keys = inspect(engine).get_foreign_keys(table_name, schema="public")
    matching = [
        item
        for item in foreign_keys
        if item["constrained_columns"] == ["expense_case_id"]
        and item["referred_table"] == "expense_cases"
        and item["referred_columns"] == ["id"]
    ]
    assert len(matching) == 1
    assert str(matching[0].get("options", {}).get("ondelete", "")).upper() == "CASCADE"


def _assert_composite_foreign_key(
    engine: Engine,
    table_name: str,
    constrained_columns: tuple[str, ...],
    referred_table: str,
    referred_columns: tuple[str, ...] = ("tenant_id", "id"),
) -> None:
    foreign_keys = inspect(engine).get_foreign_keys(table_name, schema="public")
    matching = [
        item
        for item in foreign_keys
        if tuple(item["constrained_columns"]) == constrained_columns
        and item["referred_table"] == referred_table
        and tuple(item["referred_columns"]) == referred_columns
    ]
    assert len(matching) == 1
    assert str(matching[0].get("options", {}).get("ondelete", "")).upper() == "RESTRICT"


def _assert_head_schema(engine: Engine) -> None:
    names = _table_names(engine)
    assert MIGRATION_OWNED_TABLES.issubset(names)
    assert "alembic_version" in names

    with engine.connect() as connection:
        assert connection.scalar(text("SELECT version_num FROM alembic_version")) == HEAD_REVISION

    _assert_unique_constraint(
        engine,
        "expense_cases",
        "uq_expense_cases_tenant_case_no",
        ("tenant_id", "case_no"),
    )
    _assert_unique_constraint(
        engine,
        "expense_case_links",
        "uq_expense_case_links_resource",
        ("resource_type", "resource_id"),
    )
    _assert_unique_constraint(
        engine,
        "business_events",
        "uq_business_events_tenant_case_id",
        ("tenant_id", "expense_case_id", "id"),
    )
    _assert_unique_constraint(
        engine,
        "business_events",
        "uq_business_event_idempotency",
        ("tenant_id", "aggregate_type", "aggregate_id", "event_type", "idempotency_key"),
    )
    _assert_unique_constraint(
        engine,
        "auth_sessions",
        "uq_auth_sessions_token_hash",
        ("token_hash",),
    )
    _assert_unique_constraint(
        engine,
        "ai_decisions",
        "uq_ai_decisions_tenant_case_id",
        ("tenant_id", "expense_case_id", "id"),
    )
    _assert_unique_constraint(
        engine,
        "ai_decisions",
        "uq_ai_decisions_tenant_idempotency",
        ("tenant_id", "idempotency_key"),
    )
    _assert_unique_constraint(
        engine,
        "ai_decisions",
        "uq_ai_decisions_tenant_preview_decision",
        ("tenant_id", "preview_decision_id"),
    )
    _assert_unique_constraint(
        engine,
        "ai_application_preview_decisions",
        "uq_ai_application_preview_decisions_issue_request",
        ("tenant_id", "actor_id", "auth_session_id", "issue_request_id"),
    )
    _assert_unique_constraint(
        engine,
        "ai_decision_feedback",
        "uq_ai_decision_feedback_tenant_idempotency",
        ("tenant_id", "idempotency_key"),
    )
    _assert_unique_constraint(
        engine,
        "workflow_outcomes",
        "uq_workflow_outcomes_tenant_idempotency",
        ("tenant_id", "idempotency_key"),
    )

    _assert_indexes(
        engine,
        "expense_cases",
        {
            "ix_expense_cases_tenant_stage": ("tenant_id", "current_stage"),
            "ix_expense_cases_tenant_status": ("tenant_id", "status"),
        },
    )
    _assert_indexes(
        engine,
        "expense_case_links",
        {"ix_expense_case_links_tenant_case": ("tenant_id", "expense_case_id")},
    )
    _assert_indexes(
        engine,
        "business_events",
        {
            "ix_business_events_aggregate": ("aggregate_type", "aggregate_id"),
            "ix_business_events_outbox": ("delivery_status", "occurred_at"),
            "ix_business_events_tenant_case_time": (
                "tenant_id",
                "expense_case_id",
                "occurred_at",
            ),
        },
    )
    _assert_indexes(
        engine,
        "auth_sessions",
        {
            "ix_auth_sessions_principal_active": (
                "principal_type",
                "revoked_at",
                "expires_at",
            ),
            "ix_auth_sessions_tenant_username": ("tenant_id", "username"),
        },
    )
    _assert_indexes(
        engine,
        "ai_application_preview_decisions",
        {
            "ix_ai_application_preview_decisions_actor_status_expiry": (
                "tenant_id",
                "actor_id",
                "status",
                "expires_at",
            ),
            "ix_ai_application_preview_decisions_conversation": (
                "tenant_id",
                "conversation_id",
                "created_at",
            ),
        },
    )
    _assert_indexes(
        engine,
        "ai_decisions",
        {
            "ix_ai_decisions_tenant_case_time": (
                "tenant_id",
                "expense_case_id",
                "created_at",
            ),
            "ix_ai_decisions_tenant_subject": (
                "tenant_id",
                "subject_type",
                "subject_id",
            ),
        },
    )
    _assert_indexes(
        engine,
        "ai_decision_feedback",
        {
            "ix_ai_decision_feedback_tenant_decision_time": (
                "tenant_id",
                "decision_id",
                "created_at",
            ),
        },
    )
    _assert_indexes(
        engine,
        "workflow_outcomes",
        {
            "ix_workflow_outcomes_tenant_case_time": (
                "tenant_id",
                "expense_case_id",
                "effective_at",
            ),
        },
    )
    _assert_cascade_foreign_key(engine, "expense_case_links")
    _assert_cascade_foreign_key(engine, "business_events")
    _assert_composite_foreign_key(
        engine,
        "ai_decisions",
        ("tenant_id", "preview_decision_id"),
        "ai_application_preview_decisions",
    )
    _assert_composite_foreign_key(
        engine,
        "ai_decisions",
        ("tenant_id", "expense_case_id"),
        "expense_cases",
    )
    _assert_composite_foreign_key(
        engine,
        "ai_decisions",
        ("tenant_id", "expense_case_id", "business_event_id"),
        "business_events",
        ("tenant_id", "expense_case_id", "id"),
    )
    _assert_composite_foreign_key(
        engine,
        "ai_decision_feedback",
        ("tenant_id", "decision_id"),
        "ai_decisions",
    )
    _assert_composite_foreign_key(
        engine,
        "workflow_outcomes",
        ("tenant_id", "expense_case_id"),
        "expense_cases",
    )
    _assert_composite_foreign_key(
        engine,
        "workflow_outcomes",
        ("tenant_id", "expense_case_id", "decision_id"),
        "ai_decisions",
        ("tenant_id", "expense_case_id", "id"),
    )
    _assert_composite_foreign_key(
        engine,
        "workflow_outcomes",
        ("tenant_id", "expense_case_id", "business_event_id"),
        "business_events",
        ("tenant_id", "expense_case_id", "id"),
    )


def _assert_runtime_cascade(engine: Engine) -> None:
    with engine.begin() as connection:
        connection.execute(
            text(
                """
                INSERT INTO expense_cases (
                    id, tenant_id, case_no, scene_code, title, current_stage, status
                ) VALUES (
                    'migration-probe-case', 'migration-probe', 'CASE-MIGRATION-PROBE',
                    'reimbursement', '迁移级联验证', 'claiming', 'active'
                )
                """
            )
        )
        connection.execute(
            text(
                """
                INSERT INTO expense_case_links (
                    id, tenant_id, expense_case_id, resource_type, resource_id, relation_type
                ) VALUES (
                    'migration-probe-link', 'migration-probe', 'migration-probe-case',
                    'expense_claim', 'migration-probe-claim', 'claim'
                )
                """
            )
        )
        connection.execute(
            text(
                """
                INSERT INTO business_events (
                    id, tenant_id, expense_case_id, aggregate_type, aggregate_id,
                    event_type, event_version, idempotency_key, correlation_id,
                    actor_id, actor_type, payload_json, delivery_status, delivery_attempts
                ) VALUES (
                    'migration-probe-event', 'migration-probe', 'migration-probe-case',
                    'expense_claim', 'migration-probe-claim', 'claim_draft_created', 1,
                    'migration-probe-idempotency', 'migration-probe-correlation',
                    'migration-probe-user', 'user', '{}', 'pending', 0
                )
                """
            )
        )
        connection.execute(
            text("DELETE FROM expense_cases WHERE id = 'migration-probe-case'")
        )
        assert connection.scalar(
            text("SELECT COUNT(*) FROM expense_case_links WHERE id = 'migration-probe-link'")
        ) == 0
        assert connection.scalar(
            text("SELECT COUNT(*) FROM business_events WHERE id = 'migration-probe-event'")
        ) == 0


def _assert_learning_ledger_tenant_boundary(engine: Engine) -> None:
    with engine.begin() as connection:
        connection.execute(
            text(
                """
                INSERT INTO expense_cases (
                    id, tenant_id, case_no, scene_code, title, current_stage, status
                ) VALUES
                (
                    'learning-probe-case', 'learning-probe', 'CASE-LEARNING-PROBE',
                    'travel', '学习闭环迁移验证', 'claiming', 'active'
                ),
                (
                    'learning-probe-case-b', 'learning-probe', 'CASE-LEARNING-PROBE-B',
                    'travel', '学习闭环同租户第二费用单', 'claiming', 'active'
                )
                """
            )
        )
        connection.execute(
            text(
                """
                INSERT INTO business_events (
                    id, tenant_id, expense_case_id, aggregate_type, aggregate_id,
                    event_type, event_version, idempotency_key, correlation_id,
                    actor_id, actor_type, payload_json, delivery_status, delivery_attempts
                ) VALUES (
                    'learning-probe-event', 'learning-probe', 'learning-probe-case',
                    'expense_claim', 'learning-probe-claim', 'claim_draft_created', 1,
                    'learning-probe-event-key', 'learning-probe-correlation',
                    'learning-probe-user', 'user', '{}', 'pending', 0
                )
                """
            )
        )
        connection.execute(
            text(
                """
                INSERT INTO ai_decisions (
                    id, tenant_id, expense_case_id, business_event_id,
                    expense_claim_id, agent_run_id, correlation_id, subject_type,
                    subject_id, decision_type, decision_source, status,
                    automation_mode, confidence, suggestion_json, evidence_json,
                    version_json, schema_version, idempotency_key, content_fingerprint
                ) VALUES (
                    'learning-probe-decision', 'learning-probe', 'learning-probe-case',
                    'learning-probe-event', 'learning-probe-claim', NULL,
                    'learning-probe-correlation', 'expense_claim', 'learning-probe-claim',
                    'expense_application_prefill', 'hybrid', 'executed', 'prefill', 0.9,
                    '{}', '{}', '{}', 1, 'learning-probe-decision-key',
                    'sha256:aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa'
                )
                """
            )
        )
        connection.execute(
            text(
                """
                INSERT INTO ai_decision_feedback (
                    id, tenant_id, decision_id, expense_claim_id, correlation_id,
                    feedback_type, action_type, actor_id, actor_type, evidence_source,
                    verification_status, final_value_json, changed_fields_json, idempotency_key,
                    content_fingerprint
                ) VALUES (
                    'learning-probe-feedback', 'learning-probe', 'learning-probe-decision',
                    'learning-probe-claim', 'learning-probe-correlation', 'accepted',
                    'save_draft', 'learning-probe-user', 'user',
                    'client_action_confirmation', 'client_observed', '{}', '[]',
                    'learning-probe-feedback-key',
                    'sha256:bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb'
                )
                """
            )
        )
        connection.execute(
            text(
                """
                INSERT INTO workflow_outcomes (
                    id, tenant_id, expense_case_id, decision_id, business_event_id,
                    expense_claim_id, correlation_id, outcome_type, outcome_status,
                    actor_id, actor_type, result_json, idempotency_key,
                    content_fingerprint
                ) VALUES (
                    'learning-probe-outcome', 'learning-probe', 'learning-probe-case',
                    'learning-probe-decision', 'learning-probe-event',
                    'learning-probe-claim', 'learning-probe-correlation', 'draft_saved',
                    'recorded', 'learning-probe-user', 'user', '{}',
                    'learning-probe-outcome-key',
                    'sha256:cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc'
                )
                """
            )
        )

    with pytest.raises(IntegrityError):
        with engine.begin() as connection:
            connection.execute(
                text(
                    """
                    INSERT INTO ai_decisions (
                        id, tenant_id, expense_case_id, expense_claim_id,
                        correlation_id, subject_type, subject_id, decision_type,
                        decision_source, status, automation_mode, suggestion_json,
                        evidence_json, version_json, schema_version, idempotency_key,
                        content_fingerprint
                    ) VALUES (
                        'cross-tenant-decision', 'other-tenant', 'learning-probe-case',
                        'cross-tenant-claim', 'cross-tenant-correlation', 'expense_claim',
                        'cross-tenant-claim', 'expense_application_prefill', 'heuristic',
                        'executed', 'prefill', '{}', '{}', '{}', 1,
                        'cross-tenant-key',
                        'sha256:dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd'
                    )
                    """
                )
            )

    with pytest.raises(IntegrityError):
        with engine.begin() as connection:
            connection.execute(
                text(
                    """
                    INSERT INTO ai_decision_feedback (
                        id, tenant_id, decision_id, expense_claim_id, correlation_id,
                        feedback_type, action_type, actor_id, actor_type, evidence_source,
                        verification_status, training_eligible, final_value_json,
                        changed_fields_json, idempotency_key, content_fingerprint
                    ) VALUES (
                        'unverified-training-feedback', 'learning-probe',
                        'learning-probe-decision', 'learning-probe-claim',
                        'unverified-training-correlation', 'accepted', 'save_draft',
                        'learning-probe-user', 'user', 'client_action_confirmation',
                        'client_observed', TRUE, '{}', '[]',
                        'unverified-training-feedback-key',
                        'sha256:9999999999999999999999999999999999999999999999999999999999999999'
                    )
                    """
                )
            )

    with pytest.raises(IntegrityError):
        with engine.begin() as connection:
            connection.execute(
                text(
                    """
                    INSERT INTO ai_decisions (
                        id, tenant_id, expense_case_id, business_event_id,
                        expense_claim_id, correlation_id, subject_type, subject_id,
                        decision_type, decision_source, status, automation_mode,
                        suggestion_json, evidence_json, version_json, schema_version,
                        idempotency_key, content_fingerprint
                    ) VALUES (
                        'cross-case-event-decision', 'learning-probe',
                        'learning-probe-case-b', 'learning-probe-event',
                        'cross-case-event-claim', 'cross-case-event-correlation',
                        'expense_claim', 'cross-case-event-claim',
                        'expense_application_prefill', 'heuristic', 'executed',
                        'prefill', '{}', '{}', '{}', 1,
                        'cross-case-event-key',
                        'sha256:eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee'
                    )
                    """
                )
            )

    with pytest.raises(IntegrityError):
        with engine.begin() as connection:
            connection.execute(
                text(
                    """
                    INSERT INTO workflow_outcomes (
                        id, tenant_id, expense_case_id, decision_id,
                        expense_claim_id, correlation_id, outcome_type, outcome_status,
                        actor_id, actor_type, result_json, idempotency_key,
                        content_fingerprint
                    ) VALUES (
                        'cross-case-decision-outcome', 'learning-probe',
                        'learning-probe-case-b', 'learning-probe-decision',
                        'cross-case-decision-claim', 'cross-case-decision-correlation',
                        'draft_saved', 'recorded', 'learning-probe-user', 'user', '{}',
                        'cross-case-decision-key',
                        'sha256:ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff'
                    )
                    """
                )
            )


def _create_legacy_sentinel(engine: Engine) -> None:
    with engine.begin() as connection:
        connection.execute(
            text(
                f"""
                CREATE TABLE {LEGACY_PROBE_TABLE} (
                    id VARCHAR(64) PRIMARY KEY,
                    payload VARCHAR(255) NOT NULL
                )
                """
            )
        )
        connection.execute(
            text(
                f"""
                INSERT INTO {LEGACY_PROBE_TABLE} (id, payload)
                VALUES ('legacy-sentinel', 'must-survive-migration-cycle')
                """
            )
        )


def _assert_legacy_sentinel(engine: Engine) -> None:
    assert LEGACY_PROBE_TABLE in _table_names(engine)
    with engine.connect() as connection:
        payload = connection.scalar(
            text(
                f"SELECT payload FROM {LEGACY_PROBE_TABLE} "
                "WHERE id = 'legacy-sentinel'"
            )
        )
    assert payload == "must-survive-migration-cycle"


def _assert_base_schema(engine: Engine) -> None:
    names = _table_names(engine)
    assert not MIGRATION_OWNED_TABLES.intersection(names)
    assert "alembic_version" in names
    with engine.connect() as connection:
        assert connection.scalar(text("SELECT COUNT(*) FROM alembic_version")) == 0


@pytest.mark.parametrize(
    "database_url",
    [
        "sqlite+pysqlite:///:memory:",
        "postgresql+psycopg://postgres:postgres@x-financial-local-postgres:5432/x_financial",
        "postgresql+psycopg://probe:probe@migration-probe-123:5432/x_financial",
    ],
)
def test_disposable_database_guard_rejects_unsafe_urls(database_url: str) -> None:
    with pytest.raises(RuntimeError):
        _require_disposable_probe_url(database_url)


def test_alembic_migration_cycle_on_disposable_postgres(
    migration_database_url: str,
) -> None:
    engine = create_engine(migration_database_url, poolclass=NullPool)
    try:
        assert _table_names(engine) == set(), "迁移测试必须从全新空库开始"
        assert validate_migration_state(engine).revision is None

        _upgrade_head(migration_database_url)
        _assert_head_schema(engine)
        assert validate_migration_state(engine).revision == HEAD_REVISION
        create_legacy_schema(engine)
        assert "expense_claims" in _table_names(engine)

        _upgrade_head(migration_database_url)
        _assert_head_schema(engine)
        _assert_learning_ledger_tenant_boundary(engine)
        _assert_runtime_cascade(engine)

        _create_legacy_sentinel(engine)
        _downgrade_base(migration_database_url)

        _assert_base_schema(engine)
        _assert_legacy_sentinel(engine)
        assert validate_migration_state(engine).revision is None

        with engine.begin() as connection:
            connection.execute(text("CREATE TABLE expense_cases (id VARCHAR(36) PRIMARY KEY)"))
        with pytest.raises(MigrationPreflightError, match="migration-owned tables exist"):
            validate_migration_state(engine)
        with engine.begin() as connection:
            connection.execute(text("DROP TABLE expense_cases"))
        assert validate_migration_state(engine).revision is None

        _upgrade_head(migration_database_url)
        _assert_head_schema(engine)
        _assert_legacy_sentinel(engine)
    finally:
        engine.dispose()