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")