services/api/db.py

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)