Header in serializing context (confluentinc#1120)

ldom · web-flow · commit e1c3227e46c2 · 2021-05-18T14:40:55.000-07:00
SerializationContext doesn't support message headers before, add headers support to SerializationContext.
diff --git a/src/confluent_kafka/deserializing_consumer.py b/src/confluent_kafka/deserializing_consumer.py
@@ -130,7 +130,7 @@ def poll(self, timeout=-1):
         if msg.error() is not None:
             raise ConsumeError(msg.error(), kafka_message=msg)
 
-        ctx = SerializationContext(msg.topic(), MessageField.VALUE)
+        ctx = SerializationContext(msg.topic(), MessageField.VALUE, msg.headers())
         value = msg.value()
         if self._value_deserializer is not None:
             try:
diff --git a/src/confluent_kafka/serialization/__init__.py b/src/confluent_kafka/serialization/__init__.py
@@ -56,10 +56,13 @@ class SerializationContext(object):
         field (MessageField): Describes what part of the message is
             being serialized.
 
+        headers (list): List of message header tuples. Defaults to None.
+
     """
-    def __init__(self, topic, field):
+    def __init__(self, topic, field, headers=None):
         self.topic = topic
         self.field = field
+        self.headers = headers
 
 
 class SerializationError(KafkaException):
diff --git a/src/confluent_kafka/serializing_producer.py b/src/confluent_kafka/serializing_producer.py
@@ -160,7 +160,7 @@ def produce(self, topic, key=None, value=None, partition=-1,
             KafkaException: For all other errors
 
         """
-        ctx = SerializationContext(topic, MessageField.KEY)
+        ctx = SerializationContext(topic, MessageField.KEY, headers)
         if self._key_serializer is not None:
             try:
                 key = self._key_serializer(key, ctx)
diff --git a/tests/README.md b/tests/README.md
@@ -18,7 +18,7 @@ A python3 env suitable for running tests:
 
     $ python3 -m venv venv_test
     $ source venv_test/bin/activate
-    $ pip install -r test/requirements.txt
+    $ pip install -r tests/requirements.txt
     $ python setup.py build
     $ python setup.py install
 
diff --git a/tests/integration/schema_registry/test_avro_serializers.py b/tests/integration/schema_registry/test_avro_serializers.py
@@ -124,7 +124,7 @@ def test_delivery_report_serialization(kafka_cluster, load_file, avsc, data, rec
 
     def assert_cb(err, msg):
         actual = value_deserializer(msg.value(),
-                                    SerializationContext(topic, MessageField.VALUE))
+                                    SerializationContext(topic, MessageField.VALUE, msg.headers()))
 
         if record_type == "record":
             assert [v == actual[k] for k, v in data.items()]
diff --git a/tests/schema_registry/test_avro_serializer.py b/tests/schema_registry/test_avro_serializer.py
@@ -111,9 +111,11 @@ def test_avro_serializer_record_subject_name_strategy(load_avsc):
                                      conf={'subject.name.strategy':
                                            record_subject_name_strategy})
 
-    ctx = SerializationContext('test_subj', MessageField.VALUE)
+    ctx = SerializationContext('test_subj', MessageField.VALUE, [])
     assert test_serializer._subject_name_func(ctx,
                                               test_serializer._schema_name) == 'python.test.basic'
+    assert ctx is not None
+    assert not ctx.headers
 
 
 def test_avro_serializer_record_subject_name_strategy_primitive(load_avsc):
@@ -127,9 +129,10 @@ def test_avro_serializer_record_subject_name_strategy_primitive(load_avsc):
                                      conf={'subject.name.strategy':
                                            record_subject_name_strategy})
 
-    ctx = SerializationContext('test_subj', MessageField.VALUE)
+    ctx = SerializationContext('test_subj', MessageField.VALUE, [('header1', 'header value 1'), ])
     assert test_serializer._subject_name_func(ctx,
                                               test_serializer._schema_name) == 'int'
+    assert ('header1', 'header value 1') in ctx.headers
 
 
 def test_avro_serializer_topic_record_subject_name_strategy(load_avsc):
@@ -162,6 +165,8 @@ def test_avro_serializer_topic_record_subject_name_strategy_primitive(load_avsc)
     ctx = SerializationContext('test_subj', MessageField.VALUE)
     assert test_serializer._subject_name_func(
         ctx, test_serializer._schema_name) == 'test_subj-int'
+    assert ctx is not None
+    assert ctx.headers is None
 
 
 def test_avro_serializer_subject_name_strategy_default(load_avsc):
