fix(producer): TaskProducer can produce to partitions

clients/python/src/examples/tasks.py

@@ -140,6 +140,6 @@ def producer_factory() -> KafkaProducer:
     for i in range(production_count):
         logger.debug(f"Producing message {i} onto topic {destination_topic}...")
         producer.produce(
-            topic=Topic(destination_topic),
+            dest=Topic(destination_topic),
             payload=KafkaPayload(key=None, value=payload, headers=[]),
         )

clients/python/src/taskbroker_client/types.py

@@ -5,7 +5,7 @@
 
 from arroyo.backends.abstract import ProducerFuture
 from arroyo.backends.kafka import KafkaPayload
-from arroyo.types import BrokerValue, Topic
+from arroyo.types import BrokerValue, Partition, Topic
 from sentry_protos.taskbroker.v1.taskbroker_pb2 import TaskActivation, TaskActivationStatus
 
 TaskHeaders = dict[str, str]
@@ -38,7 +38,7 @@ class ProducerProtocol(Protocol):
     """Interface for producers that tasks depend on."""
 
     def produce(
-        self, topic: Topic, payload: KafkaPayload
+        self, dest: Topic | Partition, payload: KafkaPayload
     ) -> ProducerFuture[BrokerValue[KafkaPayload]]: ...
 
 

clients/python/src/taskbroker_client/worker/producer.py

@@ -5,7 +5,7 @@
 
 from arroyo.backends.abstract import ProducerFuture, SimpleProducerFuture
 from arroyo.backends.kafka import KafkaPayload
-from arroyo.types import BrokerValue, Topic
+from arroyo.types import BrokerValue, Partition, Topic
 
 from taskbroker_client.constants import TASK_PRODUCER_MAX_PENDING_FUTURES
 from taskbroker_client.metrics import MetricsBackend, NoOpMetricsBackend
@@ -68,7 +68,7 @@ def collect_futures() -> set[ProducerFuture[BrokerValue[KafkaPayload]]]:
 
     def produce(
         self,
-        topic: Topic,
+        dest: Topic | Partition,
         payload: KafkaPayload,
         callbacks: Sequence[Callable[[Future[BrokerValue[KafkaPayload]]], Any]] = [],
     ) -> None:
@@ -79,12 +79,12 @@ def produce(
         to the future via the `callbacks` arg.
 
         Args:
-            topic: Topic to produce to.
+            dest: Topic (or specific partition) to produce to.
             payload: KafkaPayload to produce.
             callbacks: List of Callables to add to the future as done callbacks. The future itself
                        is the only arg passed to the callback.
         """
-        future = self._get().produce(topic, payload)
+        future = self._get().produce(dest, payload)
         self.track_future(future)
         if callbacks:
             # Arroyo producers can return a SimpleProducerFuture,