From 863ade71a0ca92d54cb2f1e0edc97bc924f581b3 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 26 Oct 2026 16:51:14 +1100 Subject: [PATCH] feat(db): Alembic migrations and registry UNIQUE constraints (#84, #81) #84: add alembic.ini and alembic/ with a baseline (0001) matching the fallback bootstrap schema, using IF NOT EXISTS so existing databases adopt it. bootstrap() now finds alembic.ini and runs upgrade head on start. The version table lives in the uploads schema. #81: migration 0002 adds UNIQUE (file_hash) and UNIQUE (file_name) to csv_files, the constraints dynamic_loader's race handling relied on but which never existed. If duplicates already exist, the upgrade stops and lists them instead of deleting anything. The fallback schema gains the same constraints. Loaders: a dynamic upload that loses the race now returns duplicate_file (previously a generic DB error), and a T&E load that loses it returns a structured error and rolls back its rows instead of raising a 500. Tests (tests/test_migrations.py, integration): upgrade/downgrade on a scratch schema, adoption of a fallback-created schema, refusal on existing duplicates, DB-level uniqueness, and deterministic lost-race tests for both loaders. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_019YPoXzBSj89ZHgw7nfTmbB --- alembic.ini | 9 + alembic/env.py | 40 +++++ alembic/script.py.mako | 21 +++ alembic/versions/0001_baseline.py | 59 +++++++ alembic/versions/0002_registry_unique.py | 53 ++++++ api/db.py | 4 +- api/services/dynamic_loader.py | 13 ++ api/services/te_loader.py | 26 ++- scripts/lint.sh | 5 + tests/test_migrations.py | 207 +++++++++++++++++++++++ 10 files changed, 428 insertions(+), 9 deletions(-) create mode 100644 alembic.ini create mode 100644 alembic/env.py create mode 100644 alembic/script.py.mako create mode 100644 alembic/versions/0001_baseline.py create mode 100644 alembic/versions/0002_registry_unique.py create mode 100644 tests/test_migrations.py diff --git a/alembic.ini b/alembic.ini new file mode 100644 index 0000000..486cc2b --- /dev/null +++ b/alembic.ini @@ -0,0 +1,9 @@ +# Alembic configuration for the API's uploads schema (#84). +# The database URL and target schema come from api/config.py (libpq env vars), +# so there is no sqlalchemy.url here. Run from the repo root: +# alembic upgrade head +# api/db.py bootstrap() runs the same upgrade on every API start. + +[alembic] +script_location = %(here)s/alembic +path_separator = os diff --git a/alembic/env.py b/alembic/env.py new file mode 100644 index 0000000..0b04f5e --- /dev/null +++ b/alembic/env.py @@ -0,0 +1,40 @@ +"""Alembic environment for the uploads schema. + +Connection settings come from api/config.py, the same values the API pool uses. +The version table lives inside the uploads schema, so each schema carries its +own migration state (this matters once there is one schema per tenant). +""" + +import sys +from pathlib import Path + +from alembic import context +from sqlalchemy import create_engine, text +from sqlalchemy.engine import URL +from sqlalchemy.pool import NullPool + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from api.config import settings # noqa: E402 + +if context.is_offline_mode(): + raise RuntimeError("Offline (--sql) migrations are not supported; run against a database.") + +url = URL.create( + "postgresql+psycopg2", + username=settings.PG_USER, + password=settings.PG_PASSWORD or None, + host=settings.PG_HOST, + port=settings.PG_PORT, + database=settings.PG_DATABASE, +) +schema = settings.UPLOADS_SCHEMA + +engine = create_engine(url, poolclass=NullPool) +with engine.connect() as connection: + quoted = connection.dialect.identifier_preparer.quote(schema) + connection.execute(text(f"CREATE SCHEMA IF NOT EXISTS {quoted}")) + connection.commit() + context.configure(connection=connection, version_table_schema=schema) + with context.begin_transaction(): + context.run_migrations() diff --git a/alembic/script.py.mako b/alembic/script.py.mako new file mode 100644 index 0000000..417d68f --- /dev/null +++ b/alembic/script.py.mako @@ -0,0 +1,21 @@ +"""${message} + +Revision ID: ${up_revision} +Revises: ${down_revision | comma,n} +Create Date: ${create_date} +""" + +from alembic import op + +revision = ${repr(up_revision)} +down_revision = ${repr(down_revision)} +branch_labels = None +depends_on = None + + +def upgrade() -> None: + ${upgrades if upgrades else "pass"} + + +def downgrade() -> None: + ${downgrades if downgrades else "pass"} diff --git a/alembic/versions/0001_baseline.py b/alembic/versions/0001_baseline.py new file mode 100644 index 0000000..5380f22 --- /dev/null +++ b/alembic/versions/0001_baseline.py @@ -0,0 +1,59 @@ +"""Baseline: csv_files registry and audit_log, as created by the fallback bootstrap. + +Uses IF NOT EXISTS so databases already created by api/db.py's fallback +bootstrap adopt this revision without changes. + +Revision ID: 0001 +Revises: +Create Date: 2026-10-26 +""" + +from alembic import op + +from api.config import settings + +revision = "0001" +down_revision = None +branch_labels = None +depends_on = None + + +def _schema() -> str: + return op.get_bind().dialect.identifier_preparer.quote(settings.UPLOADS_SCHEMA) + + +def upgrade() -> None: + s = _schema() + op.execute( + f""" + CREATE TABLE IF NOT EXISTS {s}.csv_files ( + id BIGSERIAL PRIMARY KEY, + file_name TEXT NOT NULL, + file_hash TEXT NOT NULL, + table_name TEXT NOT NULL, + mode TEXT NOT NULL, + row_count INTEGER NOT NULL DEFAULT 0, + column_names TEXT[] NOT NULL DEFAULT '{{}}', + created_at TIMESTAMPTZ NOT NULL DEFAULT now() + ) + """ + ) + op.execute( + f""" + CREATE TABLE IF NOT EXISTS {s}.audit_log ( + id BIGSERIAL PRIMARY KEY, + action TEXT NOT NULL, + file_id BIGINT, + file_name TEXT, + table_name TEXT, + mode TEXT, + performed_at TIMESTAMPTZ NOT NULL DEFAULT now() + ) + """ + ) + + +def downgrade() -> None: + s = _schema() + op.execute(f"DROP TABLE IF EXISTS {s}.audit_log") + op.execute(f"DROP TABLE IF EXISTS {s}.csv_files") diff --git a/alembic/versions/0002_registry_unique.py b/alembic/versions/0002_registry_unique.py new file mode 100644 index 0000000..57d0d47 --- /dev/null +++ b/alembic/versions/0002_registry_unique.py @@ -0,0 +1,53 @@ +"""Registry uniqueness: one row per file hash and per file name (#81). + +dynamic_loader.py relies on these constraints to settle concurrent identical +uploads, but they were never created. If a database already holds duplicates +(possible because of that race), the upgrade stops and lists them instead of +deleting anything. + +Note: UNIQUE (file_name) matches today's behaviour (a repeated filename is +refused or replaced). ART decision 2 (versioned datasets) revisits it in M5. + +Revision ID: 0002 +Revises: 0001 +Create Date: 2026-10-26 +""" + +from alembic import op +from sqlalchemy import text + +from api.config import settings + +revision = "0002" +down_revision = "0001" +branch_labels = None +depends_on = None + + +def _schema() -> str: + return op.get_bind().dialect.identifier_preparer.quote(settings.UPLOADS_SCHEMA) + + +def upgrade() -> None: + s = _schema() + bind = op.get_bind() + for column in ("file_hash", "file_name"): + dupes = bind.execute( + # column comes from the fixed tuple above; s is quoted by the dialect. + text(f"SELECT {column}, count(*) FROM {s}.csv_files GROUP BY {column} HAVING count(*) > 1") # nosec B608 + ).fetchall() + if dupes: + listed = ", ".join(f"{value!r} x{count}" for value, count in dupes[:10]) + raise RuntimeError( + f"Cannot add UNIQUE ({column}) to {settings.UPLOADS_SCHEMA}.csv_files: " + f"{len(dupes)} duplicated value(s): {listed}. Remove the duplicate registry " + f"rows (keep the one whose table_name still exists), then restart the API." + ) + op.execute(f"ALTER TABLE {s}.csv_files ADD CONSTRAINT csv_files_file_hash_key UNIQUE (file_hash)") + op.execute(f"ALTER TABLE {s}.csv_files ADD CONSTRAINT csv_files_file_name_key UNIQUE (file_name)") + + +def downgrade() -> None: + s = _schema() + op.execute(f"ALTER TABLE {s}.csv_files DROP CONSTRAINT IF EXISTS csv_files_file_name_key") + op.execute(f"ALTER TABLE {s}.csv_files DROP CONSTRAINT IF EXISTS csv_files_file_hash_key") diff --git a/api/db.py b/api/db.py index e3b04b0..6dcdbed 100644 --- a/api/db.py +++ b/api/db.py @@ -198,8 +198,8 @@ def _bootstrap_fallback_schema() -> None: """ CREATE TABLE IF NOT EXISTS {}.csv_files ( id BIGSERIAL PRIMARY KEY, - file_name TEXT NOT NULL, - file_hash TEXT NOT NULL, + file_name TEXT NOT NULL UNIQUE, + file_hash TEXT NOT NULL UNIQUE, table_name TEXT NOT NULL, mode TEXT NOT NULL, row_count INTEGER NOT NULL DEFAULT 0, diff --git a/api/services/dynamic_loader.py b/api/services/dynamic_loader.py index 1d56d3c..7b1d457 100644 --- a/api/services/dynamic_loader.py +++ b/api/services/dynamic_loader.py @@ -105,6 +105,19 @@ def upload_dynamic( try: return _do_upload(file_name, types, overwrite, logs, schema, rows, file_hash) + except psycopg2.errors.UniqueViolation: + # #81: an identical concurrent upload won a race that _do_upload could + # not settle itself (e.g. both creating csv_ at once). Nothing + # from this request was committed. + _log(logs, "duplicate_check", "Lost race to a concurrent identical upload", "warn") + return { + "status": "duplicate_file", + "reason": "content", + "existingFileName": file_name, + "tableName": f"csv_{file_hash[:16]}", + "existingRowCount": 0, + "logs": logs, + } except psycopg2.Error as exc: # Mirror te_loader's structured-error contract: an unexpected DB error # (e.g. a column type edge case not caught by cast_value) should not diff --git a/api/services/te_loader.py b/api/services/te_loader.py index 1646e7c..d9386f2 100644 --- a/api/services/te_loader.py +++ b/api/services/te_loader.py @@ -22,6 +22,7 @@ import hashlib import time +import psycopg2.errors from psycopg2 import sql from psycopg2.extras import execute_values @@ -290,13 +291,24 @@ def _err(row_number: int, column: str, value: str, reason: str) -> None: # *dynamic*-mode upload of the same file also has its csv_ # table dropped instead of orphaned. _cleanup_prior_registry_entries(cur, file_name, file_hash) - cur.execute( - sql.SQL( - "INSERT INTO {}.csv_files (file_name, file_hash, table_name, mode, row_count, column_names) " - "VALUES (%s, %s, %s, 'te', %s, %s) RETURNING id" - ).format(sql.Identifier(settings.UPLOADS_SCHEMA)), - (file_name, file_hash, f"{settings.TE_SCHEMA}.{target_table}", inserted, columns), - ) + try: + cur.execute( + sql.SQL( + "INSERT INTO {}.csv_files (file_name, file_hash, table_name, mode, row_count, column_names) " + "VALUES (%s, %s, %s, 'te', %s, %s) RETURNING id" + ).format(sql.Identifier(settings.UPLOADS_SCHEMA)), + (file_name, file_hash, f"{settings.TE_SCHEMA}.{target_table}", inserted, columns), + ) + except psycopg2.errors.UniqueViolation: + # #81: a concurrent load of the same file registered first. + # Roll back this load's rows too, so nothing lands twice. + conn.rollback() + _log(logs, "duplicate_check", "Lost race to a concurrent load of the same file", "warn") + return { + "status": "error", + "message": "The same file was loaded concurrently and that load finished first. Nothing was imported by this request.", + "logs": logs, + } file_id = cur.fetchone()[0] conn.commit() diff --git a/scripts/lint.sh b/scripts/lint.sh index d5e6a48..a11bf4e 100644 --- a/scripts/lint.sh +++ b/scripts/lint.sh @@ -12,6 +12,10 @@ PY_SRC=( api/main.py api/ratelimit.py api/services/audit.py + api/db.py + alembic/env.py + alembic/versions/0001_baseline.py + alembic/versions/0002_registry_unique.py ) PY_TESTS=( @@ -23,6 +27,7 @@ PY_TESTS=( tests/test_issue_05_import_summary.py tests/test_api_hardening.py tests/test_evals_tier_e_x.py + tests/test_migrations.py ) # Collect scripts/*.py dynamically diff --git a/tests/test_migrations.py b/tests/test_migrations.py new file mode 100644 index 0000000..61648c3 --- /dev/null +++ b/tests/test_migrations.py @@ -0,0 +1,207 @@ +"""Alembic migrations (#84) and registry uniqueness (#81). + +Integration: needs a live PostgreSQL. Fails (never skips) without one. +""" + +from __future__ import annotations + +import hashlib +import os +import threading +import time +import unittest +import uuid +from pathlib import Path + +import psycopg2 +import psycopg2.errors +import pytest +from alembic import command +from alembic.config import Config +from fastapi.testclient import TestClient + +from api.config import settings +from api.db import Conn +from api.main import app +from api.services.dynamic_loader import upload_dynamic +from api.services.te_loader import upload_te + +_HELP = ("Start PostgreSQL and ensure api/ can reach it (PGHOST/PGPORT/PGUSER/" + "PGPASSWORD/PGDATABASE). See scripts/start-api.ps1.") +_ALEMBIC_INI = str(Path(__file__).resolve().parent.parent / "alembic.ini") + + +def _connect(): + try: + return psycopg2.connect(host=settings.PG_HOST, port=settings.PG_PORT, user=settings.PG_USER, + password=settings.PG_PASSWORD, dbname=settings.PG_DATABASE) + except psycopg2.OperationalError as exc: + raise AssertionError(f"Database unreachable: {exc}\n{_HELP}") from None + + +def _registry_rows(column: str, value: str) -> int: + with Conn() as conn, conn.cursor() as cur: + cur.execute(f'SELECT count(*) FROM "{settings.UPLOADS_SCHEMA}".csv_files WHERE {column} = %s', (value,)) + return cur.fetchone()[0] + + +@pytest.mark.integration +class MigrationsOnScratchSchema(unittest.TestCase): + """Upgrade and downgrade run against a throwaway schema, never the live one.""" + + def setUp(self): + self.original = settings.UPLOADS_SCHEMA + settings.UPLOADS_SCHEMA = f"mig_test_{uuid.uuid4().hex[:10]}" + self.cfg = Config(_ALEMBIC_INI) + self.conn = _connect() + self.conn.autocommit = True + + def tearDown(self): + with self.conn.cursor() as cur: + cur.execute(f'DROP SCHEMA IF EXISTS "{settings.UPLOADS_SCHEMA}" CASCADE') + self.conn.close() + settings.UPLOADS_SCHEMA = self.original + + def _constraints(self) -> set[str]: + with self.conn.cursor() as cur: + cur.execute( + "SELECT conname FROM pg_constraint c JOIN pg_namespace n ON n.oid = c.connamespace " + "WHERE n.nspname = %s AND c.contype = 'u'", (settings.UPLOADS_SCHEMA,)) + return {r[0] for r in cur.fetchall()} + + def _tables(self) -> set[str]: + with self.conn.cursor() as cur: + cur.execute("SELECT table_name FROM information_schema.tables WHERE table_schema = %s", + (settings.UPLOADS_SCHEMA,)) + return {r[0] for r in cur.fetchall()} + + def test_upgrade_head_creates_tables_and_unique_constraints(self): + command.upgrade(self.cfg, "head") + self.assertTrue({"csv_files", "audit_log", "alembic_version"} <= self._tables()) + self.assertEqual(self._constraints(), {"csv_files_file_hash_key", "csv_files_file_name_key"}) + + def test_downgrade_reverses_each_step(self): + command.upgrade(self.cfg, "head") + command.downgrade(self.cfg, "0001") + self.assertEqual(self._constraints(), set()) + self.assertIn("csv_files", self._tables()) + command.downgrade(self.cfg, "base") + self.assertNotIn("csv_files", self._tables()) + + def test_baseline_adopts_a_fallback_created_schema(self): + # A deployed database created by the old fallback bootstrap already has + # the tables. Upgrading must adopt them, not fail on CREATE TABLE. + with self.conn.cursor() as cur: + cur.execute(f'CREATE SCHEMA "{settings.UPLOADS_SCHEMA}"') + cur.execute(f'CREATE TABLE "{settings.UPLOADS_SCHEMA}".csv_files (id BIGSERIAL PRIMARY KEY, ' + "file_name TEXT NOT NULL, file_hash TEXT NOT NULL, table_name TEXT NOT NULL, " + "mode TEXT NOT NULL, row_count INTEGER NOT NULL DEFAULT 0, " + "column_names TEXT[] NOT NULL DEFAULT '{}', created_at TIMESTAMPTZ NOT NULL DEFAULT now())") + command.upgrade(self.cfg, "head") + self.assertIn("csv_files_file_hash_key", self._constraints()) + + def test_existing_duplicates_stop_the_upgrade_with_remediation(self): + command.upgrade(self.cfg, "0001") + with self.conn.cursor() as cur: + for name in ("a.csv", "b.csv"): + cur.execute(f'INSERT INTO "{settings.UPLOADS_SCHEMA}".csv_files ' + "(file_name, file_hash, table_name, mode) VALUES (%s, 'samehash', 't', 'dynamic')", + (name,)) + with self.assertRaisesRegex(RuntimeError, "duplicated value.*samehash"): + command.upgrade(self.cfg, "head") + self.assertEqual(self._constraints(), set()) + + +@pytest.mark.integration +class RegistryUniqueness(unittest.TestCase): + """The live uploads schema, brought to head by the API's own bootstrap.""" + + @classmethod + def setUpClass(cls): + cls.ctx = TestClient(app) + try: + cls.client = cls.ctx.__enter__() + except Exception as exc: # noqa: BLE001 + raise AssertionError(f"Could not start the API: {type(exc).__name__}: {exc}\n{_HELP}") from None + + @classmethod + def tearDownClass(cls): + cls.ctx.__exit__(None, None, None) + + def setUp(self): + self.tag = f"{os.getpid()}_{uuid.uuid4().hex[:8]}" + + def test_bootstrap_brings_the_schema_to_head(self): + with Conn() as conn, conn.cursor() as cur: + cur.execute(f'SELECT version_num FROM "{settings.UPLOADS_SCHEMA}".alembic_version') + self.assertEqual(cur.fetchone()[0], "0002") + + def test_duplicate_hash_is_rejected_by_the_database(self): + insert = (f'INSERT INTO "{settings.UPLOADS_SCHEMA}".csv_files (file_name, file_hash, table_name, mode) ' + "VALUES (%s, %s, 't', 'dynamic')") + with Conn() as conn, conn.cursor() as cur: + cur.execute(insert, (f"a_{self.tag}.csv", f"h_{self.tag}")) + with self.assertRaises(psycopg2.errors.UniqueViolation): + cur.execute(insert, (f"b_{self.tag}.csv", f"h_{self.tag}")) + conn.rollback() + + def _against_uncommitted_competitor(self, fn, name: str, content: str, table: str, mode: str) -> dict: + """Run fn while another transaction holds an uncommitted registry row for + the same file. fn's own duplicate check cannot see that row, so it gets as + far as registering and blocks on the UNIQUE index; committing the + competitor then forces the race outcome deterministically.""" + file_hash = hashlib.sha256(content.encode("utf-8")).hexdigest() + competitor = _connect() + result: dict = {} + try: + with competitor.cursor() as cur: + cur.execute(f'INSERT INTO "{settings.UPLOADS_SCHEMA}".csv_files ' + "(file_name, file_hash, table_name, mode) VALUES (%s, %s, %s, %s)", + (name, file_hash, table, mode)) + worker = threading.Thread(target=lambda: result.update(fn())) + worker.start() + watcher = _connect() + watcher.autocommit = True + try: + for _ in range(200): # up to 10 s for the worker to block on the index + with watcher.cursor() as cur: + cur.execute("SELECT count(*) FROM pg_locks WHERE NOT granted") + if cur.fetchone()[0]: + break + time.sleep(0.05) + else: + raise AssertionError("worker never blocked on the registry's UNIQUE index") + finally: + watcher.close() + competitor.commit() + worker.join(timeout=60) + return result + finally: + with competitor.cursor() as cur: + cur.execute(f'DELETE FROM "{settings.UPLOADS_SCHEMA}".csv_files WHERE file_name = %s', (name,)) + competitor.commit() + competitor.close() + + def test_dynamic_upload_losing_the_race_reports_duplicate_and_commits_nothing(self): + name = f"race_{self.tag}.csv" + content = f"a,b\n1,{self.tag}\n2,{self.tag}\n" + table = f"csv_{hashlib.sha256(content.encode()).hexdigest()[:16]}" + result = self._against_uncommitted_competitor( + lambda: upload_dynamic(name, content, None, False), name, content, table, "dynamic") + self.assertEqual(result["status"], "duplicate_file") + with Conn() as conn, conn.cursor() as cur: + cur.execute("SELECT to_regclass(%s)", (f'"{settings.UPLOADS_SCHEMA}".{table}',)) + self.assertIsNone(cur.fetchone()[0], "the losing upload's table must be rolled back") + + def test_te_load_losing_the_race_reports_error_and_rolls_back_rows(self): + name = f"race_te_{self.tag}.csv" + content = "name,org_type,country\n" + "\n".join( + f"RACE_{self.tag}_{i},government,AU" for i in range(20)) + "\n" + result = self._against_uncommitted_competitor( + lambda: upload_te(name, content, "organisations"), name, content, + f"{settings.TE_SCHEMA}.organisations", "te") + self.assertEqual(result["status"], "error") + with Conn() as conn, conn.cursor() as cur: + cur.execute(f'SELECT count(*) FROM "{settings.TE_SCHEMA}".organisations WHERE name LIKE %s', + (f"RACE_{self.tag}_%",)) + self.assertEqual(cur.fetchone()[0], 0, "the losing load must roll back its rows")