diff --git a/hindsight-api/hindsight_api/migrations.py b/hindsight-api/hindsight_api/migrations.py index 45245087..7bc4e4bc 100644 --- a/hindsight-api/hindsight_api/migrations.py +++ b/hindsight-api/hindsight_api/migrations.py @@ -18,6 +18,7 @@ No alembic.ini required - all configuration is done programmatically. import hashlib import logging import os +import threading import time from pathlib import Path @@ -33,6 +34,13 @@ logger = logging.getLogger(__name__) # Advisory lock ID for migrations (arbitrary unique number) MIGRATION_LOCK_ID = 123456789 +# Alembic's command.upgrade() is NOT thread-safe: it uses module-level global +# proxies (context._proxy, script) that get overwritten when two threads call +# upgrade() concurrently. This causes migrations to target the wrong schema +# and crash with "relation already exists" or KeyError: 'script'. +# Serialize all Alembic invocations with a process-level lock. +_alembic_lock = threading.Lock() + def _detect_vector_extension(conn, vector_extension: str = "pgvector") -> str: """ @@ -144,9 +152,12 @@ def _run_migrations_internal(database_url: str, script_location: str, schema: st if schema: alembic_cfg.set_main_option("target_schema", schema) - # Run migrations + # Run migrations under a process-level lock. Alembic uses module-level + # global proxies that are not thread-safe, so concurrent command.upgrade() + # calls from different threads corrupt each other's context. try: - command.upgrade(alembic_cfg, "head") + with _alembic_lock: + command.upgrade(alembic_cfg, "head") except ResolutionError as e: # This happens during rolling deployments when a newer version of the code # has already run migrations, and this older replica doesn't have the new diff --git a/hindsight-api/tests/test_migrations_thread_safety.py b/hindsight-api/tests/test_migrations_thread_safety.py new file mode 100644 index 00000000..da73336d --- /dev/null +++ b/hindsight-api/tests/test_migrations_thread_safety.py @@ -0,0 +1,48 @@ +import threading +import time + +from hindsight_api import migrations + + +def test_run_migrations_internal_serializes_alembic_upgrade(monkeypatch): + max_concurrent_upgrades = 0 + active_upgrades = 0 + active_lock = threading.Lock() + start_barrier = threading.Barrier(2) + + def fake_upgrade(_cfg, _revision): + nonlocal max_concurrent_upgrades, active_upgrades + with active_lock: + active_upgrades += 1 + max_concurrent_upgrades = max(max_concurrent_upgrades, active_upgrades) + time.sleep(0.05) + with active_lock: + active_upgrades -= 1 + + monkeypatch.setattr(migrations.command, "upgrade", fake_upgrade) + + errors = [] + + def run_in_thread(schema): + try: + start_barrier.wait() + migrations._run_migrations_internal( + "postgresql://user:pass@localhost/db", + "/tmp/alembic", + schema=schema, + ) + except Exception as exc: # pragma: no cover - diagnostic path + errors.append(exc) + + threads = [ + threading.Thread(target=run_in_thread, args=("tenant_alpha",)), + threading.Thread(target=run_in_thread, args=("tenant_beta",)), + ] + + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + assert not errors + assert max_concurrent_upgrades == 1