company-ai-platform/tests/test_schema_reconciliation.py

import json
from pathlib import Path
from types import SimpleNamespace

from alembic import command
from alembic.config import Config
import pytest
from sqlalchemy import create_engine, inspect, text
from sqlalchemy.dialects import postgresql, sqlite
from sqlalchemy.schema import CreateTable

from app.core.config import get_settings
from app.modules.ai_memory.models import AIMemoryEntry
from app.modules.business.models import MarketWatchlist
from app.tools import reconcile_platform_schema as schema_tool


def _database(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
    name: str,
    revision: str,
):
    database_path = tmp_path / f"{name}.db"
    database_url = f"sqlite:///{database_path.as_posix()}"
    monkeypatch.setenv("DATABASE_URL", database_url)
    get_settings.cache_clear()
    config = Config("alembic.ini")
    command.upgrade(config, revision)
    return create_engine(database_url), config


def _remove_revision(engine) -> None:
    with engine.begin() as connection:
        connection.execute(text("DELETE FROM alembic_version"))


def test_remote_drift_allowlist_snapshot_is_explicit() -> None:
    operation_counts: dict[str, int] = {}
    for item in schema_tool.ALLOWED_DRIFT:
        operation_counts[item.operation] = operation_counts.get(item.operation, 0) + 1

    assert operation_counts == {
        "add_column": 61,
        "add_constraint": 2,
        "add_fk": 2,
        "add_index": 47,
        "add_table": 2,
        "modify_default": 8,
        "modify_nullable": 1,
        "remove_constraint": 3,
        "remove_index": 9,
        "remove_table": 1,
    }

    clean_005_only = {
        schema_tool.DriftKey(
            "remove_constraint",
            "ai_memory_entries",
            "uq_ai_memory_owner_fingerprint",
        ),
        schema_tool.DriftKey(
            "remove_constraint",
            "market_watchlists",
            "uq_market_watchlist_owner_symbol",
        ),
        schema_tool.DriftKey(
            "modify_nullable",
            "work_tasks",
            "source_system",
        ),
        *{
            schema_tool.DriftKey("modify_default", table_name, column_name)
            for table_name, columns in schema_tool._MODIFY_DEFAULTS.items()
            for column_name in columns
        },
    }
    assert len(schema_tool.ALLOWED_DRIFT - clean_005_only) == 125


@pytest.mark.parametrize(
    ("table", "constraint_name"),
    [
        (AIMemoryEntry.__table__, "uq_ai_memory_owner_fingerprint"),
        (MarketWatchlist.__table__, "uq_market_watchlist_owner_symbol"),
    ],
)
def test_owner_unique_constraints_use_postgres_nulls_not_distinct(
    table,
    constraint_name: str,
) -> None:
    constraint = next(
        item for item in table.constraints if item.name == constraint_name
    )

    assert (
        constraint.dialect_options["postgresql"]["nulls_not_distinct"] is True
    )
    postgres_ddl = str(CreateTable(table).compile(dialect=postgresql.dialect()))
    sqlite_ddl = str(CreateTable(table).compile(dialect=sqlite.dialect()))
    assert "UNIQUE NULLS NOT DISTINCT" in postgres_ddl
    assert "NULLS NOT DISTINCT" not in sqlite_ddl


def test_clean_005_upgrades_to_head_with_zero_metadata_drift(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    engine, config = _database(
        tmp_path,
        monkeypatch,
        "clean-upgrade",
        schema_tool.PREVIOUS_REVISION,
    )
    try:
        command.upgrade(config, "head")
        with engine.connect() as connection:
            audit = schema_tool.audit_connection(connection)

        assert audit.revision_rows == (schema_tool.TARGET_REVISION,)
        assert audit.observed_drift == frozenset()
        assert "approval_requests" not in inspect(engine).get_table_names()
        command.check(config)
    finally:
        engine.dispose()
        get_settings.cache_clear()


def test_baseline_stamps_and_upgrades_in_one_external_transaction(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    engine, _ = _database(
        tmp_path,
        monkeypatch,
        "baseline-success",
        schema_tool.PREVIOUS_REVISION,
    )
    try:
        with engine.begin() as connection:
            connection.execute(
                text("CREATE TABLE approval_requests (id INTEGER PRIMARY KEY)")
            )
        _remove_revision(engine)
        before = schema_tool.audit_engine(engine)

        assert before.eligible is True
        assert schema_tool.DriftKey(
            "remove_table",
            "approval_requests",
            "",
        ) in before.observed_drift

        after = schema_tool.apply_baseline(
            engine,
            before.fingerprint,
            require_postgresql=False,
        )

        assert after.revision_rows == (schema_tool.TARGET_REVISION,)
        assert after.observed_drift == frozenset()
        assert "approval_requests" not in inspect(engine).get_table_names()
    finally:
        engine.dispose()
        get_settings.cache_clear()


def test_baseline_failure_rolls_back_atomic_stamp(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    engine, _ = _database(
        tmp_path,
        monkeypatch,
        "baseline-rollback",
        schema_tool.PREVIOUS_REVISION,
    )
    try:
        _remove_revision(engine)
        before = schema_tool.audit_engine(engine)

        def fail_upgrade(*_args, **_kwargs) -> None:
            raise RuntimeError("injected migration failure")

        monkeypatch.setattr(schema_tool.command, "upgrade", fail_upgrade)
        with pytest.raises(RuntimeError, match="injected migration failure"):
            schema_tool.apply_baseline(
                engine,
                before.fingerprint,
                require_postgresql=False,
            )

        with engine.connect() as connection:
            revisions = tuple(
                connection.execute(
                    text("SELECT version_num FROM alembic_version")
                ).scalars()
            )
        assert revisions == ()
    finally:
        engine.dispose()
        get_settings.cache_clear()


def test_baseline_rejects_changed_fingerprint_before_stamp(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    engine, _ = _database(
        tmp_path,
        monkeypatch,
        "baseline-fingerprint",
        schema_tool.PREVIOUS_REVISION,
    )
    try:
        _remove_revision(engine)
        with pytest.raises(schema_tool.ReconciliationError, match="fingerprint changed"):
            schema_tool.apply_baseline(
                engine,
                "0" * 64,
                require_postgresql=False,
            )
        with engine.connect() as connection:
            assert schema_tool._revision_rows(connection) == ()
    finally:
        engine.dispose()
        get_settings.cache_clear()


def test_preflight_fingerprint_includes_full_server_default_expression() -> None:
    engine = create_engine("sqlite://")
    try:
        with engine.begin() as connection:
            connection.execute(
                text(
                    """
                    CREATE TABLE projects (
                        id INTEGER PRIMARY KEY,
                        lifecycle_state TEXT DEFAULT 'draft'
                    )
                    """
                )
            )
        with engine.connect() as connection:
            first = schema_tool._schema_fingerprint(connection)

        with engine.begin() as connection:
            connection.execute(text("DROP TABLE projects"))
            connection.execute(
                text(
                    """
                    CREATE TABLE projects (
                        id INTEGER PRIMARY KEY,
                        lifecycle_state TEXT DEFAULT 'active'
                    )
                    """
                )
            )
        with engine.connect() as connection:
            second = schema_tool._schema_fingerprint(connection)

        assert first != second
    finally:
        engine.dispose()


def test_preflight_fingerprint_binds_impacted_row_counts_without_data() -> None:
    engine = create_engine("sqlite://")
    private_value = "must-not-appear-in-reconciliation-output"
    try:
        with engine.begin() as connection:
            connection.execute(
                text(
                    """
                    CREATE TABLE approval_requests (
                        id INTEGER PRIMARY KEY,
                        private_payload TEXT
                    )
                    """
                )
            )
        before = schema_tool.audit_engine(engine)

        with engine.begin() as connection:
            connection.execute(
                text(
                    """
                    INSERT INTO approval_requests (id, private_payload)
                    VALUES (1, :private_payload)
                    """
                ),
                {"private_payload": private_value},
            )
        after = schema_tool.audit_engine(engine)

        assert before.fingerprint != after.fingerprint
        assert dict(before.impacted_table_row_counts) == {"approval_requests": 0}
        assert dict(after.impacted_table_row_counts) == {"approval_requests": 1}
        public_payload = after.public_dict()
        assert public_payload["impacted_table_row_counts"] == {
            "approval_requests": 1
        }
        assert "impacted_table_content_digests" not in public_payload
        assert private_value not in json.dumps(public_payload)
    finally:
        engine.dispose()


def test_content_fingerprint_is_order_stable_and_detects_same_count_change() -> None:
    engine = create_engine("sqlite://")
    try:
        with engine.begin() as connection:
            connection.execute(
                text(
                    """
                    CREATE TABLE projects (
                        project_code TEXT,
                        private_payload TEXT
                    )
                    """
                )
            )
            connection.execute(
                text(
                    """
                    INSERT INTO projects (project_code, private_payload)
                    VALUES ('P-1', 'first'), ('P-2', 'second')
                    """
                )
            )
        with engine.connect() as connection:
            first = schema_tool._schema_fingerprint(connection)

        with engine.begin() as connection:
            connection.execute(text("DELETE FROM projects"))
            connection.execute(
                text(
                    """
                    INSERT INTO projects (project_code, private_payload)
                    VALUES ('P-2', 'second'), ('P-1', 'first')
                    """
                )
            )
        with engine.connect() as connection:
            reordered = schema_tool._schema_fingerprint(connection)

        with engine.begin() as connection:
            connection.execute(
                text(
                    """
                    UPDATE projects
                    SET private_payload = 'changed'
                    WHERE project_code = 'P-1'
                    """
                )
            )
        with engine.connect() as connection:
            changed = schema_tool._schema_fingerprint(connection)

        assert reordered == first
        assert changed != first
    finally:
        engine.dispose()


def test_apply_rejects_same_row_count_content_change_after_dry_run(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    engine, _ = _database(
        tmp_path,
        monkeypatch,
        "baseline-content-fingerprint",
        schema_tool.PREVIOUS_REVISION,
    )
    try:
        with engine.begin() as connection:
            connection.execute(
                text(
                    """
                    INSERT INTO feishu_event_receipts (
                        event_key,
                        source,
                        event_id,
                        message_id,
                        received_at
                    )
                    VALUES (
                        'legacy-content-key',
                        'webhook',
                        'event-before',
                        'message-stable',
                        '2026-07-27 00:00:00'
                    )
                    """
                )
            )
        _remove_revision(engine)
        before = schema_tool.audit_engine(engine)

        with engine.begin() as connection:
            connection.execute(
                text(
                    """
                    UPDATE feishu_event_receipts
                    SET event_id = 'event-after'
                    WHERE event_key = 'legacy-content-key'
                    """
                )
            )
        changed = schema_tool.audit_engine(engine)

        assert changed.impacted_table_row_counts == before.impacted_table_row_counts
        assert (
            changed.impacted_table_content_digests
            != before.impacted_table_content_digests
        )
        assert changed.fingerprint != before.fingerprint
        with pytest.raises(schema_tool.ReconciliationError, match="fingerprint changed"):
            schema_tool.apply_baseline(
                engine,
                before.fingerprint,
                require_postgresql=False,
            )
        with engine.connect() as connection:
            assert schema_tool._revision_rows(connection) == ()
    finally:
        engine.dispose()
        get_settings.cache_clear()


@pytest.mark.parametrize(
    ("database_url", "legacy_database_url"),
    [
        ("sqlite:///local-platform.db", None),
        (
            "postgresql://platform@database/platform",
            "postgresql+psycopg://legacy@database:5432/platform",
        ),
    ],
)
def test_cli_dry_run_rejects_unsafe_target_before_engine_creation(
    database_url: str,
    legacy_database_url: str | None,
    monkeypatch: pytest.MonkeyPatch,
    capsys: pytest.CaptureFixture[str],
) -> None:
    engine_creation_count = 0

    def record_engine_creation(*_args, **_kwargs):
        nonlocal engine_creation_count
        engine_creation_count += 1
        raise AssertionError("unsafe target reached create_engine")

    monkeypatch.setattr(
        schema_tool,
        "get_settings",
        lambda: SimpleNamespace(
            database_url=database_url,
            legacy_database_url=legacy_database_url,
        ),
    )
    monkeypatch.setattr(schema_tool, "create_engine", record_engine_creation)
    monkeypatch.setattr("sys.argv", ["reconcile_platform_schema"])

    with pytest.raises(SystemExit) as exit_info:
        schema_tool.main()

    assert exit_info.value.code == 2
    assert engine_creation_count == 0
    output = json.loads(capsys.readouterr().out)
    assert output["applied"] is False
    assert "error" in output
    assert database_url not in output["error"]
    assert legacy_database_url is None or legacy_database_url not in output["error"]


def test_postgres_below_15_is_ineligible_and_apply_precondition_rejects() -> None:
    audit = schema_tool.BaselineAudit(
        dialect="postgresql",
        fingerprint="a" * 64,
        impacted_table_row_counts=(),
        impacted_table_content_digests=(),
        postgresql_server_version_num=140012,
        postgresql_version_supported=False,
        revision_rows=(),
        observed_drift=frozenset(),
        unexpected_drift=frozenset(),
        approval_rows=None,
        approval_inbound_foreign_keys=0,
        schema_privileges_ok=True,
        table_ownership_ok=True,
    )

    assert audit.eligible is False
    assert audit.public_dict()["postgresql_server_version_num"] == 140012
    assert audit.public_dict()["postgresql_version_supported"] is False
    with pytest.raises(schema_tool.ReconciliationError, match="PostgreSQL 15"):
        schema_tool._validate_apply_preconditions(audit, audit.fingerprint)


def test_nonempty_approval_table_is_never_removed(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    engine, _ = _database(
        tmp_path,
        monkeypatch,
        "baseline-nonempty-approval",
        schema_tool.PREVIOUS_REVISION,
    )
    try:
        with engine.begin() as connection:
            connection.execute(
                text("CREATE TABLE approval_requests (id INTEGER PRIMARY KEY)")
            )
            connection.execute(text("INSERT INTO approval_requests (id) VALUES (1)"))
        _remove_revision(engine)
        before = schema_tool.audit_engine(engine)

        assert before.eligible is False
        assert before.approval_rows == 1
        with pytest.raises(schema_tool.ReconciliationError, match="not empty"):
            schema_tool.apply_baseline(
                engine,
                before.fingerprint,
                require_postgresql=False,
            )
        with engine.connect() as connection:
            assert connection.scalar(
                text("SELECT COUNT(*) FROM approval_requests")
            ) == 1
            assert schema_tool._revision_rows(connection) == ()
    finally:
        engine.dispose()
        get_settings.cache_clear()


def test_reconciliation_downgrade_is_fail_closed_and_irreversible(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    engine, config = _database(
        tmp_path,
        monkeypatch,
        "irreversible-downgrade",
        "202607270001",
    )
    try:
        with engine.connect() as connection:
            before = schema_tool._schema_fingerprint(connection)

        with pytest.raises(RuntimeError, match="irreversible"):
            command.downgrade(config, schema_tool.PREVIOUS_REVISION)

        with engine.connect() as connection:
            after = schema_tool._schema_fingerprint(connection)
            revisions = schema_tool._revision_rows(connection)
        assert after == before
        assert revisions == ("202607270001",)
    finally:
        engine.dispose()
        get_settings.cache_clear()


def test_reconciliation_backfills_required_legacy_values(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    engine, config = _database(
        tmp_path,
        monkeypatch,
        "required-backfill",
        schema_tool.PREVIOUS_REVISION,
    )
    try:
        with engine.begin() as connection:
            connection.execute(
                text("ALTER TABLE work_tasks DROP COLUMN source_system")
            )
            connection.execute(
                text(
                    """
                    INSERT INTO work_tasks (
                        code,
                        title,
                        status,
                        priority,
                        created_at,
                        updated_at
                    )
                    VALUES (
                        'legacy-task',
                        'Legacy task',
                        'todo',
                        'P2',
                        '2026-07-27 00:00:00',
                        '2026-07-27 00:00:00'
                    )
                    """
                )
            )

        command.upgrade(config, "head")

        with engine.connect() as connection:
            source_system = connection.scalar(
                text(
                    """
                    SELECT source_system
                    FROM work_tasks
                    WHERE code = 'legacy-task'
                    """
                )
            )
        columns = {
            column["name"]: column
            for column in inspect(engine).get_columns("work_tasks")
        }
        assert source_system == "internal"
        assert columns["source_system"]["nullable"] is False
        assert columns["source_system"]["default"] is None
        assert columns["is_active"]["default"] is None
    finally:
        engine.dispose()
        get_settings.cache_clear()