64 lines
2.5 KiB
Python
64 lines
2.5 KiB
Python
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()
|