Rewrite Unit of Work interface methods

Author: febus982

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,57 @@
     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,
+    ):
+        self._repositories[name] = repository_class(
+            session=self._session_handler.scoped_session(), model_class=model_class
+        )
+
+    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,8 @@
+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 +13,21 @@ 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")

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](