Add support for sending messages with keys

Author: cnweaver

adc/producer.py

@@ -36,7 +36,8 @@ def write(self,
               msg: Union[bytes, 'Serializable'],
               headers: Optional[Union[dict, list]] = None,
               delivery_callback: Optional[DeliveryCallback] = log_delivery_errors,
-              topic: Optional[str] = None) -> None:
+              topic: Optional[str] = None,
+              key: Optional[Union[str, bytes]] = None) -> None:
         if isinstance(msg, Serializable):
             msg = msg.serialize()
         if topic is None:
@@ -48,10 +49,10 @@ def write(self,
                                 "or specify the topic argument to write()")
         self.logger.debug("writing message to %s", topic)
         if delivery_callback is not None:
-            self._producer.produce(topic, msg, headers=headers,
+            self._producer.produce(topic, msg, headers=headers, key=key,
                                    on_delivery=delivery_callback)
         else:
-            self._producer.produce(topic, msg, headers=headers)
+            self._producer.produce(topic, msg, headers=headers, key=key,)
 
     def flush(self, timeout: timedelta = timedelta(seconds=10)) -> int:
         """Attempt to flush enqueued messages. Return the number of messages still

tests/test_kafka_integration.py

@@ -3,7 +3,7 @@
 import time
 import unittest
 from datetime import datetime, timedelta
-from typing import List
+from typing import List, Optional
 
 import docker
 import pytest
@@ -59,6 +59,31 @@ def test_round_trip(self):
         self.assertEqual(msg.topic(), topic)
         self.assertEqual(msg.value(), b"can you hear me?")
 
+    def test_message_with_key(self):
+        """Try writing a message into the Kafka broker, and try pulling the same
+        message back out.
+
+        """
+        topic = "test_message_with_key"
+        # Push one message in...
+        simple_write_msg(self.kafka, topic, "can you hear me?", key="test_msg")
+        # ... and pull it back out.
+        consumer = adc.consumer.Consumer(adc.consumer.ConsumerConfig(
+            broker_urls=[self.kafka.address],
+            group_id="test_consumer",
+            auth=self.kafka.auth,
+        ))
+        consumer.subscribe(topic)
+        stream = consumer.stream()
+
+        msg = next(stream)
+        if msg.error() is not None:
+            raise Exception(msg.error())
+
+        self.assertEqual(msg.topic(), topic)
+        self.assertEqual(msg.value(), b"can you hear me?")
+        self.assertEqual(msg.key(), b"test_msg")
+
     def test_reset_to_end(self):
         # Write a few messages.
         topic = "test_reset_to_end"
@@ -569,13 +594,13 @@ def get_or_create_docker_network(self):
         return self.docker_client.networks.create(name="adc-integration-test")
 
 
-def simple_write_msg(conn: KafkaDockerConnection, topic: str, msg: str):
+def simple_write_msg(conn: KafkaDockerConnection, topic: str, msg: str, key: Optional[str] = None):
     producer = adc.producer.Producer(adc.producer.ProducerConfig(
         broker_urls=[conn.address],
         topic=topic,
         auth=conn.auth,
     ))
-    producer.write(msg)
+    producer.write(msg, key=key)
     producer.flush()