from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine from sqlalchemy.orm import sessionmaker, declarative_base from config import settings DATABASE_URL = settings.database_url if not DATABASE_URL: raise ValueError( "DATABASE_URL not found in environment variables. " "Please set it in .env file. Example: " "postgresql+asyncpg://user:password@localhost:5432/kairo_db" ) # Create async engine engine = create_async_engine( DATABASE_URL, echo=False, # Set to True for SQL logging future=True, pool_pre_ping=True, # Verify connections before using pool_size=10, max_overflow=20, ) # Create async session factory async_session = sessionmaker( engine, class_=AsyncSession, expire_on_commit=False, autoflush=False ) # Base class for all models Base = declarative_base() async def get_db() -> AsyncSession: """Dependency for getting async database session""" async with async_session() as session: try: yield session finally: await session.close() async def init_db(): """Initialize database by creating all tables""" async with engine.begin() as conn: await conn.run_sync(Base.metadata.create_all) async def close_db(): """Close database connections""" await engine.dispose()