bstein-dev-home/backend/atlas_portal/db.py

from __future__ import annotations

from contextlib import contextmanager
from typing import Any, Iterator

import psycopg
from psycopg.rows import dict_row

from . import settings


def configured() -> bool:
    return bool(settings.PORTAL_DATABASE_URL)


@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:
        yield conn


def ensure_schema() -> None:
    if not settings.PORTAL_DATABASE_URL:
        return
    with connect() as conn:
        conn.execute(
            """
            CREATE TABLE IF NOT EXISTS access_requests (
              request_code TEXT PRIMARY KEY,
              username TEXT NOT NULL,
              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
            )
            """
        )
        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 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(
            """
            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_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'
            """
        )