Merge pull request #38 from febus982/unit_of_work_rewrite

Unit of work rewrite

Author: febus982

README.md

@@ -182,13 +182,13 @@ be used for repository operations, **assumed the same bind is used for all the r
 ```python
 class MyRepo(SQLAlchemyRepository):
     _model = MyModel
-class MyOtherRepo(SQLAlchemyRepository):
-    _model = MyOtherModel
 
 bind = sa_manager.get_bind()
-uow = UnitOfWork(bind, (MyRepo, MyOtherRepo))
+uow = UnitOfWork(bind)
+uow.register_repository("repo_a", MyRepo)
+uow.register_repository("repo_b", SQLAlchemyRepository, MyOtherModel)
 
 with uow.transaction():
-    uow.MyRepo.save(some_model)
-    uow.MyOtherRepo.save(some_other_model)
+    uow.repository("repo_a").save(some_model)
+    uow.repository("repo_b").save(some_other_model)
 ```

docs/lifecycle.md

@@ -34,12 +34,16 @@ What you can do is:
 * Save the repositories in global variables and start a thread / asyncio task to handle
   a scoped request (e.g. one thread per HTTP request)
 
-What you cannot do is:
+What you should not do is:
 
 * Get a list of models
 * Save the models using `save()` in parallel threads / tasks (each task will have a different session)
 
-/// tip | The recommendation is of course to try to use a single repository instance, where possible.
+/// warning | Remember: Concurrent writes to the db can cause undesired scenarios like locks and deadlocks!
+
+///
+
+/// tip | The recommendation is to try to use a single repository instance, where possible.
 
 For example a strategy similar to this would be optimal, if possible:
 

docs/repository/uow.md

@@ -10,20 +10,21 @@ be used for repository operations.
 ```python
 class MyRepo(SQLAlchemyRepository):
     _model = MyModel
-class MyOtherRepo(SQLAlchemyRepository):
-    _model = MyOtherModel
 
 bind = sa_manager.get_bind()
-uow = UnitOfWork(bind, (MyRepo, MyOtherRepo))
+uow = UnitOfWork(bind)
+uow.register_repository("repo_a", MyRepo)
+# args and kwargs are forwarded so we can also use directly `SQLAlchemyRepository` class
+uow.register_repository("repo_b", SQLAlchemyRepository, MyOtherModel)
 
 with uow.transaction():
-    uow.MyRepo.save(some_model)
-    uow.MyOtherRepo.save(some_other_model)
+    uow.repository("repo_a").save(some_model)
+    uow.repository("repo_b").save(some_other_model)
 
 # Optionally disable the commit/rollback handling
 with uow.transaction(read_only=True):
-    model1 = uow.MyRepo.get(1)
-    model2 = uow.MyOtherRepo.get(2)
+    model1 = uow.repository("repo_a").get(1)
+    model2 = uow.repository("repo_b").get(2)
 ```
 
 /// admonition | The unit of work implementation is still experimental.

sqlalchemy_bind_manager/_unit_of_work/__init__.py

@@ -1,5 +1,6 @@
+from abc import ABC
 from contextlib import asynccontextmanager, contextmanager
-from typing import AsyncIterator, Iterable, Iterator, Type
+from typing import AsyncIterator, Dict, Generic, Iterator, Type, TypeVar, Union
 
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.orm import Session
@@ -9,41 +10,63 @@
     AsyncSessionHandler,
     SessionHandler,
 )
+from sqlalchemy_bind_manager.exceptions import RepositoryNotFound
 from sqlalchemy_bind_manager.repository import (
     SQLAlchemyAsyncRepository,
     SQLAlchemyRepository,
 )
 
+REPOSITORY = TypeVar("REPOSITORY", SQLAlchemyRepository, SQLAlchemyAsyncRepository)
+SESSION_HANDLER = TypeVar("SESSION_HANDLER", SessionHandler, AsyncSessionHandler)
 
-class UnitOfWork:
-    _session_handler: SessionHandler
 
-    def __init__(
-        self, bind: SQLAlchemyBind, repositories: Iterable[Type[SQLAlchemyRepository]]
-    ) -> None:
+class BaseUnitOfWork(Generic[REPOSITORY, SESSION_HANDLER], ABC):
+    _session_handler: SESSION_HANDLER
+    _repositories: Dict[str, REPOSITORY] = {}
+
+    def register_repository(
+        self,
+        name: str,
+        repository_class: Type[REPOSITORY],
+        model_class: Union[Type, None] = None,
+        *args,
+        **kwargs,
+    ):
+        kwargs.pop("session", None)
+        self._repositories[name] = repository_class(  # type: ignore
+            *args,
+            session=self._session_handler.scoped_session(),
+            model_class=model_class,
+            **kwargs,
+        )
+
+    def repository(self, name: str) -> REPOSITORY:
+        try:
+            return self._repositories[name]
+        except KeyError:
+            raise RepositoryNotFound(
+                "The repository has not been initialised in this unit of work"
+            )
+
+
+class UnitOfWork(BaseUnitOfWork[SQLAlchemyRepository, SessionHandler]):
+    def __init__(self, bind: SQLAlchemyBind) -> None:
         super().__init__()
         self._session_handler = SessionHandler(bind)
-        for r in repositories:
-            setattr(self, r.__name__, r(session=self._session_handler.scoped_session()))
 
     @contextmanager
     def transaction(self, read_only: bool = False) -> Iterator[Session]:
         with self._session_handler.get_session(read_only=read_only) as _s:
             yield _s
 
 
-class AsyncUnitOfWork:
-    _session_handler: AsyncSessionHandler
-
+class AsyncUnitOfWork(BaseUnitOfWork[SQLAlchemyAsyncRepository, AsyncSessionHandler]):
     def __init__(
         self,
         bind: SQLAlchemyAsyncBind,
-        repositories: Iterable[Type[SQLAlchemyAsyncRepository]],
     ) -> None:
         super().__init__()
         self._session_handler = AsyncSessionHandler(bind)
-        for r in repositories:
-            setattr(self, r.__name__, r(session=self._session_handler.scoped_session()))
 
     @asynccontextmanager
     async def transaction(self, read_only: bool = False) -> AsyncIterator[AsyncSession]:

sqlalchemy_bind_manager/exceptions.py

@@ -24,3 +24,7 @@ class UnmappedProperty(Exception):
 
 class SessionNotFound(Exception):
     pass
+
+
+class RepositoryNotFound(Exception):
+    pass

tests/unit_of_work/test_lifecycle.py

@@ -1,3 +1,10 @@
+from unittest.mock import MagicMock
+
+import pytest
+
+from sqlalchemy_bind_manager.exceptions import RepositoryNotFound
+
+
 async def test_repositories_are_initialised_with_uow_session(
     sa_bind, repository_class, model_classes, uow_class
 ):
@@ -8,15 +15,72 @@ class ChildRepoClass(repository_class):
         _model = model_classes[1]
 
     repository_classes = [RepoClass, ChildRepoClass]
-    uow = uow_class(
-        bind=sa_bind,
-        repositories=repository_classes,
-    )
+    uow = uow_class(bind=sa_bind)
+    uow.register_repository(RepoClass.__name__, RepoClass)
+    uow.register_repository("ChildRepoClass", repository_class, model_classes[1])
+
     for repo_class in repository_classes:
-        assert hasattr(uow, repo_class.__name__)
-        repo = getattr(uow, repo_class.__name__)
+        repo = uow.repository(repo_class.__name__)
+        assert repo is not None
         assert not hasattr(repo, "_session_handler")
         assert hasattr(repo, "_external_session")
         assert (
             getattr(repo, "_external_session") is uow._session_handler.scoped_session()
         )
+
+
+async def test_raises_exception_if_repository_not_found(sa_bind, uow_class):
+    uow = uow_class(bind=sa_bind)
+    with pytest.raises(RepositoryNotFound):
+        uow.repository("Not existing")
+
+
+@pytest.mark.parametrize(
+    ["submitted_args", "submitted_kwargs", "received_args", "received_kwargs"],
+    [
+        pytest.param(
+            ("1", "2"),
+            dict(a="b"),
+            ("2",),
+            dict(model_class="1", a="b"),
+            id="first_arg_model_class_if_no_kwarg",
+        ),
+        pytest.param(
+            tuple([]),
+            dict(a="b", model_class="c"),
+            tuple([]),
+            dict(model_class="c", a="b"),
+            id="model_class_rearranged_if_in_kwargs",
+        ),
+        pytest.param(
+            tuple([]),
+            dict(a="b"),
+            tuple([]),
+            dict(model_class=None, a="b"),
+            id="model_class_default_to_none",
+        ),
+        pytest.param(
+            tuple([]),
+            dict(a="b", session="c"),
+            tuple([]),
+            dict(model_class=None, a="b"),
+            id="session_removed_from_kwargs",
+        ),
+    ],
+)
+async def test_additional_arguments_are_forwarded(
+    sa_bind,
+    uow_class,
+    submitted_args: tuple,
+    submitted_kwargs: dict,
+    received_args: tuple,
+    received_kwargs: dict,
+):
+    repo = MagicMock()
+
+    uow = uow_class(bind=sa_bind)
+    uow.register_repository("r", repo, *submitted_args, **submitted_kwargs)
+
+    repo.assert_called_once_with(
+        *received_args, session=uow._session_handler.scoped_session(), **received_kwargs
+    )

tests/unit_of_work/test_session_behaviour.py

@@ -19,14 +19,14 @@ class RepoClass(repository_class):
     class ChildRepoClass(repository_class):
         _model = model_classes[1]
 
-    repository_classes = [RepoClass, ChildRepoClass]
-
     with patch.object(
         session_handler_class, "commit", return_value=None
     ) as mocked_uow_commit:
-        uow = uow_class(sa_bind, repository_classes)
-        repo1 = getattr(uow, repository_classes[0].__name__)
-        repo2 = getattr(uow, repository_classes[1].__name__)
+        uow = uow_class(sa_bind)
+        uow.register_repository(RepoClass.__name__, RepoClass)
+        uow.register_repository(ChildRepoClass.__name__, ChildRepoClass)
+        repo1 = uow.repository(RepoClass.__name__)
+        repo2 = uow.repository(ChildRepoClass.__name__)
 
         # Populate a database entry to be used for tests
         model1 = model_classes[0](
@@ -56,10 +56,11 @@ class RepoClass(repository_class):
     class OtherRepoClass(repository_class):
         _model = model_classes[0]
 
-    repository_classes = [RepoClass, OtherRepoClass]
-    uow = uow_class(sa_bind, repository_classes)
-    repo1 = getattr(uow, repository_classes[0].__name__)
-    repo2 = getattr(uow, repository_classes[1].__name__)
+    uow = uow_class(sa_bind)
+    uow.register_repository(RepoClass.__name__, RepoClass)
+    uow.register_repository(OtherRepoClass.__name__, OtherRepoClass)
+    repo1 = uow.repository(RepoClass.__name__)
+    repo2 = uow.repository(OtherRepoClass.__name__)
 
     # Populate a database entry to be used for tests
     model1 = model_classes[0](
@@ -87,10 +88,9 @@ async def test_uow_repository_operations_fail_without_transaction(
     class RepoClass(repository_class):
         _model = model_classes[0]
 
-    repository_classes = [RepoClass]
-
-    uow = uow_class(sa_bind, repository_classes)
-    repo1 = getattr(uow, repository_classes[0].__name__)
+    uow = uow_class(sa_bind)
+    uow.register_repository(RepoClass.__name__, RepoClass)
+    repo1 = uow.repository(RepoClass.__name__)
 
     # Populate a database entry to be used for tests
     model1 = model_classes[0](
@@ -115,11 +115,11 @@ class RepoClass(repository_class):
     class OtherRepoClass(repository_class):
         _model = model_classes[0]
 
-    repository_classes = [RepoClass, OtherRepoClass]
-
-    uow = uow_class(sa_bind, repository_classes)
-    repo1 = getattr(uow, repository_classes[0].__name__)
-    repo2 = getattr(uow, repository_classes[1].__name__)
+    uow = uow_class(sa_bind)
+    uow.register_repository(RepoClass.__name__, RepoClass)
+    uow.register_repository(OtherRepoClass.__name__, OtherRepoClass)
+    repo1 = uow.repository(RepoClass.__name__)
+    repo2 = uow.repository(OtherRepoClass.__name__)
 
     # Populate a database entry to be used for tests
     model1 = model_classes[0](