""" Database connection management for the Claude Code Project Tracker. """ import os from typing import AsyncGenerator from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker, AsyncSession from sqlalchemy.pool import StaticPool from app.models.base import Base # Database configuration DATABASE_URL = os.getenv("DATABASE_URL", "sqlite+aiosqlite:///./data/tracker.db") # Create async engine engine = create_async_engine( DATABASE_URL, echo=os.getenv("DEBUG", "false").lower() == "true", # Log SQL queries in debug mode connect_args={"check_same_thread": False} if "sqlite" in DATABASE_URL else {}, poolclass=StaticPool if "sqlite" in DATABASE_URL else None, ) # Create session factory async_session_maker = async_sessionmaker( engine, class_=AsyncSession, expire_on_commit=False ) async def get_db() -> AsyncGenerator[AsyncSession, None]: """ Dependency function to get database session. Used by FastAPI dependency injection to provide database sessions to route handlers. """ async with async_session_maker() as session: try: yield session except Exception: await session.rollback() raise finally: await session.close() def get_engine(): """Get the database engine.""" return engine async def init_database(): """Initialize the database by creating all tables.""" async with engine.begin() as conn: # Import all models to ensure they're registered from app.models import ( Project, Session, Conversation, Activity, WaitingPeriod, GitOperation ) # Create all tables await conn.run_sync(Base.metadata.create_all) print("Database initialized successfully!") async def close_database(): """Close database connections.""" await engine.dispose() print("Database connections closed.")