from __future__ import annotations import os from pathlib import Path from typing import Any from sqlalchemy import create_engine, event from sqlalchemy.orm import Session, sessionmaker from sqlalchemy.pool import NullPool def _normalize_database_url(url: str) -> str: if url.startswith("postgres://"): return "postgresql+psycopg://" + url[len("postgres://") :] if url.startswith("postgresql://") and "+psycopg" not in url: return "postgresql+psycopg://" + url[len("postgresql://") :] return url def _default_sqlite_url() -> str: data_dir = Path(os.getenv("CC_DATA_DIR", ".data")).resolve() data_dir.mkdir(parents=True, exist_ok=True) return f"sqlite:///{(data_dir / 'mvp_cc.db').as_posix()}" DATABASE_URL = _normalize_database_url(os.getenv("DATABASE_URL", _default_sqlite_url())) _connect_args: dict[str, Any] = {} _engine_kwargs: dict[str, Any] = {"pool_pre_ping": True, "future": True, "connect_args": _connect_args} if DATABASE_URL.startswith("sqlite"): # SQLite is used in local/demo mode where multiple services may write concurrently. # Use NullPool so short-lived requests across many local services don't exhaust # a tiny per-process QueuePool during UI polling. Keep WAL + busy timeout to # reduce "database is locked" failures. _connect_args = {"check_same_thread": False, "timeout": float(os.getenv("SQLITE_BUSY_TIMEOUT_SECONDS", "120"))} _engine_kwargs["connect_args"] = _connect_args _engine_kwargs["poolclass"] = NullPool else: _engine_kwargs.update( { "pool_size": int(os.getenv("DB_POOL_SIZE", "10")), "max_overflow": int(os.getenv("DB_MAX_OVERFLOW", "20")), "pool_timeout": int(os.getenv("DB_POOL_TIMEOUT_SECONDS", "30")), "pool_recycle": int(os.getenv("DB_POOL_RECYCLE_SECONDS", "1800")), } ) engine = create_engine(DATABASE_URL, **_engine_kwargs) if DATABASE_URL.startswith("sqlite"): @event.listens_for(engine, "connect") def _configure_sqlite_connection(dbapi_connection: Any, _connection_record: Any) -> None: cursor = dbapi_connection.cursor() cursor.execute("PRAGMA journal_mode=WAL;") cursor.execute("PRAGMA synchronous=NORMAL;") cursor.execute(f"PRAGMA busy_timeout={int(float(os.getenv('SQLITE_BUSY_TIMEOUT_SECONDS', '120')) * 1000)};") cursor.close() SessionLocal = sessionmaker(bind=engine, autoflush=False, autocommit=False, future=True) def get_session() -> Session: return SessionLocal()