DataOps-platform/tests/test_database_migrations.py

from __future__ import annotations

import os
import subprocess
import uuid
from pathlib import Path

import psycopg2
import pytest
from sqlalchemy import create_engine, inspect
from sqlalchemy.engine import make_url


ROOT = Path(__file__).resolve().parents[1]
EXPECTED_BASELINE_TABLES = {
    "data_orders",
    "data_products",
    "metadata_review_records",
    "metadata_version_history",
    "task_list",
    "users",
}
EXPECTED_UPGRADED_TABLES = {
    "datasource_credentials",
    "datasource_credential_audit_events",
    "workflow_schedules",
    "workflow_runs",
    "workflow_task_runs",
    "workflow_engine_bindings",
    "workflow_plan_audits",
    "runner_task_executions",
    "workflow_canary_evidence",
    "workflow_gateway_operations",
}


def test_alembic_configuration_is_environment_only():
    ini = (ROOT / "alembic.ini").read_text(encoding="utf-8")
    env = (ROOT / "migrations" / "env.py").read_text(encoding="utf-8")

    assert "sqlalchemy.url" not in ini
    assert "SQLALCHEMY_DATABASE_URI" in env
    assert "DATABASE_URL" in env
    assert "password" not in env.lower()


def test_baseline_migration_is_non_destructive():
    baseline = (
        ROOT / "migrations" / "versions" / "20260716_01_baseline.py"
    ).read_text(encoding="utf-8")

    assert "def upgrade" in baseline
    assert "def downgrade" in baseline
    assert "drop_table" not in baseline
    for table in EXPECTED_BASELINE_TABLES:
        assert table in baseline


@pytest.mark.integration
def test_alembic_upgrade_is_repeatable_and_downgrade_preserves_tables():
    admin_url = os.environ.get("TEST_POSTGRES_ADMIN_URL")
    if not admin_url:
        pytest.skip("TEST_POSTGRES_ADMIN_URL is not configured")

    parsed = make_url(admin_url)
    database_name = f"dataops_migration_{uuid.uuid4().hex[:12]}"
    target_url = parsed.set(database=database_name).render_as_string(
        hide_password=False
    )

    connection = psycopg2.connect(admin_url)
    connection.autocommit = True
    try:
        with connection.cursor() as cursor:
            cursor.execute(f'CREATE DATABASE "{database_name}"')

        env = os.environ.copy()
        env["SQLALCHEMY_DATABASE_URI"] = target_url
        command = [str(ROOT / ".venv" / "bin" / "alembic"), "-c", "alembic.ini"]

        subprocess.run(command + ["upgrade", "head"], cwd=ROOT, env=env, check=True)
        subprocess.run(command + ["upgrade", "head"], cwd=ROOT, env=env, check=True)

        engine = create_engine(target_url)
        try:
            tables = set(inspect(engine).get_table_names(schema="public"))
            assert EXPECTED_BASELINE_TABLES | {"alembic_version"} <= tables
            assert EXPECTED_UPGRADED_TABLES <= tables
        finally:
            engine.dispose()

        subprocess.run(command + ["downgrade", "-1"], cwd=ROOT, env=env, check=True)

        engine = create_engine(target_url)
        try:
            tables = set(inspect(engine).get_table_names(schema="public"))
            assert EXPECTED_BASELINE_TABLES <= tables
        finally:
            engine.dispose()
    finally:
        with connection.cursor() as cursor:
            cursor.execute(
                "SELECT pg_terminate_backend(pid) FROM pg_stat_activity "
                "WHERE datname = %s AND pid <> pg_backend_pid()",
                (database_name,),
            )
            cursor.execute(f'DROP DATABASE IF EXISTS "{database_name}"')
        connection.close()