inventory/crudkit/engines.py

67 lines
2.7 KiB
Python

# engines.py
from __future__ import annotations
from typing import Type, Optional
from sqlalchemy import create_engine, event
from sqlalchemy.orm import sessionmaker, raiseload, Mapper, RelationshipProperty
from .backend import make_backend_info, BackendInfo
from .config import Config, get_config
from ._sqlite import apply_sqlite_pragmas
def build_engine(config_cls: Type[Config] | None = None):
config_cls = config_cls or get_config(None)
engine = create_engine(config_cls.DATABASE_URL, **config_cls.engine_kwargs())
apply_sqlite_pragmas(engine, config_cls.SQLITE_PRAGMAS)
return engine
def _install_nplus1_guards(SessionMaker, *, strict: bool):
if not strict:
return
@event.listens_for(SessionMaker, "do_orm_execute")
def _add_global_raiseload(execute_state):
stmt = execute_state.statement
# Only touch ORM statements (have column_descriptions)
if getattr(stmt, "column_descriptions", None):
execute_state.statement = stmt.options(raiseload("*"))
def build_sessionmaker(config_cls: Type[Config] | None = None, engine=None):
config_cls = config_cls or get_config(None)
engine = engine or build_engine(config_cls)
SessionMaker = sessionmaker(bind=engine, **config_cls.session_kwargs())
# Toggle with a config flag; default off so you can turn it on when ready
strict = bool(getattr(config_cls, "STRICT_NPLUS1", False))
_install_nplus1_guards(SessionMaker, strict=strict)
return SessionMaker
class CRUDKitRuntime:
"""
Lightweight container so CRUDKit can be given either:
- prebuilt engine/sessionmaker, or
- a Config to build them lazily
"""
def __init__(self, *, engine=None, session_factory=None, config: Optional[Type[Config]] = None):
if engine is None and session_factory is None and config is None:
config = get_config(None)
self._config = config
self._engine = engine or (build_engine(config) if config else None)
self._session_factory = session_factory or (build_sessionmaker(config, self._engine) if config else None)
@property
def engine(self):
if self._engine is None and self._config:
self._engine = build_engine(self._config)
return self._engine
@property
def session_factory(self):
if self._session_factory is None:
if self._config and self._engine:
self._session_factory = build_sessionmaker(self._config, self._engine)
return self._session_factory
@property
def backend(self) -> BackendInfo:
if not hasattr(self, "_backend_info") or self._backend_info is None:
self._backend_info = make_backend_info(self.engine)
return self._backend_info