opensource/opentelemetry-python-contrib
1
0
Fork 0
mirror of https://github.com/open-telemetry/opentelemetry-python-contrib.git synced 2025-08-01 09:13:23 +08:00
Code Issues Packages Projects Releases Wiki Activity

Allow Kafka producer headers to be dict or list (#1655)

Browse Source 
* Allow Kafka producer headers to be dict or list

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

---------

Co-authored-by: Shalev Roda <65566801+shalevr@users.noreply.github.com>
Co-authored-by: Srikanth Chekuri <srikanth.chekuri92@gmail.com>
This commit is contained in:
Rajashree Mandaogane
2023-03-10 04:59:05 -08:00
committed by GitHub
parent 419975138b
commit 88783f9632
3 changed files with 57 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
instrumentation/opentelemetry-instrumentation-confluent-kafka/src/opentelemetry/instrumentation/confluent_kafka/utils.py
View File

@ -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):
@ -60,7 +70,12 @@ class KafkaContextSetter(textmap.Setter):
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()

31
instrumentation/opentelemetry-instrumentation-confluent-kafka/tests/test_instrumentation.py
View File

@ -23,6 +23,10 @@ from opentelemetry.instrumentation.confluent_kafka import (
ProxiedConsumer,
ProxiedProducer,
)
from opentelemetry.instrumentation.confluent_kafka.utils import (
KafkaContextGetter,
KafkaContextSetter,
)
class TestConfluentKafka(TestCase):
@ -73,3 +77,30 @@ class TestConfluentKafka(TestCase):
consumer = instrumentation.instrument_consumer(consumer)
self.assertEqual(consumer.__class__, ProxiedConsumer)
self.assertTrue(hasattr(consumer, "commit"))
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)
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))