web/backend/services/main/db.py

import os
from typing import Generator, Optional

import time
import logging
from sqlalchemy import create_engine, event, text
from sqlalchemy.engine import Engine
from sqlalchemy.orm import sessionmaker, Session
from sqlalchemy.pool import StaticPool
from sqlalchemy.exc import OperationalError

from .constants import DATABASE_URL

logger = logging.getLogger(__name__)

"""
Universal database interface that provides identical behavior for PostgreSQL and SQLite.

Features:
- Auto-creates parent directory for SQLite files.
- Applies SQLite pragmas (foreign_keys=ON, journal_mode=WAL) for improved compatibility.
- Uses StaticPool for in-memory or file-based SQLite when appropriate.
- Exposes `engine`, `SessionLocal`, `get_db` dependency, and `POOL_CONFIG`.
"""

# Ensure parent directory exists for SQLite file DBs
def _ensure_sqlite_parent_dir(url: str) -> None:
    if not url or not url.startswith("sqlite"):
        return
    # strip sqlite:/// prefix
    path = url.replace("sqlite:///", "", 1)
    parent = os.path.dirname(path)
    if parent:
        os.makedirs(parent, exist_ok=True)


# Pool and engine configuration (tunable via env)
POOL_SIZE = int(os.getenv("DB_POOL_SIZE", "20"))
MAX_OVERFLOW = int(os.getenv("DB_MAX_OVERFLOW", "40"))
POOL_RECYCLE = int(os.getenv("DB_POOL_RECYCLE", "1800"))
POOL_TIMEOUT = int(os.getenv("DB_POOL_TIMEOUT", "30"))

POOL_CONFIG = {
    "pool_size": POOL_SIZE,
    "max_overflow": MAX_OVERFLOW,
    "pool_recycle": POOL_RECYCLE,
    "pool_timeout": POOL_TIMEOUT,
    "pool_pre_ping": True,
}


def get_engine(database_url: Optional[str] = None) -> Engine:
    """
    Create and return a SQLAlchemy Engine configured for the given database URL.
    This function ensures SQLite-specific pragmas and connection args are applied.
    Includes retry logic for database connection failures during startup.
    """
    url = database_url or DATABASE_URL
    _ensure_sqlite_parent_dir(url)

    # Retry database connection during startup (helps with Docker initialization timing)
    if url.startswith("postgresql"):
        max_retries = 15
        retry_delay = 2

        for attempt in range(max_retries):
            try:
                logger.info(f"Attempting database connection (attempt {attempt + 1}/{max_retries})...")
                # Test the connection by creating engine and trying to connect
                test_engine = create_engine(url, pool_size=1, max_overflow=0, pool_timeout=5, future=True)
                with test_engine.connect() as conn:
                    conn.execute(text("SELECT 1"))
                test_engine.dispose()
                logger.info("Database connection successful")
                break
            except OperationalError as e:
                if attempt < max_retries - 1:
                    logger.warning(f"Database connection failed (attempt {attempt + 1}): {e}")
                    time.sleep(retry_delay)
                else:
                    logger.error(f"Database connection failed after {max_retries} attempts: {e}")
                    raise
            except Exception as e:
                logger.error(f"Unexpected error during database connection: {e}")
                raise

    if url.startswith("sqlite"):
        # For SQLite file-based DBs, use standard pooling but set connection timeout and pragmas.
        # Use StaticPool only for in-memory SQLite.
        in_memory = url in ("sqlite:///:memory:", "sqlite://")
        connect_args = {"check_same_thread": False, "timeout": int(os.getenv("SQLITE_BUSY_TIMEOUT", "5"))}

        if in_memory:
            engine = create_engine(url, connect_args=connect_args, poolclass=StaticPool, future=True)
        else:
            engine = create_engine(url, connect_args=connect_args, future=True)

        # Apply pragmas on connect for SQLite (foreign keys, WAL, busy_timeout)
        @event.listens_for(engine, "connect")
        def _sqlite_on_connect(dbapi_conn, connection_record):
            try:
                cursor = dbapi_conn.cursor()
                cursor.execute("PRAGMA foreign_keys = ON")
                cursor.execute("PRAGMA journal_mode = WAL")
                # busy_timeout in milliseconds
                busy_ms = int(os.getenv("SQLITE_BUSY_TIMEOUT_MS", "5000"))
                cursor.execute(f"PRAGMA busy_timeout = {busy_ms}")
                cursor.close()
            except Exception:
                # Best-effort; do not fail engine creation if pragmas cannot be set
                pass

        return engine

    # Default for Postgres / MySQL etc. - use pool sizing from env
    engine_kwargs = {
        "pool_size": POOL_SIZE,
        "max_overflow": MAX_OVERFLOW,
        "pool_recycle": POOL_RECYCLE,
        "pool_timeout": POOL_TIMEOUT,
        "future": True,
    }
    return create_engine(url, **engine_kwargs)


# Create global engine and session factory for convenient imports
engine = get_engine()
# Keep loaded attributes available after commit/close to avoid DetachedInstanceError
SessionLocal = sessionmaker(class_=Session, autocommit=False, autoflush=False, bind=engine, expire_on_commit=False)


def init_db(create_tables: bool = False, base_metadata=None) -> None:
    """
    Initialize the database. If `create_tables` is True and `base_metadata` is provided,
    create all tables using the provided SQLAlchemy metadata.
    """
    if create_tables:
        if base_metadata is None:
            raise ValueError("base_metadata is required to create tables")
        base_metadata.create_all(bind=engine)


def get_db() -> Generator[Session, None, None]:
    """
    FastAPI dependency that yields a SQLAlchemy Session and ensures proper close().
    """
    db = SessionLocal()
    try:
        yield db
    finally:
        try:
            db.close()
        except Exception:
            pass