feat: api playground, other tweaks

surfsense_backend/scripts/check_migration_flow.py

@@ -0,0 +1,142 @@
+"""Self-check for the alembic fast-forward/adoption flow in alembic/env.py.
+
+Verifies ``alembic upgrade head`` succeeds on the three DB states it must
+handle without replaying pre-workspace-rename history against a
+workspace-shape schema:
+
+  1. fresh   -- empty database (fast-forward: create_all + stamp head)
+  2. bootstrap -- create_all-created schema + zero_publication, no alembic
+     history (adoption: stamp head)
+  3. midcrash -- bootstrap schema whose alembic_version is stuck at a
+     pre-rename revision from a failed replay (adoption: stamp head)
+
+Run:  python scripts/check_migration_flow.py
+"""
+
+import asyncio
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+BACKEND_DIR = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(BACKEND_DIR))
+
+ADMIN_URL = os.getenv(
+    "ADMIN_DATABASE_URL", "postgresql://postgres:postgres@localhost:5432/postgres"
+)
+SCRATCH_DB = "surfsense_check_migration_flow"
+SCRATCH_URL = ADMIN_URL.rsplit("/", 1)[0] + f"/{SCRATCH_DB}"
+SCRATCH_URL_ASYNC = SCRATCH_URL.replace("postgresql://", "postgresql+asyncpg://")
+
+
+async def recreate_scratch_db() -> None:
+    import asyncpg
+
+    admin = await asyncpg.connect(ADMIN_URL)
+    await admin.execute(f'DROP DATABASE IF EXISTS "{SCRATCH_DB}" WITH (FORCE)')
+    await admin.execute(f'CREATE DATABASE "{SCRATCH_DB}"')
+    await admin.close()
+
+
+async def drop_scratch_db() -> None:
+    import asyncpg
+
+    admin = await asyncpg.connect(ADMIN_URL)
+    await admin.execute(f'DROP DATABASE IF EXISTS "{SCRATCH_DB}" WITH (FORCE)')
+    await admin.close()
+
+
+def run_alembic_upgrade() -> None:
+    env = dict(os.environ, DATABASE_URL=SCRATCH_URL_ASYNC)
+    subprocess.run(
+        [sys.executable, "-m", "alembic", "upgrade", "head"],
+        cwd=BACKEND_DIR,
+        env=env,
+        check=True,
+    )
+
+
+async def bootstrap_schema() -> None:
+    """Mimic app startup bootstrap: create_all + ensure_publication, no stamp."""
+    from sqlalchemy import text
+    from sqlalchemy.ext.asyncio import create_async_engine
+
+    from app.db import Base
+    from app.zero_publication import ensure_publication
+
+    engine = create_async_engine(SCRATCH_URL_ASYNC)
+    async with engine.begin() as conn:
+        await conn.execute(text("CREATE EXTENSION IF NOT EXISTS vector"))
+        await conn.execute(text("CREATE EXTENSION IF NOT EXISTS pg_trgm"))
+        await conn.run_sync(Base.metadata.create_all)
+        await conn.run_sync(ensure_publication)
+    await engine.dispose()
+
+
+async def set_version(version: str | None) -> None:
+    import asyncpg
+
+    conn = await asyncpg.connect(SCRATCH_URL)
+    if version is None:
+        await conn.execute("DROP TABLE IF EXISTS alembic_version")
+    else:
+        await conn.execute(
+            "CREATE TABLE IF NOT EXISTS alembic_version ("
+            "version_num VARCHAR(32) NOT NULL PRIMARY KEY)"
+        )
+        await conn.execute("DELETE FROM alembic_version")
+        await conn.execute(
+            "INSERT INTO alembic_version (version_num) VALUES ($1)", version
+        )
+    await conn.close()
+
+
+async def assert_at_head() -> None:
+    import asyncpg
+    from alembic.script import ScriptDirectory
+
+    head = ScriptDirectory(str(BACKEND_DIR / "alembic")).get_current_head()
+    conn = await asyncpg.connect(SCRATCH_URL)
+    version = await conn.fetchval("SELECT version_num FROM alembic_version")
+    workspaces = await conn.fetchval("SELECT to_regclass('workspaces')")
+    publication = await conn.fetchval(
+        "SELECT 1 FROM pg_publication WHERE pubname = 'zero_publication'"
+    )
+    await conn.close()
+    assert version == head, f"expected version {head}, got {version}"
+    assert workspaces, "workspaces table missing"
+    assert publication, "zero_publication missing"
+
+
+async def main() -> None:
+    try:
+        # 1. Fresh empty DB -> fast-forward.
+        await recreate_scratch_db()
+        run_alembic_upgrade()
+        await assert_at_head()
+        print("OK: fresh DB fast-forwards to head")
+
+        # 2. Bootstrap-created schema, no alembic history -> adoption.
+        await recreate_scratch_db()
+        await bootstrap_schema()
+        await set_version(None)
+        run_alembic_upgrade()
+        await assert_at_head()
+        print("OK: bootstrap-created schema adopted (stamped head)")
+
+        # 3. Bootstrap schema stuck at a pre-rename revision -> adoption.
+        await set_version("4")
+        run_alembic_upgrade()
+        await assert_at_head()
+        print("OK: pre-rename stuck revision adopted (stamped head)")
+
+        # Re-run must be a clean no-op.
+        run_alembic_upgrade()
+        await assert_at_head()
+        print("OK: repeat upgrade is a no-op")
+    finally:
+        await drop_scratch_db()
+
+
+asyncio.run(main())