fix(kafka): avoid overriding methods in Producers and Consumers [backport 2.0] (#7007)

Author: github-actions[bot]

ddtrace/contrib/kafka/patch.py

@@ -55,9 +55,6 @@ def __init__(self, config, *args, **kwargs):
             else config.get("metadata.broker.list")
         )
 
-    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.
     def __bool__(self):
@@ -72,12 +69,6 @@ def __init__(self, config, *args, **kwargs):
         self._group_id = config.get("group.id", "")
         self._auto_commit = asbool(config.get("enable.auto.commit", True))
 
-    def poll(self, timeout=None):
-        return super(TracedConsumer, self).poll(timeout)
-
-    def commit(self, message=None, *args, **kwargs):
-        return super(TracedConsumer, self).commit(message, args, kwargs)
-
 
 def patch():
     if getattr(confluent_kafka, "_datadog_patch", False):
@@ -230,12 +221,16 @@ def traced_commit(func, instance, args, kwargs):
         return func(*args, **kwargs)
 
     if config._data_streams_enabled:
-        message = get_argument_value(args, kwargs, 0, "message")
-        offsets = kwargs.get("offsets", [])
+        message = get_argument_value(args, kwargs, 0, "message", True)
+        # message and offset are mutually exclusive. Only one parameter can be passed.
         if message is not None:
             offsets = [TopicPartition(message.topic(), message.partition(), offset=message.offset())]
-        for offset in offsets:
-            pin.tracer.data_streams_processor.track_kafka_commit(
-                instance._group_id, offset.topic, offset.partition, offset.offset or -1, time.time()
-            )
+        else:
+            offsets = get_argument_value(args, kwargs, 1, "offsets", True)
+
+        if offsets:
+            for offset in offsets:
+                pin.tracer.data_streams_processor.track_kafka_commit(
+                    instance._group_id, offset.topic, offset.partition, offset.offset or -1, time.time()
+                )
     return func(*args, **kwargs)

ddtrace/internal/utils/__init__.py

@@ -19,6 +19,7 @@ def get_argument_value(
     kwargs,  # type: Dict[str, Any]
     pos,  # type: int
     kw,  # type: str
+    optional=False,  # type: bool
 ):
     # type: (...) -> Optional[Any]
     """
@@ -41,6 +42,8 @@ def get_argument_value(
         try:
             return args[pos]
         except IndexError:
+            if optional:
+                return None
             raise ArgumentError("%s (at position %d)" % (kw, pos))
 
 

releasenotes/notes/fix-kafka-commit-offsets-d86039c29e04259e.yaml

@@ -0,0 +1,4 @@
+---
+fixes:
+  - |
+    kafka: Fixes ``ValueError`` raised when ``Consumer.commit(offsets=...)`` is called.

tests/contrib/kafka/test_kafka.py

@@ -4,6 +4,7 @@
 
 import confluent_kafka
 from confluent_kafka import KafkaException
+from confluent_kafka import TopicPartition
 from confluent_kafka import admin as kafka_admin
 import mock
 import pytest
@@ -190,6 +191,36 @@ def test_message(producer, consumer, tombstone, kafka_topic):
         message = consumer.poll(1.0)
 
 
+@pytest.mark.snapshot(ignores=["metrics.kafka.message_offset"])
+def test_commit(producer, consumer, kafka_topic):
+    producer.produce(kafka_topic, PAYLOAD, key=KEY)
+    producer.flush()
+    message = None
+    while message is None:
+        message = consumer.poll(1.0)
+    consumer.commit(message)
+
+
+@pytest.mark.snapshot(ignores=["metrics.kafka.message_offset"])
+def test_commit_with_offset(producer, consumer, kafka_topic):
+    producer.produce(kafka_topic, PAYLOAD, key=KEY)
+    producer.flush()
+    message = None
+    while message is None:
+        message = consumer.poll(1.0)
+    consumer.commit(offsets=[TopicPartition(kafka_topic)])
+
+
+@pytest.mark.snapshot(ignores=["metrics.kafka.message_offset"])
+def test_commit_with_only_async_arg(producer, consumer, kafka_topic):
+    producer.produce(kafka_topic, PAYLOAD, key=KEY)
+    producer.flush()
+    message = None
+    while message is None:
+        message = consumer.poll(1.0)
+    consumer.commit(asynchronous=False)
+
+
 @pytest.mark.snapshot(
     token="tests.contrib.kafka.test_kafka.test_service_override", ignores=["metrics.kafka.message_offset"]
 )

tests/snapshots/tests.contrib.kafka.test_kafka.test_commit.json

@@ -0,0 +1,68 @@
+[[
+  {
+    "name": "kafka.consume",
+    "service": "kafka",
+    "resource": "kafka.consume",
+    "trace_id": 0,
+    "span_id": 1,
+    "parent_id": 0,
+    "type": "worker",
+    "meta": {
+      "_dd.base_service": "",
+      "_dd.p.dm": "-0",
+      "component": "kafka",
+      "kafka.group_id": "test_group",
+      "kafka.message_key": "test_key",
+      "kafka.received_message": "True",
+      "kafka.tombstone": "False",
+      "kafka.topic": "test_commit",
+      "language": "python",
+      "messaging.system": "kafka",
+      "runtime-id": "4cf20234df69485781c43482565d9d82",
+      "span.kind": "consumer"
+    },
+    "metrics": {
+      "_dd.measured": 1,
+      "_dd.top_level": 1,
+      "_dd.tracer_kr": 1.0,
+      "_sampling_priority_v1": 1,
+      "kafka.message_offset": -1,
+      "kafka.partition": 0,
+      "process_id": 222
+    },
+    "duration": 187877000,
+    "start": 1695310282564811000
+  }],
+[
+  {
+    "name": "kafka.produce",
+    "service": "kafka",
+    "resource": "kafka.produce",
+    "trace_id": 1,
+    "span_id": 1,
+    "parent_id": 0,
+    "type": "worker",
+    "meta": {
+      "_dd.base_service": "",
+      "_dd.p.dm": "-0",
+      "component": "kafka",
+      "kafka.message_key": "test_key",
+      "kafka.tombstone": "False",
+      "kafka.topic": "test_commit",
+      "language": "python",
+      "messaging.kafka.bootstrap.servers": "localhost:29092",
+      "messaging.system": "kafka",
+      "runtime-id": "4cf20234df69485781c43482565d9d82",
+      "span.kind": "producer"
+    },
+    "metrics": {
+      "_dd.measured": 1,
+      "_dd.top_level": 1,
+      "_dd.tracer_kr": 1.0,
+      "_sampling_priority_v1": 1,
+      "kafka.partition": -1,
+      "process_id": 222
+    },
+    "duration": 135000,
+    "start": 1695310277502282000
+  }]]

tests/snapshots/tests.contrib.kafka.test_kafka.test_commit_with_offset.json

@@ -0,0 +1,68 @@
+[[
+  {
+    "name": "kafka.consume",
+    "service": "kafka",
+    "resource": "kafka.consume",
+    "trace_id": 0,
+    "span_id": 1,
+    "parent_id": 0,
+    "type": "worker",
+    "meta": {
+      "_dd.base_service": "",
+      "_dd.p.dm": "-0",
+      "component": "kafka",
+      "kafka.group_id": "test_group",
+      "kafka.message_key": "test_key",
+      "kafka.received_message": "True",
+      "kafka.tombstone": "False",
+      "kafka.topic": "test_commit_with_offset",
+      "language": "python",
+      "messaging.system": "kafka",
+      "runtime-id": "9a7c9425ae8745c1bb8d35d2514f8958",
+      "span.kind": "consumer"
+    },
+    "metrics": {
+      "_dd.measured": 1,
+      "_dd.top_level": 1,
+      "_dd.tracer_kr": 1.0,
+      "_sampling_priority_v1": 1,
+      "kafka.message_offset": -1,
+      "kafka.partition": 0,
+      "process_id": 6340
+    },
+    "duration": 307084000,
+    "start": 1695311385900323000
+  }],
+[
+  {
+    "name": "kafka.produce",
+    "service": "kafka",
+    "resource": "kafka.produce",
+    "trace_id": 1,
+    "span_id": 1,
+    "parent_id": 0,
+    "type": "worker",
+    "meta": {
+      "_dd.base_service": "",
+      "_dd.p.dm": "-0",
+      "component": "kafka",
+      "kafka.message_key": "test_key",
+      "kafka.tombstone": "False",
+      "kafka.topic": "test_commit_with_offset",
+      "language": "python",
+      "messaging.kafka.bootstrap.servers": "localhost:29092",
+      "messaging.system": "kafka",
+      "runtime-id": "9a7c9425ae8745c1bb8d35d2514f8958",
+      "span.kind": "producer"
+    },
+    "metrics": {
+      "_dd.measured": 1,
+      "_dd.top_level": 1,
+      "_dd.tracer_kr": 1.0,
+      "_sampling_priority_v1": 1,
+      "kafka.partition": -1,
+      "process_id": 6340
+    },
+    "duration": 99000,
+    "start": 1695311382892915000
+  }]]

tests/snapshots/tests.contrib.kafka.test_kafka.test_commit_with_only_async_arg.json

@@ -0,0 +1,68 @@
+[[
+  {
+    "name": "kafka.consume",
+    "service": "kafka",
+    "resource": "kafka.consume",
+    "trace_id": 0,
+    "span_id": 1,
+    "parent_id": 0,
+    "type": "worker",
+    "meta": {
+      "_dd.base_service": "",
+      "_dd.p.dm": "-0",
+      "component": "kafka",
+      "kafka.group_id": "test_group",
+      "kafka.message_key": "test_key",
+      "kafka.received_message": "True",
+      "kafka.tombstone": "False",
+      "kafka.topic": "test_commit_with_only_async_arg",
+      "language": "python",
+      "messaging.system": "kafka",
+      "runtime-id": "15fa67b4ad1e498aa489305814d88a75",
+      "span.kind": "consumer"
+    },
+    "metrics": {
+      "_dd.measured": 1,
+      "_dd.top_level": 1,
+      "_dd.tracer_kr": 1.0,
+      "_sampling_priority_v1": 1,
+      "kafka.message_offset": -1,
+      "kafka.partition": 0,
+      "process_id": 35657
+    },
+    "duration": 343409000,
+    "start": 1695323048560969000
+  }],
+[
+  {
+    "name": "kafka.produce",
+    "service": "kafka",
+    "resource": "kafka.produce",
+    "trace_id": 1,
+    "span_id": 1,
+    "parent_id": 0,
+    "type": "worker",
+    "meta": {
+      "_dd.base_service": "",
+      "_dd.p.dm": "-0",
+      "component": "kafka",
+      "kafka.message_key": "test_key",
+      "kafka.tombstone": "False",
+      "kafka.topic": "test_commit_with_only_async_arg",
+      "language": "python",
+      "messaging.kafka.bootstrap.servers": "localhost:29092",
+      "messaging.system": "kafka",
+      "runtime-id": "15fa67b4ad1e498aa489305814d88a75",
+      "span.kind": "producer"
+    },
+    "metrics": {
+      "_dd.measured": 1,
+      "_dd.top_level": 1,
+      "_dd.tracer_kr": 1.0,
+      "_sampling_priority_v1": 1,
+      "kafka.partition": -1,
+      "process_id": 35657
+    },
+    "duration": 280000,
+    "start": 1695323045564363000
+  }]]