Exception propagation (#141)

Author: s3rius

docs/guide/state-and-deps.md

@@ -181,3 +181,33 @@ Taskiq have an ability to add new first-level dependencies using brokers.
 The AsyncBroker interface has a function called `add_dependency_context` and you can add
 more default dependencies to the taskiq. This may be useful for libraries if you want to
 add new dependencies to users.
+
+
+### Exception handling
+
+Dependencies can handle exceptions that happen in tasks. This feature is handy if you want your system to be more atomic.
+
+For example, if you open a database transaction in your dependency and want to commit it only if the function is completed successfully.
+
+```python
+async def get_transaction(db_driver: DBDriver = TaskiqDepends(get_driver)) -> AsyncGenerator[Transaction, None]:
+    trans = db_driver.begin_transaction():
+    try:
+        # Here we give transaction to our dependant function.
+        yield trans
+    # If exception was found in dependant function,
+    # we rollback our transaction.
+    except Exception:
+        await trans.rollback()
+        return
+    # Here we commit if everything is fine.
+    await trans.commit()
+```
+
+If you don't want to propagate exceptions in dependencies, you can add `--no-propagate-errors` option to `worker` command.
+
+```bash
+taskiq worker my_file:broker --no-propagate-errors
+```
+
+In this case, no exception will ever going to be propagated to any dependency.

taskiq/brokers/inmemory_broker.py

@@ -88,12 +88,13 @@ class InMemoryBroker(AsyncBroker):
     It's useful for local development, if you don't want to setup real broker.
     """
 
-    def __init__(
+    def __init__(  # noqa: WPS211
         self,
         sync_tasks_pool_size: int = 4,
         max_stored_results: int = 100,
         cast_types: bool = True,
         max_async_tasks: int = 30,
+        propagate_exceptions: bool = True,
     ) -> None:
         super().__init__()
         self.result_backend = InmemoryResultBackend(
@@ -105,6 +106,7 @@ def __init__(
             executor=self.executor,
             validate_params=cast_types,
             max_async_tasks=max_async_tasks,
+            propagate_exceptions=propagate_exceptions,
         )
         self._running_tasks: "Set[asyncio.Task[Any]]" = set()
 

taskiq/cli/worker/args.py

@@ -41,6 +41,7 @@ class WorkerArgs:
     receiver: str = "taskiq.receiver:Receiver"
     receiver_arg: List[Tuple[str, str]] = field(default_factory=list)
     max_prefetch: int = 0
+    no_propagate_errors: bool = False
 
     @classmethod
     def from_cli(  # noqa: WPS213
@@ -138,6 +139,16 @@ def from_cli(  # noqa: WPS213
                 " with pydantic."
             ),
         )
+        parser.add_argument(
+            "--no-propagate-errors",
+            action="store_true",
+            dest="no_propagate_errors",
+            help=(
+                "If this parameter is on,"
+                " all errors that happen in tasks "
+                " won't be propagated to generator dependencies."
+            ),
+        )
         parser.add_argument(
             "--max-threadpool-threads",
             type=int,

taskiq/cli/worker/run.py

@@ -133,6 +133,7 @@ def interrupt_handler(signum: int, _frame: Any) -> None:
                 validate_params=not args.no_parse,
                 max_async_tasks=args.max_async_tasks,
                 max_prefetch=args.max_prefetch,
+                propagate_exceptions=not args.no_propagate_errors,
                 **receiver_args,
             )
             loop.run_until_complete(receiver.listen())

taskiq/receiver/receiver.py

@@ -46,13 +46,15 @@ def __init__(  # noqa: WPS211
         validate_params: bool = True,
         max_async_tasks: "Optional[int]" = None,
         max_prefetch: int = 0,
+        propagate_exceptions: bool = True,
     ) -> None:
         self.broker = broker
         self.executor = executor
         self.validate_params = validate_params
         self.task_signatures: Dict[str, inspect.Signature] = {}
         self.task_hints: Dict[str, Dict[str, Any]] = {}
         self.dependency_graphs: Dict[str, DependencyGraph] = {}
+        self.propagate_exceptions = propagate_exceptions
         for task in self.broker.available_tasks.values():
             self.task_signatures[task.task_name] = inspect.signature(task.original_func)
             self.task_hints[task.task_name] = get_type_hints(task.original_func)
@@ -213,7 +215,14 @@ async def run_task(  # noqa: C901, WPS210
         # Stop the timer.
         execution_time = time() - start_time
         if dep_ctx:
-            await dep_ctx.close()
+            args = (None, None, None)
+            if found_exception and self.propagate_exceptions:
+                args = (  # type: ignore
+                    type(found_exception),
+                    found_exception,
+                    found_exception.__traceback__,
+                )
+            await dep_ctx.close(*args)
 
         # Assemble result.
         result: "TaskiqResult[Any]" = TaskiqResult(