Implement regex for consumer_groups (#14382)

yzhan289 · web-flow · commit 2cbd09e53814 · 2023-04-14T13:50:23.000-04:00
* Implement regex for consumer_groups

* Add validated files

* Match original PR

* Fix tab
diff --git a/kafka_consumer/assets/configuration/spec.yaml b/kafka_consumer/assets/configuration/spec.yaml
@@ -62,8 +62,28 @@ files:
       - name: consumer_groups
         description: |
           Each level is optional. Any empty values are fetched from the Kafka cluster.
-          You can have empty partitions (example: <CONSUMER_NAME_2>), topics (example: <CONSUMER_NAME_3>),
-          and even consumer_groups. If you omit consumer_groups, you must set `monitor_unlisted_consumer_groups` to true.
+          You can have empty partitions (example: <CONSUMER_NAME_2>), topics
+          (example: <CONSUMER_NAME_3>),
+          and even consumer_groups. If you do not specify `consumer_groups`,
+          you must set `monitor_unlisted_consumer_groups` to true. 
+          
+          If both `consumer_groups` and `monitor_unlisted_consumer_groups` are used,
+          then `consumer_groups` is ignored.
+        value:
+          type: object
+          example:
+            <CONSUMER_NAME_1>:
+              <TOPIC_NAME_1>: [0, 1, 4, 12]
+            <CONSUMER_NAME_2>:
+              <TOPIC_NAME_2>: []
+            <CONSUMER_NAME_3>: {}
+      - name: consumer_groups_regex
+        description: |
+          Set consumer groups and topics as regex strings.
+          If both `consumer_groups_regex` and `consumer_groups` are filled out,
+          both configuration options will be used. 
+          If `monitor_unlisted_consumer_groups` is enabled, then
+          `consumer_groups_regex` is ignored.
         value:
           type: object
           example:
@@ -74,8 +94,8 @@ files:
             <CONSUMER_NAME_3>: {}
       - name: monitor_unlisted_consumer_groups
         description: |
-          Setting monitor_unlisted_consumer_groups to `true` tells the check to discover all consumer groups
-          and fetch all their known offsets. If this is not set to true, you must specify consumer_groups.
+          Setting `monitor_unlisted_consumer_groups` to `true` tells the check to discover all consumer groups
+          and fetch all their known offsets. If this is not set to true, you must specify `consumer_groups`.
 
           WARNING: This feature requires that your Kafka brokers be version >= 0.10.2. It is impossible to
           support this feature on older brokers because they do not provide a way to determine the mapping
diff --git a/kafka_consumer/datadog_checks/kafka_consumer/client.py b/kafka_consumer/datadog_checks/kafka_consumer/client.py
@@ -120,30 +120,7 @@ def get_consumer_offsets(self):
         # {(consumer_group, topic, partition): offset}
         consumer_offsets = {}
 
-        if self.config._monitor_unlisted_consumer_groups:
-            # Get all consumer groups
-            consumer_groups = []
-            consumer_groups_future = self.kafka_client.list_consumer_groups()
-            self.log.debug('MONITOR UNLISTED CG FUTURES: %s', consumer_groups_future)
-            try:
-                list_consumer_groups_result = consumer_groups_future.result()
-                self.log.debug('MONITOR UNLISTED FUTURES RESULT: %s', list_consumer_groups_result)
-
-                consumer_groups.extend(
-                    valid_consumer_group.group_id for valid_consumer_group in list_consumer_groups_result.valid
-                )
-            except Exception as e:
-                self.log.error("Failed to collect consumer offsets %s", e)
-
-        elif self.config._consumer_groups:
-            self._validate_consumer_groups()
-            consumer_groups = self.config._consumer_groups
-
-        else:
-            raise ConfigurationError(
-                "Cannot fetch consumer offsets because no consumer_groups are specified and "
-                "monitor_unlisted_consumer_groups is %s." % self.config._monitor_unlisted_consumer_groups
-            )
+        consumer_groups = self._get_consumer_groups()
 
         for future in self._get_consumer_offset_futures(consumer_groups):
             try:
@@ -178,6 +155,31 @@ def get_consumer_offsets(self):
 
         return consumer_offsets
 
+    def _get_consumer_groups(self):
+        if self.config._monitor_unlisted_consumer_groups or self.config._consumer_groups_regex:
+            # Get all consumer groups
+            consumer_groups = []
+            consumer_groups_future = self.kafka_client.list_consumer_groups()
+            self.log.debug('MONITOR UNLISTED CG FUTURES: %s', consumer_groups_future)
+            try:
+                list_consumer_groups_result = consumer_groups_future.result()
+                self.log.debug('MONITOR UNLISTED FUTURES RESULT: %s', list_consumer_groups_result)
+
+                consumer_groups.extend(
+                    valid_consumer_group.group_id for valid_consumer_group in list_consumer_groups_result.valid
+                )
+            except Exception as e:
+                self.log.error("Failed to collect consumer groups: %s", e)
+            return consumer_groups
+        elif self.config._consumer_groups:
+            self._validate_consumer_groups()
+            return self.config._consumer_groups
+        else:
+            raise ConfigurationError(
+                "Cannot fetch consumer offsets because no consumer_groups are specified and "
+                "monitor_unlisted_consumer_groups is %s." % self.config._monitor_unlisted_consumer_groups
+            )
+
     def _get_consumer_offset_futures(self, consumer_groups):
         topics = self.kafka_client.list_topics(timeout=self.config._request_timeout)
         # {(consumer_group, topic, partition): offset}
@@ -221,24 +223,77 @@ def _get_topic_partitions(self, topics, consumer_group):
 
             partitions = list(topics.topics[topic].partitions.keys())
 
-            for partition in partitions:
-                # Get all topic-partition combinations allowed based on config
-                # if topics is None => collect all topics and partitions for the consumer group
-                # if partitions is None => collect all partitions from the consumer group's topic
-                if not self.config._monitor_unlisted_consumer_groups and self.config._consumer_groups.get(
-                    consumer_group
+            if self.config._monitor_unlisted_consumer_groups:
+                for partition in partitions:
+                    topic_partition = TopicPartition(topic, partition)
+                    self.log.debug("TOPIC PARTITION: %s", topic_partition)
+                    yield topic_partition
+
+            elif self.config._consumer_groups_regex:
+                for filtered_topic_partition in self._get_regex_filtered_topic_partitions(
+                    consumer_group, topic, partitions
                 ):
-                    if (
-                        self.config._consumer_groups[consumer_group]
-                        and topic not in self.config._consumer_groups[consumer_group]
-                    ):
+                    topic_partition = TopicPartition(filtered_topic_partition[0], filtered_topic_partition[1])
+                    self.log.debug("TOPIC PARTITION: %s", topic_partition)
+                    yield topic_partition
+
+            if self.config._consumer_groups:
+                for partition in partitions:
+                    # Get all topic-partition combinations allowed based on config
+                    # if topics is None => collect all topics and partitions for the consumer group
+                    # if partitions is None => collect all partitions from the consumer group's topic
+                    if self.config._consumer_groups.get(consumer_group):
+                        if (
+                            self.config._consumer_groups[consumer_group]
+                            and topic not in self.config._consumer_groups[consumer_group]
+                        ):
+                            self.log.debug(
+                                "Partition %s skipped because the topic %s is not in the consumer_group.",
+                                partition,
+                                topic,
+                            )
+                            continue
+                        if (
+                            self.config._consumer_groups[consumer_group].get(topic)
+                            and partition not in self.config._consumer_groups[consumer_group][topic]
+                        ):
+                            self.log.debug(
+                                "Partition %s skipped because it is not defined in the consumer group for the topic %s",
+                                partition,
+                                topic,
+                            )
+                            continue
+
+                    topic_partition = TopicPartition(topic, partition)
+                    self.log.debug("TOPIC PARTITION: %s", topic_partition)
+                    yield topic_partition
+
+    def _get_regex_filtered_topic_partitions(self, consumer_group, topic, partitions):
+        for partition in partitions:
+            # Do a regex filtering here for consumer groups
+            for consumer_group_compiled_regex in self.config._consumer_groups_compiled_regex:
+                if not consumer_group_compiled_regex.match(consumer_group):
+                    return
+
+                consumer_group_topics_regex = self.config._consumer_groups_compiled_regex.get(
+                    consumer_group_compiled_regex
+                )
+
+                # If topics is empty, return all combinations of topic and partition
+                if not consumer_group_topics_regex:
+                    yield (topic, partition)
+
+                # Do a regex filtering here for topics
+                for topic_regex in consumer_group_topics_regex:
+                    if not topic_regex.match(topic):
                         self.log.debug(
                             "Partition %s skipped because the topic %s is not in the consumer_group.", partition, topic
                         )
                         continue
+
                     if (
-                        self.config._consumer_groups[consumer_group].get(topic)
-                        and partition not in self.config._consumer_groups[consumer_group][topic]
+                        consumer_group_topics_regex.get(topic_regex)
+                        and partition not in consumer_group_topics_regex[topic_regex]
                     ):
                         self.log.debug(
                             "Partition %s skipped because it is not defined in the consumer group for the topic %s",
@@ -247,6 +302,4 @@ def _get_topic_partitions(self, topics, consumer_group):
                         )
                         continue
 
-                self.log.debug("TOPIC PARTITION: %s", TopicPartition(topic, partition))
-
-                yield TopicPartition(topic, partition)
+                    yield (topic, partition)
diff --git a/kafka_consumer/datadog_checks/kafka_consumer/config.py b/kafka_consumer/datadog_checks/kafka_consumer/config.py
@@ -2,25 +2,29 @@
 # All rights reserved
 # Licensed under a 3-clause BSD style license (see LICENSE)
 import os
+import re
 
 from datadog_checks.base import ConfigurationError, is_affirmative
 from datadog_checks.kafka_consumer.constants import CONTEXT_UPPER_BOUND, DEFAULT_KAFKA_TIMEOUT
 
 
 class KafkaConfig:
-    def __init__(self, init_config, instance) -> None:
+    def __init__(self, init_config, instance, log) -> None:
         self.instance = instance
         self.init_config = init_config
         self._context_limit = int(init_config.get('max_partition_contexts', CONTEXT_UPPER_BOUND))
+        self.log = log
         self._custom_tags = instance.get('tags', [])
         self._monitor_unlisted_consumer_groups = is_affirmative(instance.get('monitor_unlisted_consumer_groups', False))
         self._monitor_all_broker_highwatermarks = is_affirmative(
             instance.get('monitor_all_broker_highwatermarks', False)
         )
         self._consumer_groups = instance.get('consumer_groups', {})
+        self._consumer_groups_regex = instance.get('consumer_groups_regex', {})
 
-        self._kafka_connect_str = instance.get('kafka_connect_str')
+        self._consumer_groups_compiled_regex = self._compile_regex(self._consumer_groups_regex)
 
+        self._kafka_connect_str = instance.get('kafka_connect_str')
         self._kafka_version = instance.get('kafka_client_api_version')
         if isinstance(self._kafka_version, str):
             self._kafka_version = tuple(map(int, self._kafka_version.split(".")))
@@ -72,3 +76,31 @@ def validate_config(self):
 
             elif self._sasl_oauth_token_provider.get("client_secret") is None:
                 raise ConfigurationError("The `client_secret` setting of `auth_token` reader is required")
+
+        # If `monitor_unlisted_consumer_groups` is set to true and
+        # using `consumer_groups`, we prioritize `monitor_unlisted_consumer_groups`
+        if self._monitor_unlisted_consumer_groups and (self._consumer_groups or self._consumer_groups_regex):
+            self.log.warning(
+                "Using both monitor_unlisted_consumer_groups and consumer_groups or consumer_groups_regex, "
+                "so all consumer groups will be collected."
+            )
+
+        if self._consumer_groups and self._consumer_groups_regex:
+            self.log.warning("Using consumer_groups and consumer_groups_regex, will combine the two config options.")
+
+    def _compile_regex(self, consumer_groups_regex):
+        patterns = {}
+
+        for consumer_group_regex in consumer_groups_regex:
+            consumer_group_pattern = re.compile(consumer_group_regex)
+            patterns[consumer_group_pattern] = {}
+
+            topics_regex = consumer_groups_regex.get(consumer_group_regex)
+
+            for topic_regex in topics_regex:
+                topic_pattern = re.compile(topic_regex)
+
+                partitions = self._consumer_groups_regex[consumer_group_regex][topic_regex]
+                patterns[consumer_group_pattern].update({topic_pattern: partitions})
+
+        return patterns
diff --git a/kafka_consumer/datadog_checks/kafka_consumer/config_models/defaults.py b/kafka_consumer/datadog_checks/kafka_consumer/config_models/defaults.py
@@ -22,6 +22,10 @@ def instance_consumer_groups(field, value):
     return get_default_field_value(field, value)
 
 
+def instance_consumer_groups_regex(field, value):
+    return get_default_field_value(field, value)
+
+
 def instance_disable_generic_tags(field, value):
     return False
 
diff --git a/kafka_consumer/datadog_checks/kafka_consumer/config_models/instance.py b/kafka_consumer/datadog_checks/kafka_consumer/config_models/instance.py
@@ -41,6 +41,7 @@ class Config:
         allow_mutation = False
 
     consumer_groups: Optional[Mapping[str, Any]]
+    consumer_groups_regex: Optional[Mapping[str, Any]]
     disable_generic_tags: Optional[bool]
     empty_default_hostname: Optional[bool]
     kafka_client_api_version: Optional[str]
diff --git a/kafka_consumer/datadog_checks/kafka_consumer/data/conf.yaml.example b/kafka_consumer/datadog_checks/kafka_consumer/data/conf.yaml.example
@@ -56,8 +56,13 @@ instances:
 
     ## @param consumer_groups - mapping - optional
     ## Each level is optional. Any empty values are fetched from the Kafka cluster.
-    ## You can have empty partitions (example: <CONSUMER_NAME_2>), topics (example: <CONSUMER_NAME_3>),
-    ## and even consumer_groups. If you omit consumer_groups, you must set `monitor_unlisted_consumer_groups` to true.
+    ## You can have empty partitions (example: <CONSUMER_NAME_2>), topics
+    ## (example: <CONSUMER_NAME_3>),
+    ## and even consumer_groups. If you do not specify `consumer_groups`,
+    ## you must set `monitor_unlisted_consumer_groups` to true. 
+    ##
+    ## If both `consumer_groups` and `monitor_unlisted_consumer_groups` are used,
+    ## then `consumer_groups` is ignored.
     #
     # consumer_groups:
     #   <CONSUMER_NAME_1>:
@@ -70,9 +75,27 @@ instances:
     #     <TOPIC_NAME_2>: []
     #   <CONSUMER_NAME_3>: {}
 
+    ## @param consumer_groups_regex - mapping - optional
+    ## Set consumer groups and topics as regex strings.
+    ## If both `consumer_groups_regex` and `consumer_groups` are filled out,
+    ## both configuration options will be used. 
+    ## If `monitor_unlisted_consumer_groups` is enabled, then
+    ## `consumer_groups_regex` is ignored.
+    #
+    # consumer_groups_regex:
+    #   <CONSUMER_NAME_1>:
+    #     <TOPIC_NAME_1>:
+    #     - 0
+    #     - 1
+    #     - 4
+    #     - 12
+    #   <CONSUMER_NAME_2>:
+    #     <TOPIC_NAME_2>: []
+    #   <CONSUMER_NAME_3>: {}
+
     ## @param monitor_unlisted_consumer_groups - boolean - optional - default: false
-    ## Setting monitor_unlisted_consumer_groups to `true` tells the check to discover all consumer groups
-    ## and fetch all their known offsets. If this is not set to true, you must specify consumer_groups.
+    ## Setting `monitor_unlisted_consumer_groups` to `true` tells the check to discover all consumer groups
+    ## and fetch all their known offsets. If this is not set to true, you must specify `consumer_groups`.
     ##
     ## WARNING: This feature requires that your Kafka brokers be version >= 0.10.2. It is impossible to
     ## support this feature on older brokers because they do not provide a way to determine the mapping
diff --git a/kafka_consumer/datadog_checks/kafka_consumer/kafka_consumer.py b/kafka_consumer/datadog_checks/kafka_consumer/kafka_consumer.py
@@ -23,7 +23,7 @@ class KafkaCheck(AgentCheck):
 
     def __init__(self, name, init_config, instances):
         super(KafkaCheck, self).__init__(name, init_config, instances)
-        self.config = KafkaConfig(self.init_config, self.instance)
+        self.config = KafkaConfig(self.init_config, self.instance, self.log)
         self._context_limit = self.config._context_limit
         self.client = KafkaClient(self.config, self.get_tls_context(), self.log)
         self.check_initializations.append(self.config.validate_config)
diff --git a/kafka_consumer/tests/test_integration.py b/kafka_consumer/tests/test_integration.py
