Allow Kafka producer headers to be dict or list

open-telemetry · Feb 10, 2023 · 61892c1 · 61892c1
1 parent df32e8c
commit 61892c1
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 1 deletion.
diff --git a/...nstrumentation-confluent-kafka/src/opentelemetry/instrumentation/confluent_kafka/utils.py b/...nstrumentation-confluent-kafka/src/opentelemetry/instrumentation/confluent_kafka/utils.py
@@ -60,7 +60,12 @@ def set(self, carrier: textmap.CarrierT, key: str, value: str) -> None:
 
         if value:
             value = value.encode()
-        carrier.append((key, value))
+
+        if isinstance(carrier, list):
+            carrier.append((key, value))
+
+        if isinstance(carrier, dict):
+            carrier[key] = value
 
 
 _kafka_getter = KafkaContextGetter()

diff --git a/instrumentation/opentelemetry-instrumentation-confluent-kafka/tests/test_instrumentation.py b/instrumentation/opentelemetry-instrumentation-confluent-kafka/tests/test_instrumentation.py
@@ -24,6 +24,8 @@
     ProxiedProducer,
 )
 
+from opentelemetry.instrumentation.confluent_kafka.utils import KafkaContextSetter
+
 
 class TestConfluentKafka(TestCase):
     def test_instrument_api(self) -> None:
@@ -58,3 +60,14 @@ def test_instrument_api(self) -> None:
 
         consumer = instrumentation.uninstrument_consumer(consumer)
         self.assertEqual(consumer.__class__, Consumer)
+
+    def test_context_setter(self) -> None:
+        context_setter = KafkaContextSetter()
+
+        carrier_dict = {"key1": "val1"}
+        context_setter.set(carrier_dict, "key2", "val2")
+        self.assertGreaterEqual(carrier_dict.items(), {"key2": "val2".encode()}.items())
+
+        carrier_list = [("key1", "val1")]
+        context_setter.set(carrier_list, "key2", "val2")
+        self.assertTrue(("key2", "val2".encode()) in carrier_list)