fix(kafka): fix bug when setting message.key span tag for kafka serialization classes [backport 2.3] (#7793)

Backport 0e3a293 from #7725 to 2.3.

## Description

Fixes #7674 . Sets `SerializedProducer` span tag for `message.key` by
serializing key to bytes. Does not set the tag for a
`DeserializedConsumer` unless the key is a `str` or `bytes` since we do
not have the serialization function

## Checklist

- [x] Change(s) are motivated and described in the PR description.
- [x] Testing strategy is described if automated tests are not included
in the PR.
- [x] Risk is outlined (performance impact, potential for breakage,
maintainability, etc).
- [x] Change is maintainable (easy to change, telemetry, documentation).
- [x] [Library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html)
are followed. If no release note is required, add label
`changelog/no-changelog`.
- [x] Documentation is included (in-code, generated user docs, [public
corp docs](https://github.com/DataDog/documentation/)).
- [x] Backport labels are set (if
[applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting))

## Reviewer Checklist

- [x] Title is accurate.
- [x] No unnecessary changes are introduced.
- [x] Description motivates each change.
- [x] Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes unless absolutely necessary.
- [x] Testing strategy adequately addresses listed risk(s).
- [x] Change is maintainable (easy to change, telemetry, documentation).
- [x] Release note makes sense to a user of the library.
- [x] Reviewer has explicitly acknowledged and discussed the performance
implications of this PR as reported in the benchmarks PR comment.
- [x] Backport labels are set in a manner that is consistent with the
[release branch maintenance
policy](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)
- [x] If this PR touches code that signs or publishes builds or
packages, or handles credentials of any kind, I've requested a review
from `@DataDog/security-design-and-guidance`.
- [x] This PR doesn't touch any of that.

Co-authored-by: William Conti <[email protected]>
Co-authored-by: Emmett Butler <[email protected]>

Author: 3 people

ddtrace/contrib/kafka/patch.py

@@ -9,15 +9,16 @@
 from ddtrace.ext import SpanTypes
 from ddtrace.ext import kafka as kafkax
 from ddtrace.internal import core
-from ddtrace.internal.compat import ensure_text
 from ddtrace.internal.constants import COMPONENT
 from ddtrace.internal.constants import MESSAGING_SYSTEM
+from ddtrace.internal.logger import get_logger
 from ddtrace.internal.schema import schematize_messaging_operation
 from ddtrace.internal.schema import schematize_service_name
 from ddtrace.internal.schema.span_attribute_schema import SpanDirection
 from ddtrace.internal.utils import ArgumentError
 from ddtrace.internal.utils import get_argument_value
 from ddtrace.internal.utils.formats import asbool
+from ddtrace.internal.utils.version import parse_version
 from ddtrace.pin import Pin
 
 
@@ -29,6 +30,9 @@
 )
 
 
+log = get_logger(__name__)
+
+
 config._add(
     "kafka",
     dict(
@@ -42,6 +46,13 @@ def get_version():
     return getattr(confluent_kafka, "__version__", "")
 
 
+KAFKA_VERSION_TUPLE = parse_version(get_version())
+
+
+_SerializationContext = confluent_kafka.serialization.SerializationContext if KAFKA_VERSION_TUPLE >= (1, 4, 0) else None
+_MessageField = confluent_kafka.serialization.MessageField if KAFKA_VERSION_TUPLE >= (1, 4, 0) else None
+
+
 class TracedProducerMixin:
     def __init__(self, config, *args, **kwargs):
         super(TracedProducerMixin, self).__init__(config, *args, **kwargs)
@@ -139,6 +150,7 @@ def traced_produce(func, instance, args, kwargs):
         value = None
     message_key = kwargs.get("key", "") or ""
     partition = kwargs.get("partition", -1)
+    headers = get_argument_value(args, kwargs, 6, "headers", optional=True) or {}
     with pin.tracer.trace(
         schematize_messaging_operation(kafkax.PRODUCE, provider="kafka", direction=SpanDirection.OUTBOUND),
         service=trace_utils.ext_service(pin, config.kafka),
@@ -149,7 +161,14 @@ def traced_produce(func, instance, args, kwargs):
         span.set_tag_str(COMPONENT, config.kafka.integration_name)
         span.set_tag_str(SPAN_KIND, SpanKind.PRODUCER)
         span.set_tag_str(kafkax.TOPIC, topic)
-        span.set_tag_str(kafkax.MESSAGE_KEY, ensure_text(message_key, errors="replace"))
+
+        if _SerializingProducer is not None and isinstance(instance, _SerializingProducer):
+            serialized_key = serialize_key(instance, topic, message_key, headers)
+            if serialized_key is not None:
+                span.set_tag_str(kafkax.MESSAGE_KEY, serialized_key)
+        else:
+            span.set_tag_str(kafkax.MESSAGE_KEY, message_key)
+
         span.set_tag(kafkax.PARTITION, partition)
         span.set_tag_str(kafkax.TOMBSTONE, str(value is None))
         span.set_tag(SPAN_MEASURED_KEY)
@@ -184,7 +203,15 @@ def traced_poll(func, instance, args, kwargs):
             message_key = message.key() or ""
             message_offset = message.offset() or -1
             span.set_tag_str(kafkax.TOPIC, message.topic())
-            span.set_tag_str(kafkax.MESSAGE_KEY, ensure_text(message_key, errors="replace"))
+
+            # If this is a deserializing consumer, do not set the key as a tag since we
+            # do not have the serialization function
+            if (
+                (_DeserializingConsumer is not None and not isinstance(instance, _DeserializingConsumer))
+                or isinstance(message_key, str)
+                or isinstance(message_key, bytes)
+            ):
+                span.set_tag_str(kafkax.MESSAGE_KEY, message_key)
             span.set_tag(kafkax.PARTITION, message.partition())
             span.set_tag_str(kafkax.TOMBSTONE, str(len(message) == 0))
             span.set_tag(kafkax.MESSAGE_OFFSET, message_offset)
@@ -203,3 +230,18 @@ def traced_commit(func, instance, args, kwargs):
     core.dispatch("kafka.commit.start", [instance, args, kwargs])
 
     return func(*args, **kwargs)
+
+
+def serialize_key(instance, topic, key, headers):
+    if _SerializationContext is not None and _MessageField is not None:
+        ctx = _SerializationContext(topic, _MessageField.KEY, headers)
+        if hasattr(instance, "_key_serializer") and instance._key_serializer is not None:
+            try:
+                key = instance._key_serializer(key, ctx)
+                return key
+            except Exception:
+                log.debug("Failed to set Kafka Consumer key tag: %s", str(key))
+                return None
+        else:
+            log.warning("Failed to set Kafka Consumer key tag, no method available to serialize key: %s", str(key))
+            return None

releasenotes/notes/fix-bug-for-setting-message-key-span-tag-from-kafka-serialization-classes-630b09004a97063a.yaml

@@ -0,0 +1,4 @@
+---
+fixes:
+  - |
+    kafka: Resolves ``TypeError`` raised by serializing producers and deserializing consumers when the ``message.key`` tag is set on spans.

tests/contrib/kafka/test_kafka.py

@@ -1,3 +1,4 @@
+import json
 import logging
 import os
 import time
@@ -740,3 +741,61 @@ def test_span_has_dsm_payload_hash(dummy_tracer, consumer, producer, kafka_topic
 
     assert consume_span.name == "kafka.consume"
     assert consume_span.get_tag("pathway.hash") is not None
+
+
+def test_tracing_with_serialization_works(dummy_tracer, kafka_topic):
+    def json_serializer(msg, s_obj):
+        return json.dumps(msg).encode("utf-8")
+
+    conf = {
+        "bootstrap.servers": BOOTSTRAP_SERVERS,
+        "key.serializer": json_serializer,
+        "value.serializer": json_serializer,
+    }
+    _producer = confluent_kafka.SerializingProducer(conf)
+
+    def json_deserializer(as_bytes, ctx):
+        try:
+            return json.loads(as_bytes)
+        except json.decoder.JSONDecodeError:
+            return as_bytes
+
+    conf = {
+        "bootstrap.servers": BOOTSTRAP_SERVERS,
+        "group.id": GROUP_ID,
+        "auto.offset.reset": "earliest",
+        "key.deserializer": json_deserializer,
+        "value.deserializer": json_deserializer,
+    }
+
+    _consumer = confluent_kafka.DeserializingConsumer(conf)
+    tp = TopicPartition(kafka_topic, 0)
+    tp.offset = 0  # we want to read the first message
+    _consumer.commit(offsets=[tp])
+    _consumer.subscribe([kafka_topic])
+
+    Pin.override(_producer, tracer=dummy_tracer)
+    Pin.override(_consumer, tracer=dummy_tracer)
+
+    test_string = "serializing_test"
+    PAYLOAD = {"val": test_string}
+
+    _producer.produce(kafka_topic, key={"name": "keykey"}, value=PAYLOAD)
+    _producer.flush()
+
+    message = None
+    while message is None or message.value() != PAYLOAD:
+        message = _consumer.poll(1.0)
+
+    # message comes back with expected test string
+    assert message.value() == PAYLOAD
+
+    traces = dummy_tracer.pop_traces()
+    produce_span = traces[0][0]
+    consume_span = traces[len(traces) - 1][0]
+
+    assert produce_span.get_tag("kafka.message_key") is not None
+
+    # consumer span will not have tag set since we can't serialize the deserialized key from the original type to
+    # a string
+    assert consume_span.get_tag("kafka.message_key") is None