interpreter/app/db.py

from collections.abc import Iterator

from sqlalchemy.engine import Engine
from sqlmodel import Session, SQLModel, create_engine

from app import crud as _crud
from app.config import get_settings
from app.glossary import store as _glossary_store

del _crud, _glossary_store

_engine: Engine | None = None


def make_engine(database_url: str | None = None) -> Engine:
    url = database_url or get_settings().database_url
    connect_args = {"check_same_thread": False} if url.startswith("sqlite") else {}
    return create_engine(url, connect_args=connect_args)


def get_engine() -> Engine:
    global _engine
    if _engine is None:
        _engine = make_engine()
    return _engine


def set_engine(engine: Engine | None) -> None:
    global _engine
    _engine = engine


def init_db(engine: Engine | None = None) -> None:
    SQLModel.metadata.create_all(engine or get_engine())


def get_db() -> Iterator[Session]:
    with Session(get_engine()) as session:
        yield session