from __future__ import annotations import os from pathlib import Path from sqlalchemy import create_engine from sqlalchemy.pool import NullPool from sqlalchemy.orm import sessionmaker from .models import Base DEFAULT_POSTGRES_URL = "postgresql+psycopg://openf1:openf1@localhost:5432/openf1" DEFAULT_SQLITE_URL = "sqlite:///./data/openf1.db" def _normalize_postgres_url(url: str) -> str: if url.startswith("postgresql+psycopg2://"): return url.replace("postgresql+psycopg2://", "postgresql+psycopg://", 1) if url.startswith("postgresql://"): return url.replace("postgresql://", "postgresql+psycopg://", 1) if url.startswith("postgres://"): return url.replace("postgres://", "postgresql+psycopg://", 1) return url def _build_engine_url() -> str: url = os.getenv("DATABASE_URL", DEFAULT_POSTGRES_URL) return _normalize_postgres_url(url) def _connect_args(url: str) -> dict: if url.startswith("sqlite"): return {"check_same_thread": False} return {} def _ensure_sqlite_path(url: str) -> None: if not url.startswith("sqlite"): # only create local dir for sqlite return path = url.replace("sqlite:///", "", 1) if path.startswith("./"): path = path[2:] if not path: return Path(path).parent.mkdir(parents=True, exist_ok=True) engine_url = _build_engine_url() _ensure_sqlite_path(engine_url) pooler_enabled = os.getenv("DB_POOLER", "false").lower() in {"1", "true", "yes"} engine_kwargs = {"connect_args": _connect_args(engine_url)} if pooler_enabled and engine_url.startswith("postgresql+psycopg://"): engine_kwargs["poolclass"] = NullPool engine = create_engine(engine_url, **engine_kwargs) SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine) def init_db() -> None: Base.metadata.create_all(bind=engine)