80 lines
2.6 KiB
Python
80 lines
2.6 KiB
Python
from __future__ import annotations
|
|
|
|
from collections.abc import Generator, Mapping
|
|
from typing import Any
|
|
|
|
from sqlalchemy import create_engine
|
|
from sqlalchemy.engine import Engine
|
|
from sqlalchemy.orm import Session, sessionmaker
|
|
|
|
|
|
def default_connect_args(database_url: str) -> dict[str, Any]:
|
|
return {"check_same_thread": False} if database_url.startswith("sqlite") else {}
|
|
|
|
|
|
def create_database_engine(
|
|
database_url: str,
|
|
*,
|
|
pool_pre_ping: bool = True,
|
|
connect_args: Mapping[str, Any] | None = None,
|
|
**kwargs: Any,
|
|
) -> Engine:
|
|
merged_connect_args = dict(default_connect_args(database_url))
|
|
if connect_args:
|
|
merged_connect_args.update(connect_args)
|
|
return create_engine(database_url, pool_pre_ping=pool_pre_ping, connect_args=merged_connect_args, **kwargs)
|
|
|
|
|
|
class DatabaseHandle:
|
|
def __init__(self, database_url: str, *, engine: Engine | None = None) -> None:
|
|
self.database_url = database_url
|
|
self.engine = engine or create_database_engine(database_url)
|
|
self.SessionLocal = sessionmaker(bind=self.engine, autoflush=False, autocommit=False, expire_on_commit=False)
|
|
|
|
def session(self) -> Session:
|
|
return self.SessionLocal()
|
|
|
|
def dependency(self) -> Generator[Session, None, None]:
|
|
with self.SessionLocal() as session:
|
|
yield session
|
|
|
|
def dispose(self) -> None:
|
|
self.engine.dispose()
|
|
|
|
|
|
_default_database: DatabaseHandle | None = None
|
|
|
|
|
|
def configure_database(database_url: str, *, engine: Engine | None = None, dispose_previous: bool = False) -> DatabaseHandle:
|
|
global _default_database
|
|
if engine is None and _default_database is not None and _default_database.database_url == database_url:
|
|
return _default_database
|
|
previous_database = _default_database
|
|
_default_database = DatabaseHandle(database_url, engine=engine)
|
|
if dispose_previous and previous_database is not None and previous_database is not _default_database:
|
|
previous_database.dispose()
|
|
return _default_database
|
|
|
|
|
|
def set_database(handle: DatabaseHandle) -> DatabaseHandle:
|
|
global _default_database
|
|
_default_database = handle
|
|
return handle
|
|
|
|
|
|
def reset_database(*, dispose: bool = False) -> None:
|
|
global _default_database
|
|
if dispose and _default_database is not None:
|
|
_default_database.dispose()
|
|
_default_database = None
|
|
|
|
|
|
def get_database() -> DatabaseHandle:
|
|
if _default_database is None:
|
|
raise RuntimeError("GovOPlaN database is not configured")
|
|
return _default_database
|
|
|
|
|
|
def get_session() -> Generator[Session, None, None]:
|
|
yield from get_database().dependency()
|