From 5e46189feef2d1152939d3a09641bbd4822473f8 Mon Sep 17 00:00:00 2001 From: Wei Chen <182598863+cw-jlu@users.noreply.github.com> Date: Thu, 12 Nov 2026 21:33:11 +0800 Subject: [PATCH] Propagate terminal pipeline errors after retry result suppression --- taskiq_pipelines/middleware.py | 38 ++------ tests/test_retries.py | 170 +++++++++++++++++++++++++++++++++ 2 files changed, 177 insertions(+), 31 deletions(-) create mode 100644 tests/test_retries.py diff --git a/taskiq_pipelines/middleware.py b/taskiq_pipelines/middleware.py index b4da54f..0fb89c6 100644 --- a/taskiq_pipelines/middleware.py +++ b/taskiq_pipelines/middleware.py @@ -21,7 +21,7 @@ async def post_save( # noqa: PLR0911 result: "TaskiqResult[Any]", ) -> None: """ - Handle post-execute event. + Handle saved task results. This is the heart of pipelines. Here we decide what to do next. @@ -29,11 +29,12 @@ async def post_save( # noqa: PLR0911 If the message have pipeline labels we can calculate our next step. + Only saved errors terminate the pipeline. Attempts suppressed by retry + middleware do not reach this hook. + :param message: current message. :param result: result of the execution. """ - if result.is_err: - return if CURRENT_STEP not in message.labels: return current_step_num = int(message.labels[CURRENT_STEP]) @@ -52,6 +53,9 @@ async def post_save( # noqa: PLR0911 if current_step_num + 1 >= len(steps_data): logger.debug("Pipeline is completed.") return + if result.is_err: + await self.fail_pipeline(steps_data[-1].task_id, result.error) + return next_step_data = steps_data[current_step_num + 1] try: next_step = parse_step( @@ -82,34 +86,6 @@ async def post_save( # noqa: PLR0911 return await self.fail_pipeline(steps_data[-1].task_id) - async def on_error( - self, - message: "TaskiqMessage", - result: "TaskiqResult[Any]", - exception: BaseException, - ) -> None: - """ - Handles on_error event. - - :param message: current message. - :param result: execution result. - :param exception: found exception. - """ - if CURRENT_STEP not in message.labels: - return - current_step_num = int(message.labels[CURRENT_STEP]) - if PIPELINE_DATA not in message.labels: - logger.warning("Pipeline data not found. Execution flow is broken.") - return - pipe_data = message.labels[PIPELINE_DATA] - try: - steps = pydantic.TypeAdapter(list[DumpedStep]).validate_json(pipe_data) - except ValueError: - return - if current_step_num == len(steps) - 1: - return - await self.fail_pipeline(steps[-1].task_id, result.error) - async def fail_pipeline( self, last_task_id: str, diff --git a/tests/test_retries.py b/tests/test_retries.py new file mode 100644 index 0000000..43d399b --- /dev/null +++ b/tests/test_retries.py @@ -0,0 +1,170 @@ +import asyncio +from typing import Any + +import pytest +from taskiq import ( + InMemoryBroker, + SimpleRetryMiddleware, + TaskiqMessage, + TaskiqMiddleware, + TaskiqResult, +) +from taskiq.exceptions import NoResultError + +from taskiq_pipelines import Pipeline, PipelineMiddleware + + +class AttemptObserver(TaskiqMiddleware): + """Observe an attempt after every middleware has processed its error.""" + + def __init__(self) -> None: + self.failed_attempt_finished = asyncio.Event() + + async def post_execute( + self, + message: TaskiqMessage, + result: TaskiqResult[Any], + ) -> None: + """Notify the test after the failing attempt has finished error hooks.""" + if result.is_err: + self.failed_attempt_finished.set() + + +@pytest.mark.parametrize("pipeline_first", [True, False]) +async def test_retry_does_not_publish_a_premature_pipeline_error( + pipeline_first: bool, +) -> None: + observer = AttemptObserver() + middlewares: list[TaskiqMiddleware] = [ + PipelineMiddleware(), + SimpleRetryMiddleware(default_retry_count=3), + ] + if not pipeline_first: + middlewares.reverse() + broker = InMemoryBroker().with_middlewares(observer, *middlewares) + retry_started = asyncio.Event() + release_retry = asyncio.Event() + attempts = 0 + following_calls: list[int] = [] + + @broker.task(retry_on_error=True) + async def flaky(value: int) -> int: + nonlocal attempts + attempts += 1 + if attempts == 1: + raise ValueError("temporary failure") + retry_started.set() + await release_retry.wait() + return value + 1 + + @broker.task + async def following(value: int) -> int: + following_calls.append(value) + return value * 2 + + try: + sent = await Pipeline(broker, flaky).call_next(following).kiq(1) + await asyncio.wait_for(observer.failed_attempt_finished.wait(), timeout=3) + await asyncio.wait_for(retry_started.wait(), timeout=3) + assert not await broker.result_backend.is_result_ready(sent.task_id) + assert following_calls == [] + + release_retry.set() + result = await asyncio.wait_for( + sent.wait_result(check_interval=0.01), + timeout=3, + ) + assert not result.is_err + assert result.return_value == 4 + assert attempts == 2 + assert following_calls == [2] + finally: + release_retry.set() + await broker.shutdown() + + +@pytest.mark.parametrize("pipeline_first", [True, False]) +@pytest.mark.parametrize("retry_on_error", [True, False]) +async def test_terminal_failure_still_fails_the_pipeline( + pipeline_first: bool, + retry_on_error: bool, +) -> None: + middlewares: list[TaskiqMiddleware] = [ + PipelineMiddleware(), + SimpleRetryMiddleware(default_retry_count=3), + ] + if not pipeline_first: + middlewares.reverse() + broker = InMemoryBroker().with_middlewares(*middlewares) + attempts = 0 + following_calls: list[int] = [] + + @broker.task(retry_on_error=retry_on_error) + async def failing(value: int) -> int: + nonlocal attempts + attempts += 1 + raise ValueError("terminal failure") + + @broker.task + async def following(value: int) -> int: + following_calls.append(value) + return value + + try: + sent = await Pipeline(broker, failing).call_next(following).kiq(1) + result = await asyncio.wait_for( + sent.wait_result(check_interval=0.01), + timeout=3, + ) + assert result.is_err + assert isinstance(result.error, ValueError) + assert str(result.error) == "terminal failure" + assert attempts == (3 if retry_on_error else 1) + assert following_calls == [] + finally: + await broker.shutdown() + + +async def test_no_result_does_not_fail_the_pipeline() -> None: + observer = AttemptObserver() + broker = InMemoryBroker().with_middlewares(observer, PipelineMiddleware()) + + @broker.task + async def suppressed() -> int: + raise NoResultError + + @broker.task + async def following(value: int) -> int: + return value + + try: + sent = await Pipeline(broker, suppressed).call_next(following).kiq() + await asyncio.wait_for(observer.failed_attempt_finished.wait(), timeout=3) + assert not await broker.result_backend.is_result_ready(sent.task_id) + finally: + await broker.shutdown() + + +async def test_last_step_failure_keeps_its_saved_result() -> None: + broker = InMemoryBroker().with_middlewares(PipelineMiddleware()) + + @broker.task + async def start() -> int: + return 1 + + @broker.task + async def failing(value: int) -> int: + raise ValueError("last step failed") + + try: + sent = await Pipeline(broker, start).call_next(failing).kiq() + result = await asyncio.wait_for( + sent.wait_result(check_interval=0.01), + timeout=3, + ) + assert result.is_err + assert isinstance(result.error, ValueError) + assert str(result.error) == "last step failed" + assert result.labels["_pipe_current_step"] == 1 + finally: + await broker.shutdown()