feat(taskbroker-client): expose task enqueue producer future

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)

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,

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

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()