db: move migrations to job and cap pools

2026-01-22 14:11:44 -03:00 · 2026-01-22 14:11:44 -03:00 · bd8fa1fca5
commit bd8fa1fca5
parent e444a52b3d
5 changed files with 185 additions and 100 deletions
--- a/backend/atlas_portal/app_factory.py
+++ b/backend/atlas_portal/app_factory.py
@ -7,7 +7,6 @@ from flask import Flask, jsonify, send_from_directory
 from flask_cors import CORS
 from werkzeug.middleware.proxy_fix import ProxyFix

-from .db import ensure_schema
 from .routes import access_requests, account, admin_access, ai, auth_config, health, lab, monero


@ -16,8 +15,6 @@ def create_app() -> Flask:
    app.wsgi_app = ProxyFix(app.wsgi_app, x_for=1, x_proto=1, x_host=1, x_port=1)
    CORS(app, resources={r"/api/*": {"origins": "*"}})

-    ensure_schema()
-
    health.register(app)
    auth_config.register(app)
    account.register(app)
--- a/backend/atlas_portal/db.py
+++ b/backend/atlas_portal/db.py
@ -5,121 +5,190 @@ from typing import Any, Iterator

 import psycopg
 from psycopg.rows import dict_row
+from psycopg_pool import ConnectionPool

 from . import settings


+MIGRATION_LOCK_ID = 982731
+_pool: ConnectionPool | None = None
+
+
 def configured() -> bool:
    return bool(settings.PORTAL_DATABASE_URL)


+def _pool_kwargs() -> dict[str, Any]:
+    options = (
+        f"-c lock_timeout={settings.PORTAL_DB_LOCK_TIMEOUT_SEC}s "
+        f"-c statement_timeout={settings.PORTAL_DB_STATEMENT_TIMEOUT_SEC}s "
+        f"-c idle_in_transaction_session_timeout={settings.PORTAL_DB_IDLE_IN_TX_TIMEOUT_SEC}s"
+    )
+    return {
+        "connect_timeout": settings.PORTAL_DB_CONNECT_TIMEOUT_SEC,
+        "application_name": "atlas_portal",
+        "options": options,
+        "row_factory": dict_row,
+    }
+
+
+def _get_pool() -> ConnectionPool:
+    global _pool
+    if _pool is None:
+        if not settings.PORTAL_DATABASE_URL:
+            raise RuntimeError("portal database not configured")
+        _pool = ConnectionPool(
+            conninfo=settings.PORTAL_DATABASE_URL,
+            min_size=settings.PORTAL_DB_POOL_MIN,
+            max_size=settings.PORTAL_DB_POOL_MAX,
+            kwargs=_pool_kwargs(),
+        )
+    return _pool
+
+
@contextmanager
 def connect() -> Iterator[psycopg.Connection[Any]]:
    if not settings.PORTAL_DATABASE_URL:
        raise RuntimeError("portal database not configured")
-    with psycopg.connect(settings.PORTAL_DATABASE_URL, row_factory=dict_row) as conn:
+    with _get_pool().connection() as conn:
+        conn.row_factory = dict_row
        yield conn


-def ensure_schema() -> None:
-    if not settings.PORTAL_DATABASE_URL:
+def _try_advisory_lock(conn: psycopg.Connection[Any], lock_id: int) -> bool:
+    row = conn.execute("SELECT pg_try_advisory_lock(%s)", (lock_id,)).fetchone()
+    return bool(row and row[0])
+
+
+def _release_advisory_lock(conn: psycopg.Connection[Any], lock_id: int) -> None:
+    try:
+        conn.execute("SELECT pg_advisory_unlock(%s)", (lock_id,))
+    except Exception:
+        pass
+
+
+def run_migrations() -> None:
+    if not settings.PORTAL_DATABASE_URL or not settings.PORTAL_RUN_MIGRATIONS:
        return
    with connect() as conn:
-        conn.execute(
-            """
-            CREATE TABLE IF NOT EXISTS access_requests (
-              request_code TEXT PRIMARY KEY,
-              username TEXT NOT NULL,
-              first_name TEXT,
-              last_name TEXT,
-              contact_email TEXT,
-              note TEXT,
-              status TEXT NOT NULL,
-              email_verification_token_hash TEXT,
-              email_verification_sent_at TIMESTAMPTZ,
-              email_verified_at TIMESTAMPTZ,
-              created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
-              decided_at TIMESTAMPTZ,
-              decided_by TEXT,
-              initial_password TEXT,
-              initial_password_revealed_at TIMESTAMPTZ,
-              provision_attempted_at TIMESTAMPTZ
+        try:
+            conn.execute(f"SET lock_timeout = '{settings.PORTAL_DB_LOCK_TIMEOUT_SEC}s'")
+            conn.execute(f"SET statement_timeout = '{settings.PORTAL_DB_STATEMENT_TIMEOUT_SEC}s'")
+        except Exception:
+            pass
+        if not _try_advisory_lock(conn, MIGRATION_LOCK_ID):
+            return
+        try:
+            conn.execute(
+                """
+                CREATE TABLE IF NOT EXISTS access_requests (
+                  request_code TEXT PRIMARY KEY,
+                  username TEXT NOT NULL,
+                  first_name TEXT,
+                  last_name TEXT,
+                  contact_email TEXT,
+                  note TEXT,
+                  status TEXT NOT NULL,
+                  email_verification_token_hash TEXT,
+                  email_verification_sent_at TIMESTAMPTZ,
+                  email_verified_at TIMESTAMPTZ,
+                  created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
+                  decided_at TIMESTAMPTZ,
+                  decided_by TEXT,
+                  initial_password TEXT,
+                  initial_password_revealed_at TIMESTAMPTZ,
+                  provision_attempted_at TIMESTAMPTZ,
+                  welcome_email_sent_at TIMESTAMPTZ,
+                  approval_flags TEXT[],
+                  approval_note TEXT,
+                  denial_note TEXT
+                )
+                """
            )
-            """
-        )
-        conn.execute("ALTER TABLE access_requests ADD COLUMN IF NOT EXISTS initial_password TEXT")
-        conn.execute("ALTER TABLE access_requests ADD COLUMN IF NOT EXISTS initial_password_revealed_at TIMESTAMPTZ")
-        conn.execute("ALTER TABLE access_requests ADD COLUMN IF NOT EXISTS provision_attempted_at TIMESTAMPTZ")
-        conn.execute("ALTER TABLE access_requests ADD COLUMN IF NOT EXISTS email_verification_token_hash TEXT")
-        conn.execute("ALTER TABLE access_requests ADD COLUMN IF NOT EXISTS email_verification_sent_at TIMESTAMPTZ")
-        conn.execute("ALTER TABLE access_requests ADD COLUMN IF NOT EXISTS email_verified_at TIMESTAMPTZ")
-        conn.execute("ALTER TABLE access_requests ADD COLUMN IF NOT EXISTS welcome_email_sent_at TIMESTAMPTZ")
-        conn.execute("ALTER TABLE access_requests ADD COLUMN IF NOT EXISTS first_name TEXT")
-        conn.execute("ALTER TABLE access_requests ADD COLUMN IF NOT EXISTS last_name TEXT")
-        conn.execute("ALTER TABLE access_requests ADD COLUMN IF NOT EXISTS approval_flags TEXT[]")
-        conn.execute("ALTER TABLE access_requests ADD COLUMN IF NOT EXISTS approval_note TEXT")
-        conn.execute("ALTER TABLE access_requests ADD COLUMN IF NOT EXISTS denial_note TEXT")
-        conn.execute(
-            """
-            CREATE TABLE IF NOT EXISTS access_request_tasks (
-              request_code TEXT NOT NULL REFERENCES access_requests(request_code) ON DELETE CASCADE,
-              task TEXT NOT NULL,
-              status TEXT NOT NULL,
-              detail TEXT,
-              updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
-              PRIMARY KEY (request_code, task)
+            conn.execute(
+                """
+                ALTER TABLE access_requests
+                  ADD COLUMN IF NOT EXISTS initial_password TEXT,
+                  ADD COLUMN IF NOT EXISTS initial_password_revealed_at TIMESTAMPTZ,
+                  ADD COLUMN IF NOT EXISTS provision_attempted_at TIMESTAMPTZ,
+                  ADD COLUMN IF NOT EXISTS email_verification_token_hash TEXT,
+                  ADD COLUMN IF NOT EXISTS email_verification_sent_at TIMESTAMPTZ,
+                  ADD COLUMN IF NOT EXISTS email_verified_at TIMESTAMPTZ,
+                  ADD COLUMN IF NOT EXISTS welcome_email_sent_at TIMESTAMPTZ,
+                  ADD COLUMN IF NOT EXISTS first_name TEXT,
+                  ADD COLUMN IF NOT EXISTS last_name TEXT,
+                  ADD COLUMN IF NOT EXISTS approval_flags TEXT[],
+                  ADD COLUMN IF NOT EXISTS approval_note TEXT,
+                  ADD COLUMN IF NOT EXISTS denial_note TEXT
+                """
            )
-            """
-        )
-        conn.execute(
-            """
-            CREATE TABLE IF NOT EXISTS access_request_onboarding_steps (
-              request_code TEXT NOT NULL REFERENCES access_requests(request_code) ON DELETE CASCADE,
-              step TEXT NOT NULL,
-              completed_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
-              PRIMARY KEY (request_code, step)
+            conn.execute(
+                """
+                CREATE TABLE IF NOT EXISTS access_request_tasks (
+                  request_code TEXT NOT NULL REFERENCES access_requests(request_code) ON DELETE CASCADE,
+                  task TEXT NOT NULL,
+                  status TEXT NOT NULL,
+                  detail TEXT,
+                  updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
+                  PRIMARY KEY (request_code, task)
+                )
+                """
            )
-            """
-        )
-        conn.execute(
-            """
-            CREATE TABLE IF NOT EXISTS access_request_onboarding_artifacts (
-              request_code TEXT NOT NULL REFERENCES access_requests(request_code) ON DELETE CASCADE,
-              artifact TEXT NOT NULL,
-              value_hash TEXT NOT NULL,
-              created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
-              PRIMARY KEY (request_code, artifact)
+            conn.execute(
+                """
+                CREATE TABLE IF NOT EXISTS access_request_onboarding_steps (
+                  request_code TEXT NOT NULL REFERENCES access_requests(request_code) ON DELETE CASCADE,
+                  step TEXT NOT NULL,
+                  completed_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
+                  PRIMARY KEY (request_code, step)
+                )
+                """
            )
-            """
-        )
-        conn.execute(
-            """
-            CREATE INDEX IF NOT EXISTS access_requests_status_created_at
-              ON access_requests (status, created_at)
-            """
-        )
-        conn.execute(
-            """
-            CREATE INDEX IF NOT EXISTS access_request_tasks_request_code
-              ON access_request_tasks (request_code)
-            """
-        )
-        conn.execute(
-            """
-            CREATE INDEX IF NOT EXISTS access_request_onboarding_steps_request_code
-              ON access_request_onboarding_steps (request_code)
-            """
-        )
-        conn.execute(
-            """
-            CREATE INDEX IF NOT EXISTS access_request_onboarding_artifacts_request_code
-              ON access_request_onboarding_artifacts (request_code)
-            """
-        )
-        conn.execute(
-            """
-            CREATE UNIQUE INDEX IF NOT EXISTS access_requests_username_pending
-              ON access_requests (username)
-              WHERE status = 'pending'
-            """
-        )
+            conn.execute(
+                """
+                CREATE TABLE IF NOT EXISTS access_request_onboarding_artifacts (
+                  request_code TEXT NOT NULL REFERENCES access_requests(request_code) ON DELETE CASCADE,
+                  artifact TEXT NOT NULL,
+                  value_hash TEXT NOT NULL,
+                  created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
+                  PRIMARY KEY (request_code, artifact)
+                )
+                """
+            )
+            conn.execute(
+                """
+                CREATE INDEX IF NOT EXISTS access_requests_status_created_at
+                  ON access_requests (status, created_at)
+                """
+            )
+            conn.execute(
+                """
+                CREATE INDEX IF NOT EXISTS access_request_tasks_request_code
+                  ON access_request_tasks (request_code)
+                """
+            )
+            conn.execute(
+                """
+                CREATE INDEX IF NOT EXISTS access_request_onboarding_steps_request_code
+                  ON access_request_onboarding_steps (request_code)
+                """
+            )
+            conn.execute(
+                """
+                CREATE INDEX IF NOT EXISTS access_request_onboarding_artifacts_request_code
+                  ON access_request_onboarding_artifacts (request_code)
+                """
+            )
+            conn.execute(
+                """
+                CREATE UNIQUE INDEX IF NOT EXISTS access_requests_username_pending
+                  ON access_requests (username)
+                  WHERE status = 'pending'
+                """
+            )
+        finally:
+            _release_advisory_lock(conn, MIGRATION_LOCK_ID)
+
+
+def ensure_schema() -> None:
+    run_migrations()
--- a/backend/atlas_portal/migrate.py
+++ b/backend/atlas_portal/migrate.py
@ -0,0 +1,11 @@
+from __future__ import annotations
+
+from .db import run_migrations
+
+
+def main() -> None:
+    run_migrations()
+
+
+if __name__ == "__main__":
+    main()
--- a/backend/atlas_portal/settings.py
+++ b/backend/atlas_portal/settings.py
@ -60,6 +60,13 @@ ACCOUNT_ALLOWED_GROUPS = [
 ]

 PORTAL_DATABASE_URL = os.getenv("PORTAL_DATABASE_URL", "").strip()
+PORTAL_DB_POOL_MIN = int(os.getenv("PORTAL_DB_POOL_MIN", "0"))
+PORTAL_DB_POOL_MAX = int(os.getenv("PORTAL_DB_POOL_MAX", "5"))
+PORTAL_DB_CONNECT_TIMEOUT_SEC = int(os.getenv("PORTAL_DB_CONNECT_TIMEOUT_SEC", "5"))
+PORTAL_DB_LOCK_TIMEOUT_SEC = int(os.getenv("PORTAL_DB_LOCK_TIMEOUT_SEC", "5"))
+PORTAL_DB_STATEMENT_TIMEOUT_SEC = int(os.getenv("PORTAL_DB_STATEMENT_TIMEOUT_SEC", "30"))
+PORTAL_DB_IDLE_IN_TX_TIMEOUT_SEC = int(os.getenv("PORTAL_DB_IDLE_IN_TX_TIMEOUT_SEC", "10"))
+PORTAL_RUN_MIGRATIONS = _env_bool("PORTAL_RUN_MIGRATIONS", "false")

 PORTAL_ADMIN_USERS = [u.strip() for u in os.getenv("PORTAL_ADMIN_USERS", "bstein").split(",") if u.strip()]
 PORTAL_ADMIN_GROUPS = [g.strip() for g in os.getenv("PORTAL_ADMIN_GROUPS", "admin").split(",") if g.strip()]
--- a/backend/requirements.txt
+++ b/backend/requirements.txt
@ -4,3 +4,4 @@ gunicorn==21.2.0
 httpx==0.27.2
 PyJWT[crypto]==2.10.1
 psycopg[binary]==3.2.6
+psycopg-pool==3.2.6