""" Database Base Configuration =========================== CTO-201: Async SQLAlchemy setup Features: - SQLAlchemy 2.0 async engine - PostgreSQL (asyncpg) - 188 PostgreSQL - Session dependency injection 統帥鐵律 2026-03-23: - 絕對禁止 SQLite - 所有 Episodic Memory 必須使用 PostgreSQL """ from collections.abc import AsyncGenerator from contextlib import asynccontextmanager from sqlalchemy import text from sqlalchemy.ext.asyncio import ( AsyncEngine, AsyncSession, async_sessionmaker, create_async_engine, ) from sqlalchemy.orm import DeclarativeBase from src.core.config import settings # ============================================================================= # Base Model # ============================================================================= class Base(DeclarativeBase): """SQLAlchemy declarative base""" pass # ============================================================================= # Engine & Session Factory # ============================================================================= _engine: AsyncEngine | None = None _session_factory: async_sessionmaker[AsyncSession] | None = None def get_engine() -> AsyncEngine: """ Get or create async engine 統帥鐵律 2026-03-23: - 使用 DATABASE_URL (PostgreSQL) - 絕對禁止 SQLite """ global _engine if _engine is None: database_url = settings.DATABASE_URL # 統帥鐵律: 禁止 SQLite (AWOOOI 憲法) # 🔴 違反此規則必須立即報錯,禁止 fallback if "sqlite" in database_url.lower(): import structlog logger = structlog.get_logger(__name__) logger.error("sqlite_forbidden", url=database_url) raise ValueError( "SQLite is FORBIDDEN by AWOOOI Constitution. " "Set DATABASE_URL to PostgreSQL: postgresql+asyncpg://user:pass@host:5432/db" ) _engine = create_async_engine( database_url, echo=settings.DEBUG, pool_size=10, max_overflow=20, pool_pre_ping=True, ) return _engine def get_session_factory() -> async_sessionmaker[AsyncSession]: """Get or create session factory""" global _session_factory if _session_factory is None: _session_factory = async_sessionmaker( bind=get_engine(), class_=AsyncSession, expire_on_commit=False, autoflush=False, ) return _session_factory # ============================================================================= # Dependency Injection # ============================================================================= async def get_db() -> AsyncGenerator[AsyncSession, None]: """ FastAPI dependency for database session Usage: @router.get("/items") async def get_items(db: AsyncSession = Depends(get_db)): ... """ factory = get_session_factory() async with factory() as session: try: yield session await session.commit() except Exception: await session.rollback() raise @asynccontextmanager async def get_db_context() -> AsyncGenerator[AsyncSession, None]: """ Context manager for database session (non-FastAPI usage) Usage: async with get_db_context() as db: ... """ factory = get_session_factory() async with factory() as session: try: yield session await session.commit() except Exception: await session.rollback() raise # ============================================================================= # Initialization # ============================================================================= async def init_db() -> None: """ Initialize database tables Call this at application startup. """ engine = get_engine() async with engine.begin() as conn: await conn.run_sync(Base.metadata.create_all) # 2026-04-02 Claude Code: 確保 risklevel enum 包含 'high' 值 # Phase 23 新增,避免舊 DB 缺少此值導致 InvalidTextRepresentation await conn.execute( text(""" DO $$ BEGIN IF NOT EXISTS ( SELECT 1 FROM pg_enum WHERE enumtypid = 'risklevel'::regtype AND enumlabel = 'high' ) THEN ALTER TYPE risklevel ADD VALUE 'high'; END IF; END $$; """) ) async def close_db() -> None: """ Close database connections Call this at application shutdown. """ global _engine, _session_factory if _engine is not None: await _engine.dispose() _engine = None _session_factory = None