mirror of
https://github.com/Hestia-Homes/Model.git
synced 2026-07-12 13:29:04 +00:00
39 lines
1.1 KiB
Python
39 lines
1.1 KiB
Python
from __future__ import annotations
|
|
|
|
from collections.abc import Callable
|
|
from types import TracebackType
|
|
from typing import Optional
|
|
|
|
from sqlmodel import Session
|
|
|
|
from repositories.magic_plan.magic_plan_postgres_repository import (
|
|
MagicPlanPostgresRepository,
|
|
)
|
|
from repositories.uploaded_file.uploaded_file_postgres_repository import (
|
|
UploadedFilePostgresRepository,
|
|
)
|
|
|
|
|
|
class AuditGeneratorUnitOfWork:
|
|
def __init__(self, session_factory: Callable[[], Session]) -> None:
|
|
self._session_factory = session_factory
|
|
|
|
def __enter__(self) -> "AuditGeneratorUnitOfWork":
|
|
self._session = self._session_factory()
|
|
self.uploaded_file = UploadedFilePostgresRepository(self._session)
|
|
self.magic_plan = MagicPlanPostgresRepository(self._session)
|
|
return self
|
|
|
|
def __exit__(
|
|
self,
|
|
exc_type: Optional[type[BaseException]],
|
|
exc: Optional[BaseException],
|
|
tb: Optional[TracebackType],
|
|
) -> None:
|
|
try:
|
|
self._session.rollback()
|
|
finally:
|
|
self._session.close()
|
|
|
|
def commit(self) -> None:
|
|
self._session.commit()
|