Blame for app/security/database.py

from __future__ import annotations
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
from uuid import uuid4

from sqlalchemy import event, text
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker, create_async_engine
from app.security.models import Base


def _async_database_url(database_url: str) -> str:
    url = database_url.strip()
    for prefix in ("postgres://", "postgresql://", "postgresql+psycopg2://", "postgresql+psycopg://"):
        if url.startswith(prefix):
            url = "postgresql+asyncpg://" + url[len(prefix):]
            break
    parts = urlsplit(url)
    if parts.scheme != "postgresql+asyncpg":
        return url
    query = parse_qsl(parts.query, keep_blank_values=True)
    normalized = []
    has_cache_setting = False
    for key, value in query:
        if key == "sslmode":
            if value in {"require", "verify-ca", "verify-full"}:
                normalized.append(("ssl", value))
            continue
        if key == "prepared_statement_cache_size":
            has_cache_setting = True
        normalized.append((key, value))
    if not has_cache_setting:
        normalized.append(("prepared_statement_cache_size", "0"))
    return urlunsplit((parts.scheme, parts.netloc, parts.path, urlencode(normalized), parts.fragment))


class SecurityDatabase:
    def __init__(self, database_url: str) -> None:
        async_url = _async_database_url(database_url)
        options: dict[str, object] = {"pool_pre_ping": True}
        if async_url.startswith("postgresql+asyncpg://"):
            # Unique names avoid PgBouncer transaction-pool collisions even
            # when the client-side prepared-statement cache is disabled.
            options["connect_args"] = {
                "statement_cache_size": 0,
                "prepared_statement_name_func": lambda: "__asyncpg_" + uuid4().hex + "__",
            }
        self.engine: AsyncEngine = create_async_engine(async_url, **options)
        if async_url.startswith("sqlite"):
            event.listen(self.engine.sync_engine, "connect", self._configure_sqlite)
        self.session_factory = async_sessionmaker(self.engine, expire_on_commit=False, class_=AsyncSession)

    @staticmethod
    def _configure_sqlite(dbapi_connection: object, _record: object) -> None:
        cursor = dbapi_connection.cursor()  # type: ignore[attr-defined]
        cursor.execute("PRAGMA foreign_keys=ON")
        cursor.execute("PRAGMA journal_mode=WAL")
        cursor.execute("PRAGMA busy_timeout=5000")
        cursor.close()

    @staticmethod
    async def _migrate_postgres_uuid_columns(connection: object) -> None:
        conn = connection
        result = await conn.execute(text("""
            SELECT table_name, column_name, data_type
            FROM information_schema.columns
            WHERE table_schema = current_schema()
              AND (table_name, column_name) IN (
                ('api_keys', 'id'), ('api_keys', 'rotated_from_id'),
                ('audit_logs', 'api_key_id'), ('rate_limits', 'api_key_id')
              )
        """))
        types = {(r.table_name, r.column_name): r.data_type for r in result.fetchall()}
        if not types:
            return
        await conn.execute(text("ALTER TABLE IF EXISTS api_keys DROP CONSTRAINT IF EXISTS api_keys_rotated_from_id_fkey"))
        await conn.execute(text("ALTER TABLE IF EXISTS audit_logs DROP CONSTRAINT IF EXISTS audit_logs_api_key_id_fkey"))
        await conn.execute(text("ALTER TABLE IF EXISTS rate_limits DROP CONSTRAINT IF EXISTS rate_limits_api_key_id_fkey"))
        for table, column in (
            ("api_keys", "id"), ("api_keys", "rotated_from_id"),
            ("audit_logs", "api_key_id"), ("rate_limits", "api_key_id"),
        ):
            if types.get((table, column)) in {"character varying", "text", "varchar"}:
                await conn.execute(text(
                    f'ALTER TABLE "{table}" ALTER COLUMN "{column}" TYPE UUID '
                    f"""USING NULLIF("{column}"::text, '')::uuid"""
                ))
        if ("api_keys", "id") in types:
            if ("api_keys", "rotated_from_id") in types:
                await conn.execute(text("""
                    ALTER TABLE api_keys ADD CONSTRAINT api_keys_rotated_from_id_fkey
                    FOREIGN KEY (rotated_from_id) REFERENCES api_keys(id) ON DELETE SET NULL
                """))
            if ("audit_logs", "api_key_id") in types:
                await conn.execute(text("""
                    ALTER TABLE audit_logs ADD CONSTRAINT audit_logs_api_key_id_fkey
                    FOREIGN KEY (api_key_id) REFERENCES api_keys(id) ON DELETE SET NULL
                """))
            if ("rate_limits", "api_key_id") in types:
                await conn.execute(text("""
                    ALTER TABLE rate_limits ADD CONSTRAINT rate_limits_api_key_id_fkey
                    FOREIGN KEY (api_key_id) REFERENCES api_keys(id) ON DELETE CASCADE
                """))

    async def initialize(self) -> None:
        async with self.engine.begin() as connection:
            if self.engine.dialect.name == "postgresql":
                await self._migrate_postgres_uuid_columns(connection)
            await connection.run_sync(Base.metadata.create_all)

    async def close(self) -> None:
        await self.engine.dispose()

    @asynccontextmanager
    async def session(self) -> AsyncIterator[AsyncSession]:
        async with self.session_factory() as session:
            yield session