Kafka consumer : add consumer group state as tag (#17686)

* kafka consumer group state

* kafka consumer group state 1

* kafka consumer group state fix test

* kafka consumer group state linter and changelog

* kafka consumer group state update

* kafka consumer group state remove expception

* kafka consumer group state remove expception 2

* kafka consumer group state remove expception 3

* kafka consumer group state remove expception 3\4

* kafka consumer group state fix tests

Author: HadhemiDD

kafka_consumer/changelog.d/17686.added

@@ -0,0 +1 @@
+Kafka consumer : add consumer group state as tag 

kafka_consumer/datadog_checks/kafka_consumer/client.py

@@ -254,9 +254,39 @@ def _get_consumer_groups(self):
         else:
             return self.config._consumer_groups
 
+    def get_consumer_group_state(self, consumer_group):
+        consumer_group_state = ""
+        # Get the consumer group state if present
+        consumer_groups_future = self._describe_consumer_groups(consumer_group)
+        consumer_groups_result = consumer_groups_future[consumer_group].result()
+        self.log.debug(
+            "Consumer group: %s in state %s",
+            consumer_groups_result.group_id,
+            consumer_groups_result.state,
+        )
+        consumer_group_result_state = str(consumer_groups_result.state)
+        consumer_group_state = consumer_group_result_state.split('.')[1]
+
+        return consumer_group_state
+
     def _list_consumer_group_offsets(self, cg_tp):
+        """
+        :returns: A dict of futures for each group, keyed by the group id.
+                  The future result() method returns :class:`ConsumerGroupTopicPartitions`.
+
+        :rtype: dict[str, future]
+        """
         return self.kafka_client.list_consumer_group_offsets([cg_tp])
 
+    def _describe_consumer_groups(self, consumer_group):
+        """
+        :returns: A dict of futures for each group, keyed by the group_id.
+                  The future result() method returns :class:`ConsumerGroupDescription`.
+
+        :rtype: dict[str, future]
+        """
+        return self.kafka_client.describe_consumer_groups([consumer_group])
+
     def close_admin_client(self):
         self._kafka_client = None
 

kafka_consumer/datadog_checks/kafka_consumer/kafka_consumer.py

@@ -138,6 +138,7 @@ def report_consumer_offsets_and_lag(
         reported_contexts = 0
         self.log.debug("Reporting consumer offsets and lag metrics")
         for (consumer_group, topic, partition), consumer_offset in consumer_offsets.items():
+            consumer_group_state = self.client.get_consumer_group_state(consumer_group)
             if reported_contexts >= contexts_limit:
                 self.log.debug(
                     "Reported contexts number %s greater than or equal to contexts limit of %s, returning",
@@ -151,6 +152,7 @@ def report_consumer_offsets_and_lag(
                 'partition:%s' % partition,
                 'consumer_group:%s' % consumer_group,
                 'kafka_cluster_id:%s' % cluster_id,
+                'consumer_group_state:%s' % consumer_group_state,
             ]
             consumer_group_tags.extend(self.config._custom_tags)
 

kafka_consumer/tests/common.py

@@ -102,11 +102,11 @@ def assert_check_kafka(aggregator, consumer_groups):
                     aggregator.assert_metric(mname, tags=tags, count=1)
 
                 for mname in CONSUMER_METRICS:
-                    aggregator.assert_metric(
-                        mname,
-                        tags=tags + [f"consumer_group:{name}"],
-                        count=1,
-                    )
+                    aggregator.assert_metric(mname)
+                    t = tags + [f"consumer_group:{name}"]
+                    if aggregator.assert_metric_has_tags(mname, t):
+                        # Check for the tag consumer_group_state
+                        aggregator.assert_metric_has_tag_prefix(mname, tag_prefix='consumer_group_state')
 
     aggregator.assert_all_metrics_covered()
     aggregator.assert_metrics_using_metadata(get_metadata_metrics())

kafka_consumer/tests/test_unit.py

@@ -162,6 +162,7 @@ def test_when_consumer_lag_less_than_zero_then_emit_event(
     mock_client.get_consumer_offsets.return_value = consumer_offset
     mock_client.get_highwater_offsets.return_value = (highwater_offset, "cluster_id")
     mock_client.get_partitions_for_topic.return_value = ['partition1']
+    mock_client.get_consumer_group_state.return_value = "STABLE"
     mock_generic_client.return_value = mock_client
 
     # When
@@ -183,6 +184,7 @@ def test_when_consumer_lag_less_than_zero_then_emit_event(
             'partition:partition1',
             'topic:topic1',
             'kafka_cluster_id:cluster_id',
+            'consumer_group_state:STABLE',
         ],
     )
     aggregator.assert_metric(
@@ -194,6 +196,7 @@ def test_when_consumer_lag_less_than_zero_then_emit_event(
             'partition:partition1',
             'topic:topic1',
             'kafka_cluster_id:cluster_id',
+            'consumer_group_state:STABLE',
         ],
     )
     aggregator.assert_event(
@@ -208,6 +211,7 @@ def test_when_consumer_lag_less_than_zero_then_emit_event(
             'partition:partition1',
             'topic:topic1',
             'kafka_cluster_id:cluster_id',
+            'consumer_group_state:STABLE',
         ],
     )