""" 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