Repository: DataDog/dd-trace-py

Issue:
Version 1.11.0 breaks confluent-kafka `AvroConsumer` and `AvroProducer` classes
Version 1.11.0 breaks confluent-kafka `AvroConsumer` and `AvroProducer` classes, because, compared to `Consumer` and `Producer`, they have additional fields not listed in the `__slots__` of `TracedConsumer` and `TracedProducer`.

The traceback is
```
Traceback (most recent call last):
  File "<frozen runpy>", line 198, in _run_module_as_main
  File "<frozen runpy>", line 88, in _run_code
  File "__main__.py", line 521, in <module>
    AvroConsumer(
  File "/home/federico/.pyenv/versions/test-env/lib/python3.11/site-packages/confluent_kafka/avro/__init__.py", line 161, in __init__
    self._serializer = MessageSerializer(schema_registry, reader_key_schema, reader_value_schema)
    ^^^^^^^^^^^^^^^^
AttributeError: 'cimpl.Consumer' object has no attribute '_serializer'
```

[I have a fix](https://github.com/DataDog/dd-trace-py/compare/1.x...ffissore:dd-trace-py:1.x) but it's probably naive and lacks test coverage. I'll try to make it better in the coming days but don't hold your breath.

_Originally posted by @ffissore in https://github.com/DataDog/dd-trace-py/issues/5027#issuecomment-1497542203_
            


Discussion:
For anyone stumbling on this, you can disable confluent kafka patching by adding this env var: `DD_PATCH_MODULES=kafka:false`

Resolving patch:
diff --git a/ddtrace/contrib/kafka/patch.py b/ddtrace/contrib/kafka/patch.py
--- a/ddtrace/contrib/kafka/patch.py
+++ b/ddtrace/contrib/kafka/patch.py
@@ -14,7 +14,6 @@
 from ddtrace.internal.utils import ArgumentError
 from ddtrace.internal.utils import get_argument_value
 from ddtrace.pin import Pin
-from ddtrace.vendor.wrapt import ObjectProxy
 
 
 _Producer = confluent_kafka.Producer
@@ -27,43 +26,9 @@
 )
 
 
-class TracedProducer(ObjectProxy):
-    def __init__(self, *args, **kwargs):
-        producer = _Producer(*args, **kwargs)
-        super(TracedProducer, self).__init__(producer)
-        Pin().onto(self)
-
-    def produce(self, *args, **kwargs):
-        func = self.__wrapped__.produce
-        topic = get_argument_value(args, kwargs, 0, "topic")
-        try:
-            value = get_argument_value(args, kwargs, 1, "value")
-        except ArgumentError:
-            value = None
-        message_key = kwargs.get("key", "")
-        partition = kwargs.get("partition", -1)
-
-        pin = Pin.get_from(self)
-        if not pin or not pin.enabled():
-            return func(*args, **kwargs)
-
-        with pin.tracer.trace(
-            kafkax.PRODUCE,
-            service=trace_utils.ext_service(pin, config.kafka),
-            span_type=SpanTypes.WORKER,
-        ) as span:
-            span.set_tag_str(MESSAGING_SYSTEM, kafkax.SERVICE)
-            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))
-            span.set_tag(kafkax.PARTITION, partition)
-            span.set_tag_str(kafkax.TOMBSTONE, str(value is None))
-            span.set_tag(SPAN_MEASURED_KEY)
-            rate = config.kafka.get_analytics_sample_rate()
-            if rate is not None:
-                span.set_tag(ANALYTICS_SAMPLE_RATE_KEY, rate)
-            return func(*args, **kwargs)
+class TracedProducer(confluent_kafka.Producer):
+    def produce(self, topic, value=None, *args, **kwargs):
+        super(TracedProducer, self).produce(topic, value, *args, **kwargs)
 
     # in older versions of confluent_kafka, bool(Producer()) evaluates to False,
     # which makes the Pin functionality ignore it.
@@ -73,46 +38,13 @@ def __bool__(self):
     __nonzero__ = __bool__
 
 
-class TracedConsumer(ObjectProxy):
-
-    __slots__ = "_group_id"
-
-    def __init__(self, *args, **kwargs):
-        consumer = _Consumer(*args, **kwargs)
-        super(TracedConsumer, self).__init__(consumer)
-        self._group_id = get_argument_value(args, kwargs, 0, "config")["group.id"]
-        Pin().onto(self)
-
-    def poll(self, *args, **kwargs):
-        func = self.__wrapped__.poll
-        pin = Pin.get_from(self)
-        if not pin or not pin.enabled():
-            return func(*args, **kwargs)
-
-        with pin.tracer.trace(
-            kafkax.CONSUME,
-            service=trace_utils.ext_service(pin, config.kafka),
-            span_type=SpanTypes.WORKER,
-        ) as span:
-            message = func(*args, **kwargs)
-            span.set_tag_str(MESSAGING_SYSTEM, kafkax.SERVICE)
-            span.set_tag_str(COMPONENT, config.kafka.integration_name)
-            span.set_tag_str(SPAN_KIND, SpanKind.CONSUMER)
-            span.set_tag_str(kafkax.RECEIVED_MESSAGE, str(message is not None))
-            span.set_tag_str(kafkax.GROUP_ID, self._group_id)
-            if message is not None:
-                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))
-                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)
-            span.set_tag(SPAN_MEASURED_KEY)
-            rate = config.kafka.get_analytics_sample_rate()
-            if rate is not None:
-                span.set_tag(ANALYTICS_SAMPLE_RATE_KEY, rate)
-            return message
+class TracedConsumer(confluent_kafka.Consumer):
+    def __init__(self, config):
+        super(TracedConsumer, self).__init__(config)
+        self._group_id = config["group.id"]
+
+    def poll(self, timeout=1):
+        return super(TracedConsumer, self).poll(timeout)
 
 
 def patch():
@@ -120,13 +52,86 @@ def patch():
         return
     setattr(confluent_kafka, "_datadog_patch", True)
 
-    setattr(confluent_kafka, "Producer", TracedProducer)
-    setattr(confluent_kafka, "Consumer", TracedConsumer)
+    confluent_kafka.Producer = TracedProducer
+    confluent_kafka.Consumer = TracedConsumer
+
+    trace_utils.wrap(TracedProducer, "produce", traced_produce)
+    trace_utils.wrap(TracedConsumer, "poll", traced_poll)
+    Pin().onto(confluent_kafka.Producer)
+    Pin().onto(confluent_kafka.Consumer)
 
 
 def unpatch():
     if getattr(confluent_kafka, "_datadog_patch", False):
         setattr(confluent_kafka, "_datadog_patch", False)
 
-    setattr(confluent_kafka, "Producer", _Producer)
-    setattr(confluent_kafka, "Consumer", _Consumer)
+    if trace_utils.iswrapped(TracedProducer.produce):
+        trace_utils.unwrap(TracedProducer, "produce")
+    if trace_utils.iswrapped(TracedConsumer.poll):
+        trace_utils.unwrap(TracedConsumer, "poll")
+
+    confluent_kafka.Producer = _Producer
+    confluent_kafka.Consumer = _Consumer
+
+
+def traced_produce(func, instance, args, kwargs):
+    pin = Pin.get_from(instance)
+    if not pin or not pin.enabled():
+        return func(*args, **kwargs)
+
+    topic = get_argument_value(args, kwargs, 0, "topic")
+    try:
+        value = get_argument_value(args, kwargs, 1, "value")
+    except ArgumentError:
+        value = None
+    message_key = kwargs.get("key", "")
+    partition = kwargs.get("partition", -1)
+
+    with pin.tracer.trace(
+        kafkax.PRODUCE,
+        service=trace_utils.ext_service(pin, config.kafka),
+        span_type=SpanTypes.WORKER,
+    ) as span:
+        span.set_tag_str(MESSAGING_SYSTEM, kafkax.SERVICE)
+        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))
+        span.set_tag(kafkax.PARTITION, partition)
+        span.set_tag_str(kafkax.TOMBSTONE, str(value is None))
+        span.set_tag(SPAN_MEASURED_KEY)
+        rate = config.kafka.get_analytics_sample_rate()
+        if rate is not None:
+            span.set_tag(ANALYTICS_SAMPLE_RATE_KEY, rate)
+        return func(*args, **kwargs)
+
+
+def traced_poll(func, instance, args, kwargs):
+    pin = Pin.get_from(instance)
+    if not pin or not pin.enabled():
+        return func(*args, **kwargs)
+
+    with pin.tracer.trace(
+        kafkax.CONSUME,
+        service=trace_utils.ext_service(pin, config.kafka),
+        span_type=SpanTypes.WORKER,
+    ) as span:
+        message = func(*args, **kwargs)
+        span.set_tag_str(MESSAGING_SYSTEM, kafkax.SERVICE)
+        span.set_tag_str(COMPONENT, config.kafka.integration_name)
+        span.set_tag_str(SPAN_KIND, SpanKind.CONSUMER)
+        span.set_tag_str(kafkax.RECEIVED_MESSAGE, str(message is not None))
+        span.set_tag_str(kafkax.GROUP_ID, instance._group_id)
+        if message is not None:
+            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))
+            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)
+        span.set_tag(SPAN_MEASURED_KEY)
+        rate = config.kafka.get_analytics_sample_rate()
+        if rate is not None:
+            span.set_tag(ANALYTICS_SAMPLE_RATE_KEY, rate)
+        return message


Test patch:
diff --git a/tests/contrib/kafka/test_kafka_patch.py b/tests/contrib/kafka/test_kafka_patch.py
--- a/tests/contrib/kafka/test_kafka_patch.py
+++ b/tests/contrib/kafka/test_kafka_patch.py
@@ -9,17 +9,14 @@ class TestKafkaPatch(PatchTestCase.Base):
     __patch_func__ = patch
     __unpatch_func__ = unpatch
 
-    # DEV: normally, we directly patch methods, but since confluent-kafka's methods are implemented in C and
-    # directly imported, we have to patch the Producer/Consumer classes via proxy Traced Producer/Consumer classes.
-    # Because of this, we need to create instances of each proxy class to make wrapping status assertions.
     def assert_module_patched(self, confluent_kafka):
-        self.assert_wrapped(confluent_kafka.Producer({}))
-        self.assert_wrapped(confluent_kafka.Consumer({"group.id": "group_id"}))
+        self.assert_wrapped(confluent_kafka.Producer({}).produce)
+        self.assert_wrapped(confluent_kafka.Consumer({"group.id": "group_id"}).poll)
 
     def assert_not_module_patched(self, confluent_kafka):
-        self.assert_not_wrapped(confluent_kafka.Producer({}))
-        self.assert_not_wrapped(confluent_kafka.Consumer({"group.id": "group_id"}))
+        self.assert_not_wrapped(confluent_kafka.Producer({}).produce)
+        self.assert_not_wrapped(confluent_kafka.Consumer({"group.id": "group_id"}).poll)
 
     def assert_not_module_double_patched(self, confluent_kafka):
-        self.assert_not_double_wrapped(confluent_kafka.Producer({}))
-        self.assert_not_double_wrapped(confluent_kafka.Consumer({"group.id": "group_id"}))
+        self.assert_not_double_wrapped(confluent_kafka.Producer({}).produce)
+        self.assert_not_double_wrapped(confluent_kafka.Consumer({"group.id": "group_id"}).poll)