Spaces:
Running
Running
| 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) | |
| 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() | |
| 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() | |
| async def session(self) -> AsyncIterator[AsyncSession]: | |
| async with self.session_factory() as session: | |
| yield session | |