modify kafka context getter helper methods to work on dict and list

Author: mrajashree

CHANGELOG.md

@@ -7,9 +7,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+### Added
+
 - Add connection attributes to sqlalchemy connect span
   ([#1608](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1608))
 
+### Fixed
+
+- Fix confluent-kafka instrumentation by allowing Producer headers to be dict or list
+  ([#1655](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1655))
+
 ## Version 1.16.0/0.37b0 (2023-02-17)
 
 ### Added
@@ -54,8 +61,6 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   ([#1512](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1512))
 - Add commit method for ConfluentKafkaInstrumentor's ProxiedConsumer
   ([#1656](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1656))
-- Fix confluent-kafka instrumentation by allowing Producer headers to be dict or list
-  ([#1655](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1655))
 
 ## Version 1.15.0/0.36b0 (2022-12-10)
 

instrumentation/opentelemetry-instrumentation-confluent-kafka/src/opentelemetry/instrumentation/confluent_kafka/utils.py

@@ -41,16 +41,26 @@ class KafkaContextGetter(textmap.Getter):
     def get(self, carrier: textmap.CarrierT, key: str) -> Optional[List[str]]:
         if carrier is None:
             return None
-        for item_key, value in carrier:
+
+        carrier_items = carrier
+        if isinstance(carrier, dict):
+            carrier_items = carrier.items()
+
+        for item_key, value in carrier_items:
             if item_key == key:
                 if value is not None:
                     return [value.decode()]
+
         return None
 
     def keys(self, carrier: textmap.CarrierT) -> List[str]:
         if carrier is None:
             return []
-        return [key for (key, value) in carrier]
+
+        carrier_items = carrier
+        if isinstance(carrier, dict):
+            carrier_items = carrier.items()
+        return [key for (key, value) in carrier_items]
 
 
 class KafkaContextSetter(textmap.Setter):

instrumentation/opentelemetry-instrumentation-confluent-kafka/tests/test_instrumentation.py

@@ -24,6 +24,7 @@
     ProxiedProducer,
 )
 from opentelemetry.instrumentation.confluent_kafka.utils import (
+    KafkaContextGetter,
     KafkaContextSetter,
 )
 
@@ -89,3 +90,17 @@ def test_context_setter(self) -> None:
         carrier_list = [("key1", "val1")]
         context_setter.set(carrier_list, "key2", "val2")
         self.assertTrue(("key2", "val2".encode()) in carrier_list)
+
+    def test_context_getter(self) -> None:
+        context_setter = KafkaContextSetter()
+        context_getter = KafkaContextGetter()
+
+        carrier_dict = {}
+        context_setter.set(carrier_dict, "key1", "val1")
+        self.assertEqual(context_getter.get(carrier_dict, "key1"), ["val1"])
+        self.assertEqual(["key1"], context_getter.keys(carrier_dict))
+
+        carrier_list = []
+        context_setter.set(carrier_list, "key1", "val1")
+        self.assertEqual(context_getter.get(carrier_list, "key1"), ["val1"])
+        self.assertEqual(["key1"], context_getter.keys(carrier_list))