feat!: replace decorators with middlewares

Author: dimastbk

poetry.lock

@@ -36,6 +36,7 @@ pytest = ">=7.4,<9.0"
 pytest-asyncio = "^0.21.1"
 pytest-grpc = "^0.8.0"
 pytest-mock = "^3.11.1"
+pytest-cov = "^5.0.0"
 pylint = ">=2.17.5,<4.0.0"
 black = "^25.0.0"
 mypy = "^1.10.0"
@@ -93,7 +94,8 @@ markers = [
 
 [tool.coverage.run]
 omit = [
-    "pyzeebe/proto/*"
+    "pyzeebe/proto/*",
+    "tests/*",
 ]
 
 [tool.ruff]

pyproject.toml

@@ -9,9 +9,9 @@
 from pyzeebe.client.sync_client import SyncZeebeClient
 from pyzeebe.job.job import Job, JobController
 from pyzeebe.job.job_status import JobStatus
+from pyzeebe.middlewares import BaseMiddleware
 from pyzeebe.task.exception_handler import ExceptionHandler, default_exception_handler
 from pyzeebe.task.task_config import TaskConfig
-from pyzeebe.task.types import TaskDecorator
 from pyzeebe.worker.task_router import ZeebeTaskRouter
 from pyzeebe.worker.worker import ZeebeWorker
 
@@ -28,8 +28,8 @@
     "JobStatus",
     "ExceptionHandler",
     "TaskConfig",
-    "TaskDecorator",
     "ZeebeTaskRouter",
     "default_exception_handler",
     "ZeebeWorker",
+    "BaseMiddleware",
 )

pyzeebe/__init__.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
 
 from pyzeebe.job.job_status import JobStatus
 from pyzeebe.types import Headers, Variables
@@ -27,9 +27,9 @@ class Job:
     variables: Variables
     tenant_id: str | None = None
     status: JobStatus = JobStatus.Running
-    task_result = None
+    task_result: Variables | None = None
 
-    def set_task_result(self, task_result: Any) -> None:
+    def set_task_result(self, task_result: Variables) -> None:
         object.__setattr__(self, "task_result", task_result)
 
     def _set_status(self, value: JobStatus) -> None:
@@ -46,11 +46,9 @@ def __init__(self, job: Job, zeebe_adapter: ZeebeAdapter) -> None:
         self._job = job
         self._zeebe_adapter = zeebe_adapter
 
-    async def set_running_after_decorators_status(self) -> None:
-        """
-        RunningAfterDecorators status means that the task has been completed as intended and the after decorators will now run.
-        """
-        self._job._set_status(JobStatus.RunningAfterDecorators)
+    @property
+    def job(self) -> Job:
+        return self._job
 
     async def set_success_status(self, variables: Variables | None = None) -> None:
         """

pyzeebe/job/job.py

@@ -4,6 +4,5 @@
 class JobStatus(Enum):
     Running = "Running"
     Completed = "Completed"
-    RunningAfterDecorators = "RunningAfterDecorators"
     Failed = "Failed"
     ErrorThrown = "ErrorThrown"

pyzeebe/job/job_status.py

@@ -0,0 +1,4 @@
+from .base import BaseMiddleware as BaseMiddleware
+from .base import ExceptionMiddleware as ExceptionMiddleware
+from .types import ConsumeMiddlewareStack as ConsumeMiddlewareStack
+from .types import ExecuteMiddlewareStack as ExecuteMiddlewareStack

pyzeebe/middlewares/__init__.py

@@ -0,0 +1,37 @@
+from __future__ import annotations
+
+import logging
+from collections.abc import Mapping
+
+from pyzeebe.job.job import Job, JobController
+from pyzeebe.task.exception_handler import ExceptionHandler
+
+from .types import ConsumeMiddlewareStack, ExecuteMiddlewareStack
+
+logger = logging.getLogger(__name__)
+
+
+class BaseMiddleware:
+    async def consume_scope(self, call_next: ConsumeMiddlewareStack, job: Job) -> Job:
+        return await call_next(job)
+
+    async def execute_scope(self, call_next: ExecuteMiddlewareStack, job: Job, job_controller: JobController) -> Job:
+        return await call_next(job, job_controller)
+
+
+class ExceptionMiddleware(BaseMiddleware):
+    def __init__(self, exception_handlers: Mapping[type[Exception], ExceptionHandler]) -> None:
+        assert Exception in exception_handlers
+
+        self.exception_handlers = exception_handlers
+
+    async def execute_scope(self, call_next: ExecuteMiddlewareStack, job: Job, job_controller: JobController) -> Job:
+        try:
+            job = await call_next(job, job_controller)
+        except Exception as err:
+            handler = self.exception_handlers.get(type(err), self.exception_handlers[Exception])
+            await handler(err, job, job_controller)
+        else:
+            await job_controller.set_success_status(job.task_result)
+
+        return job

pyzeebe/middlewares/base.py

@@ -0,0 +1,6 @@
+from collections.abc import Awaitable, Callable
+
+from pyzeebe.job.job import Job, JobController
+
+ConsumeMiddlewareStack = Callable[[Job], Awaitable[Job]]
+ExecuteMiddlewareStack = Callable[[Job, JobController], Awaitable[Job]]

pyzeebe/middlewares/types.py

@@ -1,24 +1,19 @@
 from __future__ import annotations
 
+from dataclasses import dataclass
 from typing import Any
 
 from pyzeebe.function_tools import Function
 from pyzeebe.task.task_config import TaskConfig
 from pyzeebe.task.types import JobHandler
 
 
+@dataclass()
 class Task:
-    def __init__(self, original_function: Function[..., Any], job_handler: JobHandler, config: TaskConfig) -> None:
-        self.original_function = original_function
-        self.job_handler = job_handler
-        self.config = config
+    original_function: Function[..., Any]
+    job_handler: JobHandler
+    config: TaskConfig
 
     @property
     def type(self) -> str:
         return self.config.type
-
-    def __repr__(self) -> str:
-        return (
-            f"Task(config= {self.config}, original_function={self.original_function}, "
-            f"job_handler={self.job_handler})"
-        )

pyzeebe/task/task.py

@@ -3,7 +3,6 @@
 import functools
 import inspect
 import logging
-from collections.abc import Sequence
 from typing import Any, TypeVar
 
 from typing_extensions import ParamSpec
@@ -14,10 +13,9 @@
 from pyzeebe.function_tools.dict_tools import convert_to_dict_function
 from pyzeebe.function_tools.parameter_tools import get_job_parameter_name
 from pyzeebe.job.job import JobController
-from pyzeebe.task.exception_handler import default_exception_handler
 from pyzeebe.task.task import Task
 from pyzeebe.task.task_config import TaskConfig
-from pyzeebe.task.types import AsyncTaskDecorator, DecoratorRunner, JobHandler
+from pyzeebe.task.types import JobHandler
 from pyzeebe.types import Variables
 
 P = ParamSpec("P")
@@ -34,20 +32,10 @@ def build_task(task_function: Function[..., Any], task_config: TaskConfig) -> Ta
 def build_job_handler(task_function: Function[..., Any], task_config: TaskConfig) -> JobHandler:
     prepared_task_function = prepare_task_function(task_function, task_config)
 
-    before_decorator_runner = create_decorator_runner(task_config.before)
-    after_decorator_runner = create_decorator_runner(task_config.after)
-
     @functools.wraps(task_function)
     async def job_handler(job: Job, job_controller: JobController) -> Job:
-        job = await before_decorator_runner(job)
-        return_variables, succeeded = await run_original_task_function(
-            prepared_task_function, task_config, job, job_controller
-        )
+        return_variables = await run_original_task_function(prepared_task_function, task_config, job)
         job.set_task_result(return_variables)
-        await job_controller.set_running_after_decorators_status()
-        job = await after_decorator_runner(job)
-        if succeeded:
-            await job_controller.set_success_status(variables=return_variables)
         return job
 
     return job_handler
@@ -63,38 +51,30 @@ def prepare_task_function(task_function: Function[P, R], task_config: TaskConfig
     return task_function  # type: ignore[return-value]
 
 
-async def run_original_task_function(
-    task_function: DictFunction[...], task_config: TaskConfig, job: Job, job_controller: JobController
-) -> tuple[Variables, bool]:
-    try:
-        if task_config.variables_to_fetch is None:
-            variables: dict[str, Any] = {}
-        elif task_wants_all_variables(task_config):
-            if only_job_is_required_in_task_function(task_function):
-                variables = {}
-            else:
-                variables = {**job.variables}
+async def run_original_task_function(task_function: DictFunction[...], task_config: TaskConfig, job: Job) -> Variables:
+    if task_config.variables_to_fetch is None:
+        variables: dict[str, Any] = {}
+    elif task_wants_all_variables(task_config):
+        if only_job_is_required_in_task_function(task_function):
+            variables = {}
         else:
-            variables = {
-                k: v
-                for k, v in job.variables.items()
-                if k in task_config.variables_to_fetch or k == task_config.job_parameter_name
-            }
+            variables = {**job.variables}
+    else:
+        variables = {
+            k: v
+            for k, v in job.variables.items()
+            if k in task_config.variables_to_fetch or k == task_config.job_parameter_name
+        }
 
-        if task_config.job_parameter_name:
-            variables[task_config.job_parameter_name] = job
+    if task_config.job_parameter_name:
+        variables[task_config.job_parameter_name] = job
 
-        returned_value = await task_function(**variables)
+    returned_value = await task_function(**variables)
 
-        if returned_value is None:
-            returned_value = {}
+    if returned_value is None:
+        returned_value = {}
 
-        return returned_value, True
-    except Exception as e:
-        logger.debug("Failed job: %s. Error: %s.", job, e)
-        exception_handler = task_config.exception_handler or default_exception_handler
-        await exception_handler(e, job, job_controller)
-        return job.variables, False
+    return returned_value
 
 
 def only_job_is_required_in_task_function(task_function: DictFunction[...]) -> bool:
@@ -104,20 +84,3 @@ def only_job_is_required_in_task_function(task_function: DictFunction[...]) -> b
 
 def task_wants_all_variables(task_config: TaskConfig) -> bool:
     return task_config.variables_to_fetch == []
-
-
-def create_decorator_runner(decorators: Sequence[AsyncTaskDecorator]) -> DecoratorRunner:
-    async def decorator_runner(job: Job) -> Job:
-        for decorator in decorators:
-            job = await run_decorator(decorator, job)
-        return job
-
-    return decorator_runner
-
-
-async def run_decorator(decorator: AsyncTaskDecorator, job: Job) -> Job:
-    try:
-        return await decorator(job)
-    except Exception as e:
-        logger.warning("Failed to run decorator %s. Exception: %s", decorator, e, exc_info=True)
-        return job