Merge pull request #8 from taskiq-python/bugfix/exceptions

s3rius · web-flow · commit fa643891180e · 2023-04-26T22:41:25.000+04:00
diff --git a/.flake8 b/.flake8
@@ -88,6 +88,8 @@ ignore =
   N802,
   ; Do not perform function calls in argument defaults.
   B008,
+  ; Found except `BaseException`
+  WPS424,
 
   ; all init files
   __init__.py:
diff --git a/README.md b/README.md
@@ -125,3 +125,54 @@ with graph.sync_ctx() as ctx:
 ```
 
 The ParamInfo has the information about name and parameters signature. It's useful if you want to create a dependency that changes based on parameter name, or signature.
+
+
+## Exception propagation
+
+By default if error happens within the context, we send this error to the dependency,
+so you can close it properly. You can disable this functionality by setting `exception_propagation` parameter to `False`.
+
+Let's imagine that you want to get a database session from pool and commit after the function is done.
+
+
+```python
+async def get_session():
+    session = sessionmaker()
+
+    yield session
+
+    await session.commit()
+
+```
+
+But what if the error happened when the dependant function was called? In this case you want to rollback, instead of commit.
+To solve this problem, you can just wrap the `yield` statement in `try except` to handle the error.
+
+```python
+async def get_session():
+    session = sessionmaker()
+
+    try:
+        yield session
+    except Exception:
+        await session.rollback()
+        return
+
+    await session.commit()
+
+```
+
+**Also, as a library developer, you can disable exception propagation**. If you do so, then no exception will ever be propagated to dependencies and no such `try except` expression will ever work.
+
+
+Example of disabled propogation.
+
+```python
+
+graph = DependencyGraph(target_func)
+
+with graph.sync_ctx(exception_propagation=False) as ctx:
+    print(ctx.resolve_kwargs())
+
+
+```
diff --git a/taskiq_dependencies/ctx.py b/taskiq_dependencies/ctx.py
@@ -1,6 +1,7 @@
 import asyncio
 import inspect
 from copy import copy
+from logging import getLogger
 from typing import TYPE_CHECKING, Any, Dict, Generator, List, Optional
 
 from taskiq_dependencies.utils import ParamInfo
@@ -9,18 +10,23 @@
     from taskiq_dependencies.graph import DependencyGraph  # pragma: no cover
 
 
+logger = getLogger("taskiq.dependencies.ctx")
+
+
 class BaseResolveContext:
     """Base resolver context."""
 
     def __init__(
         self,
         graph: "DependencyGraph",
         initial_cache: Optional[Dict[Any, Any]] = None,
+        exception_propagation: bool = True,
     ) -> None:
         self.graph = graph
         self.opened_dependencies: List[Any] = []
         self.sub_contexts: "List[Any]" = []
         self.initial_cache = initial_cache or {}
+        self.propagate_excs = exception_propagation
 
     def traverse_deps(  # noqa: C901, WPS210
         self,
@@ -116,18 +122,34 @@ def __enter__(self) -> "SyncResolveContext":
         return self
 
     def __exit__(self, *args: Any) -> None:
-        self.close()
+        self.close(*args)
 
-    def close(self) -> None:
+    def close(self, *args: Any) -> None:  # noqa: C901
         """
         Close all opened dependencies.
 
         This function runs teardown of all dependencies.
+
+        :param args: exception info if any.
         """
+        exception_found = False
+        if args[1] is not None and self.propagate_excs:
+            exception_found = True
         for ctx in self.sub_contexts:
-            ctx.close()
+            ctx.close(*args)
         for dep in reversed(self.opened_dependencies):
             if inspect.isgenerator(dep):
+                if exception_found:
+                    try:
+                        dep.throw(*args)
+                    except BaseException as exc:
+                        logger.warning(
+                            "Exception found on dependency teardown %s",
+                            exc,
+                            exc_info=True,
+                        )
+                        continue
+                    continue
                 for _ in dep:  # noqa: WPS328
                     pass  # noqa: WPS420
 
@@ -201,21 +223,48 @@ async def __aenter__(self) -> "AsyncResolveContext":
         return self
 
     async def __aexit__(self, *args: Any) -> None:
-        await self.close()
+        await self.close(*args)
 
-    async def close(self) -> None:  # noqa: C901
+    async def close(self, *args: Any) -> None:  # noqa: C901
         """
         Close all opened dependencies.
 
         This function runs teardown of all dependencies.
+
+        :param args: exception info if any.
         """
+        exception_found = False
+        if args[1] is not None and self.propagate_excs:
+            exception_found = True
         for ctx in self.sub_contexts:
-            await ctx.close()  # type: ignore
+            await ctx.close(*args)  # type: ignore
         for dep in reversed(self.opened_dependencies):
             if inspect.isgenerator(dep):
+                if exception_found:
+                    try:
+                        dep.throw(*args)
+                    except BaseException as exc:
+                        logger.warning(
+                            "Exception found on dependency teardown %s",
+                            exc,
+                            exc_info=True,
+                        )
+                        continue
+                    continue
                 for _ in dep:  # noqa: WPS328
                     pass  # noqa: WPS420
             elif inspect.isasyncgen(dep):
+                if exception_found:
+                    try:
+                        await dep.athrow(*args)
+                    except BaseException as exc:
+                        logger.warning(
+                            "Exception found on dependency teardown %s",
+                            exc,
+                            exc_info=True,
+                        )
+                        continue
+                    continue
                 async for _ in dep:  # noqa: WPS328
                     pass  # noqa: WPS420
 
diff --git a/taskiq_dependencies/graph.py b/taskiq_dependencies/graph.py
@@ -41,35 +41,43 @@ def is_empty(self) -> bool:
     def async_ctx(
         self,
         initial_cache: Optional[Dict[Any, Any]] = None,
+        exception_propagation: bool = True,
     ) -> AsyncResolveContext:
         """
         Create dependency resolver context.
 
         This context is used to actually resolve dependencies.
 
         :param initial_cache: initial cache dict.
+        :param exception_propagation: If true, all found errors within
+            context will be propagated to dependencies.
         :return: new resolver context.
         """
         return AsyncResolveContext(
             self,
             initial_cache,
+            exception_propagation,
         )
 
     def sync_ctx(
         self,
         initial_cache: Optional[Dict[Any, Any]] = None,
+        exception_propagation: bool = True,
     ) -> SyncResolveContext:
         """
         Create dependency resolver context.
 
         This context is used to actually resolve dependencies.
 
         :param initial_cache: initial cache dict.
+        :param exception_propagation: If true, all found errors within
+            context will be propagated to dependencies.
         :return: new resolver context.
         """
         return SyncResolveContext(
             self,
             initial_cache,
+            exception_propagation,
         )
 
     def _build_graph(self) -> None:  # noqa: C901, WPS210
diff --git a/tests/test_graph.py b/tests/test_graph.py
@@ -310,3 +310,140 @@ def target(class_val: str = Depends(TeClass("tval"))) -> None:
 
     info: str = kwargs["class_val"]
     assert info == "tval"
+
+
+def test_exception_generators() -> None:
+
+    errors_found = 0
+
+    def my_generator() -> Generator[int, None, None]:
+        nonlocal errors_found
+        try:
+            yield 1
+        except ValueError:
+            errors_found += 1
+
+    def target(_: int = Depends(my_generator)) -> None:
+        raise ValueError()
+
+    with pytest.raises(ValueError):
+        with DependencyGraph(target=target).sync_ctx() as g:
+            target(**g.resolve_kwargs())
+
+    assert errors_found == 1
+
+
+@pytest.mark.anyio
+async def test_async_exception_generators() -> None:
+
+    errors_found = 0
+
+    async def my_generator() -> AsyncGenerator[int, None]:
+        nonlocal errors_found
+        try:
+            yield 1
+        except ValueError:
+            errors_found += 1
+
+    def target(_: int = Depends(my_generator)) -> None:
+        raise ValueError()
+
+    with pytest.raises(ValueError):
+        async with DependencyGraph(target=target).async_ctx() as g:
+            target(**(await g.resolve_kwargs()))
+
+    assert errors_found == 1
+
+
+@pytest.mark.anyio
+async def test_async_exception_generators_multiple() -> None:
+
+    errors_found = 0
+
+    async def my_generator() -> AsyncGenerator[int, None]:
+        nonlocal errors_found
+        try:
+            yield 1
+        except ValueError:
+            errors_found += 1
+
+    def target(
+        _a: int = Depends(my_generator, use_cache=False),
+        _b: int = Depends(my_generator, use_cache=False),
+        _c: int = Depends(my_generator, use_cache=False),
+    ) -> None:
+        raise ValueError()
+
+    with pytest.raises(ValueError):
+        async with DependencyGraph(target=target).async_ctx() as g:
+            target(**(await g.resolve_kwargs()))
+
+    assert errors_found == 3
+
+
+@pytest.mark.anyio
+async def test_async_exception_in_teardown() -> None:
+
+    errors_found = 0
+
+    async def my_generator() -> AsyncGenerator[int, None]:
+        nonlocal errors_found
+        try:
+            yield 1
+        except ValueError:
+            errors_found += 1
+            raise Exception()
+
+    def target(_: int = Depends(my_generator)) -> None:
+        raise ValueError()
+
+    with pytest.raises(ValueError):
+        async with DependencyGraph(target=target).async_ctx() as g:
+            target(**(await g.resolve_kwargs()))
+
+
+@pytest.mark.anyio
+async def test_async_propagation_disabled() -> None:
+
+    errors_found = 0
+
+    async def my_generator() -> AsyncGenerator[int, None]:
+        nonlocal errors_found
+        try:
+            yield 1
+        except ValueError:
+            errors_found += 1
+            raise Exception()
+
+    def target(_: int = Depends(my_generator)) -> None:
+        raise ValueError()
+
+    with pytest.raises(ValueError):
+        async with DependencyGraph(target=target).async_ctx(
+            exception_propagation=False,
+        ) as g:
+            target(**(await g.resolve_kwargs()))
+
+    assert errors_found == 0
+
+
+def test_sync_propagation_disabled() -> None:
+
+    errors_found = 0
+
+    def my_generator() -> Generator[int, None, None]:
+        nonlocal errors_found
+        try:
+            yield 1
+        except ValueError:
+            errors_found += 1
+            raise Exception()
+
+    def target(_: int = Depends(my_generator)) -> None:
+        raise ValueError()
+
+    with pytest.raises(ValueError):
+        with DependencyGraph(target=target).sync_ctx(exception_propagation=False) as g:
+            target(**(g.resolve_kwargs()))
+
+    assert errors_found == 0
