""" Database client for AI Service. [AC-AISVC-11] PostgreSQL database with SQLModel for multi-tenant data isolation. """ import logging from typing import AsyncGenerator from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine from sqlalchemy.pool import NullPool from sqlmodel import SQLModel from app.core.config import get_settings logger = logging.getLogger(__name__) settings = get_settings() engine = create_async_engine( settings.database_url, pool_size=settings.database_pool_size, max_overflow=settings.database_max_overflow, echo=settings.debug, pool_pre_ping=True, ) async_session_maker = async_sessionmaker( engine, class_=AsyncSession, expire_on_commit=False, autocommit=False, autoflush=False, ) async def init_db() -> None: """ [AC-AISVC-11] Initialize database tables. Creates all tables defined in SQLModel metadata. """ async with engine.begin() as conn: await conn.run_sync(SQLModel.metadata.create_all) logger.info("[AC-AISVC-11] Database tables initialized") async def close_db() -> None: """ Close database connections. """ await engine.dispose() logger.info("Database connections closed") async def get_session() -> AsyncGenerator[AsyncSession, None]: """ [AC-AISVC-11] Dependency injection for database session. Ensures proper session lifecycle management. """ async with async_session_maker() as session: try: yield session await session.commit() except Exception: await session.rollback() raise finally: await session.close()