claude-code-tracker/app/database/connection.py

"""
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, ToolCall,
            HookEvent, ToolError, WaitingPeriodNew, PerformanceMetric,
            CodeQualityEvent, WorkflowEvent, LearningEvent,
            EnvironmentEvent, CollaborationEvent, ProjectIntelligence
        )
        
        # Initialize hook relationships
        from app.models.hooks import add_hook_relationships
        add_hook_relationships()
        
        # 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.")