From 389d707cdc5c1cf676f9f6cf89af53b72ab4779e Mon Sep 17 00:00:00 2001 From: Rajashree Mandaogane Date: Fri, 17 Feb 2023 16:08:11 -0800 Subject: [PATCH] modify kafka context getter helper methods to work on dict and list --- .../instrumentation/confluent_kafka/utils.py | 14 ++++++++++++-- .../tests/test_instrumentation.py | 15 +++++++++++++++ 2 files changed, 27 insertions(+), 2 deletions(-) diff --git a/instrumentation/opentelemetry-instrumentation-confluent-kafka/src/opentelemetry/instrumentation/confluent_kafka/utils.py b/instrumentation/opentelemetry-instrumentation-confluent-kafka/src/opentelemetry/instrumentation/confluent_kafka/utils.py index 89fad19c1b..77fce03cd8 100644 --- a/instrumentation/opentelemetry-instrumentation-confluent-kafka/src/opentelemetry/instrumentation/confluent_kafka/utils.py +++ b/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): diff --git a/instrumentation/opentelemetry-instrumentation-confluent-kafka/tests/test_instrumentation.py b/instrumentation/opentelemetry-instrumentation-confluent-kafka/tests/test_instrumentation.py index aeccd96020..cc78258328 100644 --- a/instrumentation/opentelemetry-instrumentation-confluent-kafka/tests/test_instrumentation.py +++ b/instrumentation/opentelemetry-instrumentation-confluent-kafka/tests/test_instrumentation.py @@ -25,6 +25,7 @@ ) from opentelemetry.instrumentation.confluent_kafka.utils import ( KafkaContextSetter, + KafkaContextGetter, ) @@ -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))