Data Streams live messages (DataDog#20512)

* Data Streams: log messages from Kafka

* run lint

Author: piochelepiotr

.gitignore

@@ -98,6 +98,7 @@ tmp
 .python-version
 
 # Common IDEs & editors
+.cursor/
 .idea/
 .vscode/
 *.iml

kafka_consumer/changelog.d/20512.added

@@ -0,0 +1 @@
+kafka_consumer check can retrieve messages from Kafka and log them.

kafka_consumer/datadog_checks/kafka_consumer/client.py

@@ -172,6 +172,13 @@ def list_consumer_group_offsets(self, groups):
             offsets.append((response_offset_info.group_id, tpo))
         return offsets
 
+    def start_collecting_messages(self, start_offsets):
+        self.open_consumer('datadog_live_messages')
+        self._consumer.assign(start_offsets)
+
+    def get_next_message(self):
+        return self._consumer.poll(timeout=1)
+
     def describe_consumer_group(self, consumer_group):
         desc = self.kafka_client.describe_consumer_groups([consumer_group])[consumer_group].result()
         return desc.state.name

kafka_consumer/datadog_checks/kafka_consumer/config.py

@@ -80,6 +80,9 @@ def __init__(self, init_config, instance, log) -> None:
         ):
             self._tls_ca_cert = '/opt/datadog-agent/embedded/ssl/certs/cacert.pem'
 
+        # Data Streams live messages
+        self.live_messages_configs = instance.get('live_messages_configs', [])
+
     def validate_config(self):
         if not self._kafka_connect_str:
             raise ConfigurationError('`kafka_connect_str` is required')
@@ -124,6 +127,42 @@ def validate_config(self):
             )
 
         self._validate_consumer_groups()
+        self._validate_live_messages_configs()
+
+    def _validate_live_messages_configs(self):
+        live_messages_configs = []
+        for config in self.live_messages_configs:
+            if 'id' not in config:
+                self.log.debug('Data Streams live messages configuration has no ID')
+                continue
+            kafka = config.get('kafka', None)
+            if not kafka:
+                self.log.debug('Data Streams live messages configuration has no kafka configuration')
+                continue
+            if not (
+                'cluster' in kafka
+                and 'topic' in kafka
+                and 'partition' in kafka
+                and 'start_offset' in kafka
+                and 'n_messages' in kafka
+            ):
+                self.log.debug('Data Streams live messages configuration missing required kafka parameters.', kafka)
+                continue
+            # Only json format is supported for Data Streams live messages
+            if kafka.get('value_format', '') == '':
+                kafka['value_format'] = 'json'
+            if kafka['value_format'] != 'json':
+                self.log.debug(
+                    'Only json format is supported for Data Streams live messages, got %s', kafka['value_format']
+                )
+            if kafka.get('key_format', '') == '':
+                kafka['key_format'] = 'json'
+            if kafka['key_format'] != 'json':
+                self.log.debug(
+                    'Only json format is supported for Data Streams live messages, got %s', kafka['key_format']
+                )
+            live_messages_configs.append(config)
+        self.live_messages_configs = live_messages_configs
 
     def _compile_regex(self, consumer_groups_regex, consumer_groups):
         # Turn the dict of regex dicts into a single string and compile

kafka_consumer/datadog_checks/kafka_consumer/kafka_consumer.py

@@ -5,12 +5,16 @@
 from collections import defaultdict
 from time import time
 
+from confluent_kafka import TopicPartition
+
 from datadog_checks.base import AgentCheck, is_affirmative
 from datadog_checks.kafka_consumer.client import KafkaClient
 from datadog_checks.kafka_consumer.config import KafkaConfig
 from datadog_checks.kafka_consumer.constants import KAFKA_INTERNAL_TOPICS, OFFSET_INVALID
 
 MAX_TIMESTAMPS = 1000
+SCHEMA_REGISTRY_MAGIC_BYTE = 0x00
+DATA_STREAMS_MESSAGES_CACHE_KEY = 'get_messages_cache'
 
 
 class KafkaCheck(AgentCheck):
@@ -96,6 +100,7 @@ def check(self, _):
         )
         if self.config._close_admin_client:
             self.client.close_admin_client()
+        self.data_streams_live_message(highwater_offsets or {}, cluster_id)
 
     def get_consumer_offsets(self):
         # {(consumer_group, topic, partition): offset}
@@ -175,6 +180,30 @@ def _load_broker_timestamps(self, persistent_cache_key):
             self.log.warning('Could not read broker timestamps from cache: %s', str(e))
         return broker_timestamps
 
+    def _messages_have_been_retrieved(self, config_id):
+        """Check if messages have been retrieved for the given config ID."""
+        try:
+            content = self.read_persistent_cache(DATA_STREAMS_MESSAGES_CACHE_KEY)
+            if content:
+                config_ids = set(content.split(","))
+                return config_id in config_ids
+        except Exception as e:
+            self.log.warning('Could not read persistent cache: %s', str(e))
+        return False
+
+    def _mark_messages_retrieved(self, config_id):
+        """Mark that messages have been retrieved for the given config ID."""
+        try:
+            content = self.read_persistent_cache(DATA_STREAMS_MESSAGES_CACHE_KEY)
+            if content:
+                config_ids = set(content.split(","))
+            else:
+                config_ids = set()
+            config_ids.add(config_id)
+            self.write_persistent_cache(DATA_STREAMS_MESSAGES_CACHE_KEY, ",".join(config_ids))
+        except Exception as e:
+            self.log.warning('Could not write to persistent cache: %s', str(e))
+
     def _add_broker_timestamps(self, broker_timestamps, highwater_offsets):
         for (topic, partition), highwater_offset in highwater_offsets.items():
             timestamps = broker_timestamps["{}_{}".format(topic, partition)]
@@ -378,6 +407,78 @@ def send_event(self, title, text, tags, event_type, aggregation_key, severity='i
         }
         self.event(event_dict)
 
+    def data_streams_live_message(self, highwater_offsets, cluster_id):
+        for cfg in self.config.live_messages_configs:
+            kafka = cfg['kafka']
+            topic = kafka["topic"]
+            partition = kafka["partition"]
+            start_offset = kafka["start_offset"]
+            n_messages = kafka["n_messages"]
+            cluster = kafka["cluster"]
+            config_id = cfg["id"]
+            if self._messages_have_been_retrieved(config_id):
+                continue
+            if cluster != cluster_id:
+                continue
+            start_offsets = resolve_start_offsets(highwater_offsets, topic, partition, start_offset, n_messages)
+
+            if not start_offsets:
+                self.log.warning('Unable to get a list of partitions to read from for live messages')
+                self.send_log(
+                    {
+                        'timestamp': int(time()),
+                        'config_id': config_id,
+                        'technology': 'kafka',
+                        'cluster': str(cluster),
+                        'topic': str(topic),
+                        'live_messages_error': 'Unable to list partitions to read from',
+                        'message': "Unable to list partitions to read from",
+                    }
+                )
+                continue
+
+            self.client.start_collecting_messages(start_offsets)
+            for _ in range(n_messages):
+                message = self.client.get_next_message()
+                if message is None:
+                    self.log.debug('Live messages: no message to retrieve')
+                    self.send_log(
+                        {
+                            'timestamp': int(time()),
+                            'config_id': config_id,
+                            'technology': 'kafka',
+                            'cluster': str(cluster),
+                            'topic': str(topic),
+                            'live_messages_error': 'No more messages to retrieve',
+                            'message': "No more messages to retrieve",
+                        }
+                    )
+                    break
+                data = {
+                    'timestamp': int(time()),
+                    'technology': 'kafka',
+                    'cluster': str(cluster),
+                    'config_id': config_id,
+                    'topic': str(topic),
+                    'partition': str(message.partition()),
+                    'offset': str(message.offset()),
+                }
+                decoded_value, value_schema_id, decoded_key, key_schema_id = deserialize_message(message)
+                if decoded_value:
+                    data['message_value'] = decoded_value
+                else:
+                    data['message'] = "Message format not supported"
+                    data['live_messages_error'] = 'Message format not supported'
+                if value_schema_id:
+                    data['value_schema_id'] = str(value_schema_id)
+                if decoded_key:
+                    data['message_key'] = decoded_key
+                if key_schema_id:
+                    data['key_schema_id'] = str(key_schema_id)
+                self.send_log(data)
+            self.client.close_consumer()
+            self._mark_messages_retrieved(config_id)
+
 
 def _get_interpolated_timestamp(timestamps, offset):
     if offset in timestamps:
@@ -406,3 +507,66 @@ def _get_interpolated_timestamp(timestamps, offset):
     slope = (timestamp_after - timestamp_before) / float(offset_after - offset_before)
     timestamp = slope * (offset - offset_after) + timestamp_after
     return timestamp
+
+
+def resolve_start_offsets(highwater_offsets, target_topic, target_partition, start_offset, n_messages):
+    if int(target_partition) == -1:
+        # in this case, we get n_messages, starting at offset latest - n_messages on each partition.
+        # this doesn't match exactly to the latest messages, but if we don't do that, we could run into
+        # edge cases when some partitions don't get any traffic.
+        start_offsets = []
+        for topic, partition in highwater_offsets:
+            if topic == target_topic and highwater_offsets[(topic, partition)] >= 0:
+                start_offsets.append(
+                    TopicPartition(topic, partition, max(0, highwater_offsets[(topic, partition)] - n_messages + 1))
+                )
+                if len(start_offsets) >= n_messages:
+                    break
+        return start_offsets
+    if int(start_offset) == -1:
+        end_offset = highwater_offsets.get((target_topic, target_partition), -1)
+        return (
+            []
+            if end_offset < 0
+            else [TopicPartition(target_topic, target_partition, max(0, end_offset - n_messages + 1))]
+        )
+    return [TopicPartition(target_topic, target_partition, start_offset)]
+
+
+def deserialize_message(message):
+    try:
+        decoded_value, value_schema_id = _deserialize_bytes_maybe_schema_registry(message.value())
+    except (UnicodeDecodeError, json.JSONDecodeError):
+        return None, None, None, None
+    try:
+        decoded_key, key_schema_id = _deserialize_bytes_maybe_schema_registry(message.key())
+        return decoded_value, value_schema_id, decoded_key, key_schema_id
+    except (UnicodeDecodeError, json.JSONDecodeError):
+        return decoded_value, value_schema_id, None, None
+
+
+def _deserialize_bytes_maybe_schema_registry(message):
+    try:
+        return _deserialize_bytes(message), None
+    except (UnicodeDecodeError, json.JSONDecodeError) as e:
+        # If the message is not a valid JSON, it might be a schema registry message, that is prefixed
+        # with a magic byte and a schema ID.
+        if len(message) < 5 or message[0] != SCHEMA_REGISTRY_MAGIC_BYTE:
+            raise e
+        schema_id = int.from_bytes(message[1:5], 'big')
+        message = message[5:]  # Skip the schema ID bytes
+        return _deserialize_bytes(message), schema_id
+
+
+def _deserialize_bytes(message):
+    """Deserialize a message from Kafka. Supports JSON format.
+    Args:
+        message: Raw message bytes from Kafka
+    Returns:
+        Decoded message as a string
+    """
+    if not message:
+        return ""
+    decoded = message.decode('utf-8')
+    json.loads(decoded)
+    return decoded

kafka_consumer/tests/runners.py

@@ -38,11 +38,33 @@ def run(self):
         while not self._shutdown_event.is_set():
             for partition in PARTITIONS:
                 try:
-                    producer.produce('marvel', b"Peter Parker", partition=partition)
-                    producer.produce('marvel', b"Bruce Banner", partition=partition)
-                    producer.produce('marvel', b"Tony Stark", partition=partition)
-                    producer.produce('marvel', b"Johhny Blaze", partition=partition)
-                    producer.produce('marvel', b"\xc2BoomShakalaka", partition=partition)
+                    producer.produce(
+                        'marvel',
+                        b'{"name": "Peter Parker", "age": 18, "transaction_amount": 123, "currency": "dollar"}',
+                        partition=partition,
+                    )
+                    producer.produce(
+                        'marvel',
+                        key=b'{"name": "Bruce Banner"}',
+                        value=b'\x00\x00\x00\x01\x5e{"name": "Bruce Banner",\
+ "age": 45, "transaction_amount": 456, "currency": "dollar"}',
+                        partition=partition,
+                    )
+                    producer.produce(
+                        'marvel',
+                        b'{"name": "Tony Stark", "age": 35, "transaction_amount": 789, "currency": "dollar"}',
+                        partition=partition,
+                    )
+                    producer.produce(
+                        'marvel',
+                        b'{"name": "Johnny Blaze", "age": 30, "transaction_amount": 321, "currency": "dollar"}',
+                        partition=partition,
+                    )
+                    producer.produce(
+                        'marvel',
+                        b'{"name": "BoomShakalaka", "age": 25, "transaction_amount": 654, "currency": "dollar"}',
+                        partition=partition,
+                    )
                     producer.produce('dc', b"Diana Prince", partition=partition)
                     producer.produce('dc', b"Bruce Wayne", partition=partition)
                     producer.produce('dc', b"Clark Kent", partition=partition)

kafka_consumer/tests/test_integration.py

@@ -466,3 +466,51 @@ def test_regex_consumer_groups(
     aggregator.assert_metric("kafka.estimated_consumer_lag", count=consumer_lag_seconds_count)
 
     assert expected_warning in caplog.text
+
+
+@mock.patch('datadog_checks.kafka_consumer.kafka_consumer.time', mocked_time)
+def test_data_streams_live_messages(dd_run_check, check, kafka_instance, datadog_agent):
+    cluster_id = common.get_cluster_id()
+    kafka_instance['live_messages_configs'] = [
+        {
+            'kafka': {
+                'cluster': cluster_id,
+                'topic': 'marvel',
+                'partition': 0,
+                'start_offset': 0,
+                'n_messages': 2,
+                'value_format': 'json',
+            },
+            'id': 'config_1_id',
+        }
+    ]
+    kafka_check = check(kafka_instance)
+    dd_run_check(kafka_check)
+    expected_logs = [
+        {
+            'timestamp': 400 * 1000,
+            'technology': 'kafka',
+            'cluster': str(cluster_id),
+            'config_id': 'config_1_id',
+            'topic': 'marvel',
+            'partition': '0',
+            'offset': '0',
+            'message_value': '{"name": "Peter Parker", "age": 18, "transaction_amount": 123, "currency": "dollar"}',
+            'ddtags': 'optional:tag1',
+        },
+        {
+            'timestamp': 400 * 1000,
+            'technology': 'kafka',
+            'cluster': str(cluster_id),
+            'config_id': 'config_1_id',
+            'topic': 'marvel',
+            'partition': '0',
+            'offset': '1',
+            'message_value': '{"name": "Bruce Banner", "age": 45,\
+ "transaction_amount": 456, "currency": "dollar"}',
+            'value_schema_id': '350',
+            'message_key': '{"name": "Bruce Banner"}',
+            'ddtags': 'optional:tag1',
+        },
+    ]
+    datadog_agent.assert_logs(kafka_check.check_id, expected_logs)