core/database.py

"""Database connection and session management.

[Task]: T010
[From]: specs/001-user-auth/plan.md

[Task]: T004
[From]: specs/004-ai-chatbot/plan.md
"""
from sqlmodel import create_engine, Session
from typing import Generator

from core.config import get_settings

settings = get_settings()

# Create database engine with connection pooling
# Optimized for conversation/message table queries in Phase III
# SQLite doesn't support connection pooling, so we conditionally apply parameters
is_sqlite = settings.database_url.startswith("sqlite:")
is_postgresql = settings.database_url.startswith("postgresql:") or settings.database_url.startswith("postgres://")

if is_sqlite:
    # SQLite configuration (no pooling)
    engine = create_engine(
        settings.database_url,
        echo=settings.environment == "development",  # Log SQL in development
        connect_args={"check_same_thread": False}  # Allow multithreaded access
    )
elif is_postgresql:
    # PostgreSQL configuration with connection pooling
    engine = create_engine(
        settings.database_url,
        echo=settings.environment == "development",  # Log SQL in development
        pool_pre_ping=True,  # Verify connections before using
        pool_size=10,  # Number of connections to maintain
        max_overflow=20,  # Additional connections beyond pool_size
        pool_recycle=3600,  # Recycle connections after 1 hour (prevents stale connections)
        pool_timeout=30,  # Timeout for getting connection from pool
        connect_args={
            "connect_timeout": 10,  # Connection timeout
        }
    )
else:
    # Default configuration for other databases
    engine = create_engine(
        settings.database_url,
        echo=settings.environment == "development",
        pool_pre_ping=True
    )


def get_session() -> Generator[Session, None, None]:
    """Get database session.

    Yields:
        Session: SQLModel database session

    Example:
        ```python
        @app.get("/users")
        def read_users(session: Session = Depends(get_session)):
            users = session.exec(select(User)).all()
            return users
        ```
    """
    with Session(engine) as session:
        yield session


# Alias for compatibility with chat.py
get_db = get_session


def init_db():
    """Initialize database tables.

    Creates all tables defined in SQLModel models.
    Should be called on application startup.
    """
    from sqlmodel import SQLModel
    import models.user  # Import models to register them with SQLModel
    import models.task  # Import task model

    # Phase III: Import conversation and message models
    try:
        import models.conversation
        import models.message
    except ImportError:
        # Models not yet created (Phase 2 pending)
        pass

    SQLModel.metadata.create_all(engine)