✨ feat: add unit of work abstraction

packages/common-library/src/common_library/unit_of_work.py

@@ -0,0 +1,26 @@
+from abc import ABC, abstractmethod
+from contextlib import AbstractAsyncContextManager
+
+
+class ReadUnitOfWork:
+    """Opaque persistence scope for sequential reads."""
+
+
+class TransactionalUnitOfWork(ReadUnitOfWork):
+    """Opaque persistence scope for sequential reads and writes."""
+
+
+class UnitOfWorkFactory(ABC):
+    @abstractmethod
+    def read(
+        self,
+        *,
+        existing: ReadUnitOfWork | None = None,
+    ) -> AbstractAsyncContextManager[ReadUnitOfWork]: ...
+
+    @abstractmethod
+    def transaction(
+        self,
+        *,
+        existing: TransactionalUnitOfWork | None = None,
+    ) -> AbstractAsyncContextManager[TransactionalUnitOfWork]: ...

packages/common-library/tests/test_unit_of_work.py

@@ -0,0 +1,75 @@
+import inspect
+from contextlib import AbstractAsyncContextManager
+from types import TracebackType
+
+from common_library.unit_of_work import (
+    ReadUnitOfWork,
+    TransactionalUnitOfWork,
+    UnitOfWorkFactory,
+)
+
+
+class _ReadUnitOfWork(ReadUnitOfWork): ...
+
+
+class _TransactionalUnitOfWork(TransactionalUnitOfWork): ...
+
+
+class _UnitOfWorkContext[UnitOfWorkT: ReadUnitOfWork]:
+    def __init__(self, unit_of_work: UnitOfWorkT) -> None:
+        self._unit_of_work = unit_of_work
+
+    async def __aenter__(self) -> UnitOfWorkT:
+        return self._unit_of_work
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> bool:
+        return False
+
+
+class _IncompleteUnitOfWorkFactory(UnitOfWorkFactory): ...
+
+
+class _UnitOfWorkFactory(UnitOfWorkFactory):
+    def read(
+        self,
+        *,
+        existing: ReadUnitOfWork | None = None,
+    ) -> AbstractAsyncContextManager[ReadUnitOfWork]:
+        self._read_calls += 1
+        return _UnitOfWorkContext(existing or _ReadUnitOfWork())
+
+    def transaction(
+        self,
+        *,
+        existing: TransactionalUnitOfWork | None = None,
+    ) -> AbstractAsyncContextManager[TransactionalUnitOfWork]:
+        self._transaction_calls += 1
+        return _UnitOfWorkContext(existing or _TransactionalUnitOfWork())
+
+    def __init__(self) -> None:
+        self._read_calls = 0
+        self._transaction_calls = 0
+
+
+def test_incomplete_unit_of_work_factory_cannot_be_instantiated():
+    assert inspect.isabstract(_IncompleteUnitOfWorkFactory)
+
+
+async def test_unit_of_work_factory_contract_supports_new_and_existing_scopes():
+    factory = _UnitOfWorkFactory()
+
+    async with factory.read() as read_uow:
+        assert isinstance(read_uow, ReadUnitOfWork)
+        async with factory.read(existing=read_uow) as reused_read_uow:
+            assert reused_read_uow is read_uow
+
+    async with factory.transaction() as transactional_uow:
+        assert isinstance(transactional_uow, TransactionalUnitOfWork)
+        assert isinstance(transactional_uow, ReadUnitOfWork)
+        async with factory.transaction(existing=transactional_uow) as reused_transactional_uow:
+            assert reused_transactional_uow is transactional_uow

packages/postgres-database/src/simcore_postgres_database/unit_of_work.py

@@ -0,0 +1,86 @@
+from collections.abc import AsyncIterator
+from contextlib import AbstractAsyncContextManager, asynccontextmanager
+from dataclasses import dataclass
+
+from common_library.unit_of_work import (
+    ReadUnitOfWork,
+    TransactionalUnitOfWork,
+    UnitOfWorkFactory,
+)
+from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine
+
+
+@dataclass(frozen=True, kw_only=True, slots=True)
+class _SqlAlchemyReadUnitOfWork(ReadUnitOfWork):
+    connection: AsyncConnection
+
+
+@dataclass(frozen=True, kw_only=True, slots=True)
+class _SqlAlchemyTransactionalUnitOfWork(TransactionalUnitOfWork):
+    connection: AsyncConnection
+
+
+def get_sqlalchemy_connection(unit_of_work: ReadUnitOfWork) -> AsyncConnection:
+    if isinstance(
+        unit_of_work,
+        (_SqlAlchemyReadUnitOfWork, _SqlAlchemyTransactionalUnitOfWork),
+    ):
+        return unit_of_work.connection
+    msg = f"Expected a SQLAlchemy unit of work, got {type(unit_of_work).__name__}"
+    raise TypeError(msg)
+
+
+def get_sqlalchemy_transaction_connection(
+    unit_of_work: TransactionalUnitOfWork,
+) -> AsyncConnection:
+    if isinstance(unit_of_work, _SqlAlchemyTransactionalUnitOfWork):
+        return unit_of_work.connection
+    msg = f"Expected a SQLAlchemy transactional unit of work, got {type(unit_of_work).__name__}"
+    raise TypeError(msg)
+
+
+@asynccontextmanager
+async def _read_scope(
+    engine: AsyncEngine,
+    existing: ReadUnitOfWork | None,
+) -> AsyncIterator[ReadUnitOfWork]:
+    if existing is not None:
+        get_sqlalchemy_connection(existing)
+        yield existing
+        return
+
+    async with engine.connect() as connection:
+        yield _SqlAlchemyReadUnitOfWork(connection=connection)
+
+
+@asynccontextmanager
+async def _transaction_scope(
+    engine: AsyncEngine,
+    existing: TransactionalUnitOfWork | None,
+) -> AsyncIterator[TransactionalUnitOfWork]:
+    if existing is not None:
+        get_sqlalchemy_transaction_connection(existing)
+        yield existing
+        return
+
+    async with engine.begin() as connection:
+        yield _SqlAlchemyTransactionalUnitOfWork(connection=connection)
+
+
+@dataclass(frozen=True, kw_only=True, slots=True)
+class SqlAlchemyUnitOfWorkFactory(UnitOfWorkFactory):
+    engine: AsyncEngine
+
+    def read(
+        self,
+        *,
+        existing: ReadUnitOfWork | None = None,
+    ) -> AbstractAsyncContextManager[ReadUnitOfWork]:
+        return _read_scope(self.engine, existing)
+
+    def transaction(
+        self,
+        *,
+        existing: TransactionalUnitOfWork | None = None,
+    ) -> AbstractAsyncContextManager[TransactionalUnitOfWork]:
+        return _transaction_scope(self.engine, existing)

packages/postgres-database/tests/unit_of_work/conftest.py

@@ -0,0 +1,68 @@
+from dataclasses import dataclass
+from types import TracebackType
+from typing import cast
+
+import pytest
+from simcore_postgres_database.unit_of_work import SqlAlchemyUnitOfWorkFactory
+from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine
+
+
+@dataclass
+class _ScopeState:
+    closed: bool = False
+    committed: bool = False
+    rolled_back: bool = False
+
+
+class _ConnectionScope:
+    def __init__(self, connection: AsyncConnection, state: _ScopeState) -> None:
+        self._connection = connection
+        self._state = state
+
+    async def __aenter__(self) -> AsyncConnection:
+        return self._connection
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> bool:
+        self._state.closed = True
+        return False
+
+
+class _TransactionScope(_ConnectionScope):
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> bool:
+        self._state.closed = True
+        self._state.committed = exc_type is None
+        self._state.rolled_back = exc_type is not None
+        return False
+
+
+class _Engine:
+    def __init__(self) -> None:
+        self.connection = cast(AsyncConnection, object())
+        self.read_scopes: list[_ScopeState] = []
+        self.transaction_scopes: list[_ScopeState] = []
+
+    def connect(self) -> _ConnectionScope:
+        state = _ScopeState()
+        self.read_scopes.append(state)
+        return _ConnectionScope(self.connection, state)
+
+    def begin(self) -> _TransactionScope:
+        state = _ScopeState()
+        self.transaction_scopes.append(state)
+        return _TransactionScope(self.connection, state)
+
+
+@pytest.fixture
+def sqlalchemy_uow_factory() -> tuple[SqlAlchemyUnitOfWorkFactory, _Engine]:
+    engine = _Engine()
+    return SqlAlchemyUnitOfWorkFactory(engine=cast(AsyncEngine, engine)), engine