""" Async Database Session Management Provides async SQLAlchemy engine and session management. """ import os from contextlib import asynccontextmanager from typing import AsyncGenerator, Optional from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, create_async_engine from sqlalchemy.orm import sessionmaker from sqlalchemy.pool import NullPool from backend.core.config import settings from backend.core.exceptions import DatabaseError from backend.logging.logger import get_logger from sqlalchemy import text # Initialize logger logger = get_logger("database_session", component="db") # Global engine instance _engine: Optional[AsyncEngine] = None _session_factory: Optional[sessionmaker] = None def _is_production() -> bool: """Check if running in production mode.""" return os.getenv("AEGISLM_ENV", "development").lower() == "production" def create_database_engine() -> AsyncEngine: """ Create async database engine. Uses asyncpg for PostgreSQL with connection pooling. SQLite is NOT allowed in production - PostgreSQL is required. """ global _engine if _engine is not None: return _engine is_prod = _is_production() try: # Determine if using SQLite is_sqlite = settings.database_url.startswith("sqlite") # PRODUCTION: SQLite is NOT allowed if is_sqlite and is_prod: raise DatabaseError( "CRITICAL: SQLite is not allowed in production. " "Please use PostgreSQL with postgresql+asyncpg:// driver.", details={"database_url": "SQLite URLs are blocked in production"} ) # DEVELOPMENT: Warn about SQLite if is_sqlite: logger.warning( "SQLite is being used. This is only suitable for development. " "For production, use PostgreSQL with postgresql+asyncpg://" ) # Fix SQLite URL for async compatibility (development only) # sqlite:// -> sqlite+aiosqlite:// db_url = settings.database_url if is_sqlite and "+aiosqlite" not in db_url: db_url = db_url.replace("sqlite://", "sqlite+aiosqlite://") # Create async engine with appropriate pool settings # Note: When using NullPool, we can't use pool_size/max_overflow/pool_pre_ping engine = create_async_engine( db_url, echo=False, # Set to True for SQL debugging poolclass=NullPool, # Use NullPool for asyncpg compatibility ) # Enable SQLite foreign keys - defer to first use instead of at init # This avoids event loop issues with TestClient if is_sqlite: # Store engine first, FK will be enabled lazily on first connection _engine = engine logger.warning("SQLite database engine created (development mode only)") return _engine # PostgreSQL engine created _engine = engine logger.info(f"PostgreSQL database engine created: {settings.database_url.split('@')[-1]}") return _engine except DatabaseError: # Re-raise DatabaseError as-is raise except Exception as e: raise DatabaseError( f"Failed to create database engine: {str(e)}", details={"database_url": settings.database_url.split("@")[-1] if "@" in settings.database_url else "unknown"} ) def get_session_factory() -> sessionmaker: """ Get or create session factory. Returns: SQLAlchemy async session maker """ global _session_factory, _engine if _engine is None: _engine = create_database_engine() if _session_factory is None: _session_factory = sessionmaker( _engine, class_=AsyncSession, expire_on_commit=False, ) return _session_factory async def get_db_session() -> AsyncGenerator[AsyncSession, None]: """ Dependency for getting async database session. Yields: Async SQLAlchemy session Usage: @app.get("/items") async def get_items(db: AsyncSession = Depends(get_db_session)): ... """ factory = get_session_factory() async with factory() as session: try: yield session await session.commit() except Exception: await session.rollback() raise finally: await session.close() @asynccontextmanager async def get_db_context() -> AsyncGenerator[AsyncSession, None]: """ Context manager for database sessions. Usage: async with get_db_context() as session: result = await session.execute(...) """ factory = get_session_factory() async with factory() as session: try: yield session await session.commit() except Exception: await session.rollback() raise finally: await session.close() async def init_database() -> None: """ Initialize database tables. Creates all tables if they don't exist. For production, use Alembic migrations instead. """ from backend.db.models import Base engine = create_database_engine() async with engine.begin() as conn: # Create all tables await conn.run_sync(Base.metadata.create_all) # Dispose engine (will be recreated on next use) await engine.dispose() async def close_database() -> None: """ Close database connections and dispose engine. Call this on application shutdown. """ global _engine, _session_factory if _engine is not None: await _engine.dispose() _engine = None _session_factory = None async def check_database_connection() -> bool: """ Check database connectivity. Returns: True if connection successful, False otherwise """ try: engine = create_database_engine() async with engine.connect() as conn: await conn.execute(text("SELECT 1")) return True except Exception: return False