"""Database connection and session management.""" from typing import AsyncGenerator from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, create_async_engine, async_sessionmaker from tradingagents.api.config import settings # Create async engine engine: AsyncEngine = create_async_engine( settings.DATABASE_URL, echo=settings.ENVIRONMENT == "development", future=True, pool_pre_ping=True, ) # Create async session factory AsyncSessionLocal = async_sessionmaker( engine, class_=AsyncSession, expire_on_commit=False, autocommit=False, autoflush=False, ) async def get_db() -> AsyncGenerator[AsyncSession, None]: """ Dependency for getting database session. Yields: AsyncSession: Database session Example: @app.get("/items") async def get_items(db: AsyncSession = Depends(get_db)): result = await db.execute(select(Item)) return result.scalars().all() """ async with AsyncSessionLocal() as session: try: yield session await session.commit() except Exception: await session.rollback() raise finally: await session.close() async def init_db() -> None: """ Initialize database tables. Creates all tables defined in models. Use only for development - use Alembic migrations in production. """ from tradingagents.api.models import Base async with engine.begin() as conn: await conn.run_sync(Base.metadata.create_all) async def close_db() -> None: """Close database connections.""" await engine.dispose()