claudetools/api/database.py

"""
Database connection and session management for ClaudeTools.

This module provides the database engine configuration, session management,
and FastAPI dependency functions for database access throughout the application.
"""

from typing import Generator

from sqlalchemy import create_engine, event, text
from sqlalchemy.engine import Engine
from sqlalchemy.exc import SQLAlchemyError
from sqlalchemy.orm import Session, sessionmaker
from sqlalchemy.pool import Pool

from api.config import get_settings

# Load settings from environment
settings = get_settings()

# Create database engine with connection pooling
engine = create_engine(
    settings.DATABASE_URL,
    pool_size=settings.DATABASE_POOL_SIZE,
    max_overflow=settings.DATABASE_MAX_OVERFLOW,
    pool_pre_ping=True,
    echo=False,
    pool_recycle=3600,
    connect_args={
        "connect_timeout": 10,
    },
)


@event.listens_for(Pool, "connect")
def set_mysql_pragma(dbapi_connection, connection_record) -> None:
    """
    Set MySQL/MariaDB session variables on new connections.

    This event listener ensures consistent behavior across all database
    connections by setting session-level variables when connections are
    established from the pool.

    Args:
        dbapi_connection: The raw database connection
        connection_record: SQLAlchemy's connection record
    """
    cursor = dbapi_connection.cursor()
    cursor.execute("SET SESSION sql_mode='STRICT_TRANS_TABLES,NO_ZERO_DATE'")
    cursor.execute("SET SESSION time_zone='+00:00'")
    cursor.close()


# Session factory for creating database sessions
SessionLocal = sessionmaker(
    autocommit=False,
    autoflush=False,
    bind=engine,
    expire_on_commit=False,
)


def get_db() -> Generator[Session, None, None]:
    """
    FastAPI dependency that provides a database session.

    This function creates a new database session for each request and ensures
    proper cleanup after the request is complete. It handles both successful
    requests and exceptions, guaranteeing that sessions are always closed.

    Yields:
        Session: A SQLAlchemy database session

    Example:
        ```python
        @app.get("/users")
        def get_users(db: Session = Depends(get_db)):
            return db.query(User).all()
        ```

    Raises:
        SQLAlchemyError: Propagates any database errors after cleanup
    """
    db = SessionLocal()
    try:
        yield db
    except SQLAlchemyError:
        db.rollback()
        raise
    finally:
        db.close()


def init_db() -> None:
    """
    Initialize the database by creating all tables.

    This function should be called during application startup to ensure
    all database tables exist. It uses the Base metadata to create tables
    that don't already exist.

    Note:
        This function uses create_all() which is safe for existing tables
        (it won't recreate them). For production migrations, use Alembic.

    Raises:
        SQLAlchemyError: If there's an error creating database tables
    """
    from api.models.base import Base

    try:
        Base.metadata.create_all(bind=engine)
    except SQLAlchemyError as e:
        raise SQLAlchemyError(f"Failed to initialize database: {str(e)}") from e


def check_db_connection() -> bool:
    """
    Check if the database connection is working.

    This function attempts to execute a simple query to verify that
    the database is accessible and responding to queries.

    Returns:
        bool: True if connection is successful, False otherwise

    Example:
        ```python
        if not check_db_connection():
            logger.error("Database is not accessible")
        ```
    """
    try:
        with engine.connect() as connection:
            connection.execute(text("SELECT 1"))
        return True
    except SQLAlchemyError:
        return False