mirror of
https://github.com/Hestia-Homes/Model.git
synced 2026-07-12 13:29:04 +00:00
65 lines
2.1 KiB
Python
65 lines
2.1 KiB
Python
from collections.abc import Iterator
|
|
from contextlib import contextmanager
|
|
from typing import Any, Optional, Type
|
|
|
|
from sqlalchemy.engine import Engine
|
|
from sqlalchemy.pool import Pool
|
|
from sqlmodel import Session, create_engine
|
|
|
|
from infrastructure.postgres.config import PostgresConfig
|
|
|
|
|
|
def make_engine(
|
|
config: PostgresConfig, poolclass: Optional[Type[Pool]] = None
|
|
) -> Engine:
|
|
kwargs: dict[str, Any] = {"pool_pre_ping": config.pool_pre_ping}
|
|
if poolclass is None:
|
|
kwargs["pool_size"] = config.pool_size
|
|
kwargs["max_overflow"] = config.max_overflow
|
|
kwargs["pool_recycle"] = config.pool_recycle
|
|
else:
|
|
kwargs["poolclass"] = poolclass
|
|
return create_engine(config.url(), **kwargs)
|
|
|
|
|
|
def make_session(engine: Engine) -> Session:
|
|
return Session(engine)
|
|
|
|
|
|
@contextmanager # pyright: ignore[reportDeprecated]
|
|
def transactional_session(engine: Engine) -> Iterator[Session]:
|
|
"""Yield a session whose lifecycle owns the transaction.
|
|
|
|
On clean exit the session commits; on any exception it rolls back and
|
|
re-raises. Either way the session is closed. Callers in the application
|
|
layer can do their work inside the ``with`` block without ever invoking
|
|
``.commit()`` / ``.rollback()`` themselves -- transaction semantics stay
|
|
in the infrastructure layer.
|
|
"""
|
|
session = Session(engine)
|
|
try:
|
|
yield session
|
|
session.commit()
|
|
except Exception:
|
|
session.rollback()
|
|
raise
|
|
finally:
|
|
session.close()
|
|
|
|
|
|
@contextmanager # pyright: ignore[reportDeprecated]
|
|
def commit_scope(session: Session) -> Iterator[Session]:
|
|
"""Commit a caller-owned session on clean exit; roll back on error.
|
|
|
|
Like ``transactional_session`` but for a session the caller already holds
|
|
and will close itself. Use it to keep slow, non-DB work *outside* the
|
|
transaction: build the session, run the slow work, then enter
|
|
``commit_scope`` only for the persistence -- so a connection is checked out
|
|
(SQLModel sessions are lazy) for the shortest possible window.
|
|
"""
|
|
try:
|
|
yield session
|
|
session.commit()
|
|
except Exception:
|
|
session.rollback()
|
|
raise
|