Retry create_all once when another process creates the tables concurrently

src/fastmcp_feedback/instrumentation/storage.py

@@ -7,6 +7,7 @@ Alembic ``env.py`` (``target_metadata``) and pass their existing AsyncEngine to
 
 import asyncio
 import logging
+import random as _random
 import time
 from collections.abc import Sequence
 from datetime import UTC, datetime, timedelta
@@ -31,6 +32,7 @@ from sqlalchemy import (
     select,
 )
 from sqlalchemy.dialects.postgresql import JSONB
+from sqlalchemy.exc import DBAPIError
 from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
 
 from .events import MAX_KEY_CHARS, MAX_KIND_CHARS, EventRecord
@@ -172,10 +174,30 @@ class DatabaseSink:
         self._last_prune: float | None = None
 
     async def ensure_tables(self) -> None:
-        """Create the tables if missing. Only called when create_tables=True."""
+        """Create the tables if missing. Only called when create_tables=True.
+
+        Another process starting against the same empty database can create a
+        table between our existence check and our CREATE, which fails ours
+        (SQLite: "already exists"; Postgres: DuplicateTable or a unique
+        violation in pg_type). So on a database error, wait a moment and run
+        create_all once more; it skips whatever the other process made. A
+        second failure is raised.
+        """
+        try:
+            await self._create_all()
+        except DBAPIError as first:
+            await asyncio.sleep(_random.uniform(0.05, 0.25))
+            await self._create_all()
+            logger.info(
+                "table creation failed once (%s), retry succeeded; another "
+                "process was probably creating the same tables",
+                type(first).__name__,
+            )
+        self._tables_ready = True
+
+    async def _create_all(self) -> None:
         async with self.engine.begin() as conn:
             await conn.run_sync(self.metadata.create_all)
-        self._tables_ready = True
 
     async def _ready(self) -> None:
         if self.create_tables and not self._tables_ready:

tests/fixtures/create_tables_race_worker.py

@@ -0,0 +1,39 @@
+"""One writer process for the cross-process create_tables race test.
+
+Usage: python create_tables_race_worker.py <db url> <trigger file> <events>
+
+Builds a DatabaseSink with create_tables=True, waits until the trigger file
+exists so every worker hits the empty database at the same moment, records
+events, flushes, and prints {"sink_errors": n} as JSON. Prints "ready" once
+imports and setup are done, so the parent can release all workers together.
+"""
+
+import asyncio
+import json
+import os
+import sys
+import time
+
+from fastmcp_feedback.instrumentation import DatabaseSink, InstrumentationMiddleware
+
+
+async def main(url: str, trigger: str, count: int) -> dict:
+    sink = DatabaseSink(url, create_tables=True)
+    mw = InstrumentationMiddleware([sink])
+    print("ready", flush=True)
+    deadline = time.monotonic() + 30
+    while not os.path.exists(trigger):
+        if time.monotonic() > deadline:
+            raise SystemExit("trigger file never appeared")
+        await asyncio.sleep(0.001)
+    for i in range(count):
+        mw.record_event("race.write", key=f"{os.getpid()}-{i}")
+    await mw.flush()
+    errors = mw.dispatcher.sink_errors
+    await mw.aclose()  # also disposes the sink's engine
+    return {"sink_errors": errors}
+
+
+if __name__ == "__main__":
+    url, trigger, count = sys.argv[1], sys.argv[2], int(sys.argv[3])
+    print(json.dumps(asyncio.run(main(url, trigger, count))), flush=True)

tests/unit/test_create_tables_race.py

@@ -0,0 +1,162 @@
+"""create_tables=True when several processes start against the same empty database.
+
+Both processes see the tables missing and both CREATE them; one fails. The
+sink retries create_all once, which skips what the other process made.
+"""
+
+import json
+import subprocess
+import sys
+from pathlib import Path
+
+import pytest
+from sqlalchemy import create_engine, func, select
+from sqlalchemy.exc import OperationalError
+
+from fastmcp_feedback.instrumentation import (
+    DatabaseSink,
+    InstrumentationMiddleware,
+    build_metadata,
+)
+
+EVENTS = build_metadata().tables["ffb_events"]
+WORKER = Path(__file__).parent.parent / "fixtures" / "create_tables_race_worker.py"
+
+
+def _already_exists():
+    return OperationalError(
+        "CREATE TABLE ffb_tool_calls",
+        {},
+        Exception("table ffb_tool_calls already exists"),
+    )
+
+
+def race(url: str, trigger: Path, workers: int = 2, events: int = 20) -> list[dict]:
+    """Start the workers, release them together, return their JSON results."""
+    procs = [
+        subprocess.Popen(
+            [sys.executable, str(WORKER), url, str(trigger), str(events)],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+        )
+        for _ in range(workers)
+    ]
+    try:
+        for p in procs:
+            assert p.stdout.readline().strip() == "ready"
+        trigger.touch()
+        results = []
+        for p in procs:
+            out, err = p.communicate(timeout=30)
+            assert p.returncode == 0, err
+            results.append(json.loads(out.strip().splitlines()[-1]))
+        return results
+    finally:
+        for p in procs:
+            if p.poll() is None:
+                p.kill()
+                p.wait()
+
+
+@pytest.mark.slow
+def test_two_processes_cold_start(tmp_path):
+    workers, events = 2, 20
+    for run in range(5):
+        db = tmp_path / f"shared{run}.db"
+        results = race(
+            f"sqlite+aiosqlite:///{db}", tmp_path / f"go{run}", workers, events
+        )
+        engine = create_engine(f"sqlite:///{db}")
+        try:
+            with engine.connect() as conn:
+                stored = conn.execute(
+                    select(func.count()).select_from(EVENTS)
+                ).scalar_one()
+        finally:
+            engine.dispose()
+        assert [r["sink_errors"] for r in results] == [0] * workers, f"run {run}"
+        assert stored == workers * events, f"run {run}"
+
+
+@pytest.fixture
+async def sink(tmp_path):
+    s = DatabaseSink(f"sqlite+aiosqlite:///{tmp_path}/calls.db", create_tables=True)
+    yield s
+    await s.engine.dispose()
+
+
+class TestEnsureTablesRetry:
+    async def test_first_create_fails_retry_succeeds(self, sink, monkeypatch, caplog):
+        real = sink.metadata.create_all
+        calls = []
+
+        def create_all(bind, **kwargs):
+            calls.append(1)
+            if len(calls) == 1:
+                raise _already_exists()
+            return real(bind, **kwargs)
+
+        monkeypatch.setattr(sink.metadata, "create_all", create_all)
+        mw = InstrumentationMiddleware([sink])
+        caplog.set_level("DEBUG", logger="fastmcp_feedback.instrumentation.storage")
+        assert mw.record_event("a.b", key="k") is True
+        await mw.flush()
+        assert mw.dispatcher.sink_errors == 0
+        assert len(calls) == 2
+        assert sink._tables_ready
+        async with sink.engine.connect() as conn:
+            count = (
+                await conn.execute(select(func.count()).select_from(sink.events))
+            ).scalar_one()
+        assert count == 1
+        retry_logs = [r for r in caplog.records if "retry succeeded" in r.getMessage()]
+        assert retry_logs and all(r.levelname == "INFO" for r in retry_logs)
+        assert not [
+            r
+            for r in caplog.records
+            if r.levelname == "WARNING" and r.name.endswith("storage")
+        ]
+
+    async def test_persistent_failure_surfaces(self, sink, monkeypatch):
+        calls = []
+
+        def create_all(bind, **kwargs):
+            calls.append(1)
+            raise _already_exists()
+
+        monkeypatch.setattr(sink.metadata, "create_all", create_all)
+        mw = InstrumentationMiddleware([sink])
+        assert mw.record_event("a.b") is True
+        await mw.flush()
+        assert mw.dispatcher.sink_errors == 1
+        assert len(calls) == 2
+        assert not sink._tables_ready
+        # The next batch tries again rather than assuming the tables exist.
+        assert mw.record_event("a.c") is True
+        await mw.flush()
+        assert mw.dispatcher.sink_errors == 2
+        assert len(calls) == 4
+        await mw.aclose()
+
+    async def test_non_database_error_not_retried(self, sink, monkeypatch):
+        calls = []
+
+        def create_all(bind, **kwargs):
+            calls.append(1)
+            raise RuntimeError("bug")
+
+        monkeypatch.setattr(sink.metadata, "create_all", create_all)
+        with pytest.raises(RuntimeError):
+            await sink.ensure_tables()
+        assert len(calls) == 1
+
+    async def test_create_tables_false_never_creates(self, tmp_path, monkeypatch):
+        s = DatabaseSink(f"sqlite+aiosqlite:///{tmp_path}/none.db")
+        calls = []
+        monkeypatch.setattr(s.metadata, "create_all", lambda *a, **k: calls.append(1))
+        try:
+            await s._ready()
+        finally:
+            await s.engine.dispose()
+        assert calls == []