From d1f76e3074c9a452e1c6fd80fd28b49adda07228 Mon Sep 17 00:00:00 2001 From: Tony Le Date: Fri, 3 Jul 2026 12:49:20 -0400 Subject: [PATCH] feat(taskbroker-client): expose task enqueue producer future --- .../python/src/taskbroker_client/registry.py | 8 ++- clients/python/src/taskbroker_client/task.py | 58 +++++++++++++++---- clients/python/tests/test_registry.py | 6 +- clients/python/tests/test_task.py | 21 +++++++ 4 files changed, 80 insertions(+), 13 deletions(-) diff --git a/clients/python/src/taskbroker_client/registry.py b/clients/python/src/taskbroker_client/registry.py index 8e8da7ce..63e8ae4e 100644 --- a/clients/python/src/taskbroker_client/registry.py +++ b/clients/python/src/taskbroker_client/registry.py @@ -4,7 +4,7 @@ import logging from collections.abc import Callable from concurrent import futures -from typing import Any +from typing import Any, cast import sentry_sdk from arroyo.backends.kafka import KafkaPayload @@ -167,7 +167,9 @@ def _handle_produce_future(self, future: ProducerFuture, tags: dict[str, str]) - else: self.metrics.incr("taskworker.registry.send_task.success", tags=tags) - def send_task(self, activation: TaskActivation, wait_for_delivery: bool = False) -> None: + def send_task( + self, activation: TaskActivation, wait_for_delivery: bool = False + ) -> ProducerFuture: topic = self.topic with sentry_sdk.start_span( @@ -210,6 +212,8 @@ def send_task(self, activation: TaskActivation, wait_for_delivery: bool = False) except Exception: logger.exception("Failed to wait for delivery") + return cast(ProducerFuture, produce_future) + def _producer(self, topic: str) -> ProducerProtocol: if topic not in self._producers: self._producers[topic] = self._producer_factory(topic) diff --git a/clients/python/src/taskbroker_client/task.py b/clients/python/src/taskbroker_client/task.py index 651e7d45..ba84762d 100644 --- a/clients/python/src/taskbroker_client/task.py +++ b/clients/python/src/taskbroker_client/task.py @@ -27,7 +27,7 @@ from taskbroker_client.retry import Retry if TYPE_CHECKING: - from taskbroker_client.registry import TaskNamespace + from taskbroker_client.registry import ProducerFuture, TaskNamespace ALWAYS_EAGER = False @@ -150,21 +150,14 @@ def delay(self, *args: P.args, **kwargs: P.kwargs) -> None: """ self.apply_async(args=args, kwargs=kwargs) - def apply_async( + def _create_activation_for_async_dispatch( self, args: Any | None = None, kwargs: Any | None = None, headers: MutableMapping[str, Any] | None = None, expires: int | datetime.timedelta | None = None, countdown: int | datetime.timedelta | None = None, - **options: Any, - ) -> None: - """ - Schedule a task to run later with a set of arguments. - - The provided parameters will be msgpack encoded and stored within - a `TaskActivation` protobuf that is appended to kafka. - """ + ) -> tuple[TaskActivation, Any, Any]: if args is None: args = [] if kwargs is None: @@ -177,6 +170,26 @@ def apply_async( activation = self.create_activation( args=args, kwargs=kwargs, headers=headers, expires=expires, countdown=countdown ) + return activation, args, kwargs + + def apply_async( + self, + args: Any | None = None, + kwargs: Any | None = None, + headers: MutableMapping[str, Any] | None = None, + expires: int | datetime.timedelta | None = None, + countdown: int | datetime.timedelta | None = None, + **options: Any, + ) -> None: + """ + Schedule a task to run later with a set of arguments. + + The provided parameters will be msgpack encoded and stored within + a `TaskActivation` protobuf that is appended to kafka. + """ + activation, args, kwargs = self._create_activation_for_async_dispatch( + args=args, kwargs=kwargs, headers=headers, expires=expires, countdown=countdown + ) if ALWAYS_EAGER: self._call_func(*args, **kwargs) else: @@ -185,6 +198,31 @@ def apply_async( wait_for_delivery=self.wait_for_delivery, ) + def apply_async_with_future( + self, + args: Any | None = None, + kwargs: Any | None = None, + headers: MutableMapping[str, Any] | None = None, + expires: int | datetime.timedelta | None = None, + countdown: int | datetime.timedelta | None = None, + **options: Any, + ) -> ProducerFuture | None: + """ + Schedule a task and return the producer future for callers that need to + wait on delivery alongside other Kafka produces. + """ + activation, args, kwargs = self._create_activation_for_async_dispatch( + args=args, kwargs=kwargs, headers=headers, expires=expires, countdown=countdown + ) + if ALWAYS_EAGER: + self._call_func(*args, **kwargs) + return None + + return self._namespace.send_task( + activation, + wait_for_delivery=self.wait_for_delivery, + ) + def apply_async_testing( self, args: Any | None = None, diff --git a/clients/python/tests/test_registry.py b/clients/python/tests/test_registry.py index ab21e7ec..e691e55a 100644 --- a/clients/python/tests/test_registry.py +++ b/clients/python/tests/test_registry.py @@ -1,4 +1,5 @@ from concurrent.futures import Future +from typing import Any from unittest.mock import Mock import msgpack @@ -141,7 +142,10 @@ def simple_task() -> None: mock_producer = Mock() namespace._producers["taskworker"] = mock_producer - namespace.send_task(activation) + ret_value: Future[Any] = Future() + mock_producer.produce.return_value = ret_value + + assert namespace.send_task(activation) is ret_value assert mock_producer.produce.call_count == 1 mock_call = mock_producer.produce.call_args diff --git a/clients/python/tests/test_task.py b/clients/python/tests/test_task.py index 166cf7f2..2556b9e7 100644 --- a/clients/python/tests/test_task.py +++ b/clients/python/tests/test_task.py @@ -1,6 +1,7 @@ import contextlib import datetime from collections.abc import MutableMapping +from concurrent.futures import Future from typing import Any from unittest.mock import patch @@ -152,6 +153,26 @@ def test_func(mixed: Any) -> None: assert calls[1] == {"mixed": "str"} +def test_apply_async_with_future_returns_producer_future(task_namespace: TaskNamespace) -> None: + def test_func(*args: Any, **kwargs: Any) -> None: + pass + + task = Task( + name="test.test_func", + func=test_func, + namespace=task_namespace, + ) + producer_future: Future[Any] = Future() + + with patch.object(task_namespace, "send_task", return_value=producer_future) as mock_send: + assert task.apply_async_with_future(args=["arg2"], kwargs={"org_id": 2}) is producer_future + + assert mock_send.call_count == 1 + activation = mock_send.call_args.args[0] + expected_params = {"args": ["arg2"], "kwargs": {"org_id": 2}} + assert msgpack.unpackb(activation.parameters_bytes, raw=False) == expected_params + + def test_should_retry(task_namespace: TaskNamespace) -> None: retry = Retry(times=3, times_exceeded=LastAction.Deadletter) state = retry.initial_state()