initial commit after split

src/govoplan_core/db/session.py

@@ -0,0 +1,66 @@
+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
+
+
+_default_database: DatabaseHandle | None = None
+
+
+def configure_database(database_url: str, *, engine: Engine | None = None) -> DatabaseHandle:
+    global _default_database
+    if engine is None and _default_database is not None and _default_database.database_url == database_url:
+        return _default_database
+    _default_database = DatabaseHandle(database_url, engine=engine)
+    return _default_database
+
+
+def set_database(handle: DatabaseHandle) -> DatabaseHandle:
+    global _default_database
+    _default_database = handle
+    return handle
+
+
+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()