Clean up public API of data streams

Author: ericfirth

lib/datadog/tracing.rb

@@ -2,6 +2,7 @@
 
 require_relative 'core'
 require_relative 'tracing/pipeline'
+require_relative 'tracing/data_streams'
 
 module Datadog
   # Datadog APM tracing public API.

lib/datadog/tracing/contrib/kafka/instrumentation/consumer.rb

@@ -17,22 +17,22 @@ module InstanceMethods
               def each_message(**kwargs, &block)
                 # Wrap the block to add DSM processing for each message
                 wrapped_block = if Datadog.configuration.tracing.data_streams.enabled
-                  proc do |message|
-                    # DSM: Create checkpoint for consumed message
-                    Datadog.logger.debug { "Kafka each_message: DSM enabled for topic #{message.topic}" }
+                                  proc do |message|
+                                    # DSM: Create checkpoint for consumed message
+                                    Datadog.logger.debug { "Kafka each_message: DSM enabled for topic #{message.topic}" }
 
-                    processor = Datadog.configuration.tracing.data_streams.processor
+                                    processor = Datadog.configuration.tracing.data_streams.processor
 
-                    # Extract pathway context from message headers if available
-                    headers = message.headers || {}
-                    processor.set_consume_checkpoint('kafka', message.topic) { |key| headers[key] }
+                                    # Extract pathway context from message headers if available
+                                    headers = message.headers || {}
+                                    processor.set_consume_checkpoint('kafka', message.topic) { |key| headers[key] }
 
-                    # Call the original block if provided
-                    block.call(message) if block_given?
-                  end
-                else
-                  block
-                end
+                                    # Call the original block if provided
+                                    yield(message) if block
+                                  end
+                                else
+                                  block
+                                end
 
                 # Call the original method with wrapped block
                 super(**kwargs, &wrapped_block)
@@ -42,22 +42,22 @@ def each_message(**kwargs, &block)
               def each_batch(**kwargs, &block)
                 # Wrap the block to add DSM processing for each batch
                 wrapped_block = if Datadog.configuration.tracing.data_streams.enabled
-                  proc do |batch|
-                    # DSM: Create checkpoint for consumed batch
-                    Datadog.logger.debug { "Kafka each_batch: DSM enabled for topic #{batch.topic}" }
+                                  proc do |batch|
+                                    # DSM: Create checkpoint for consumed batch
+                                    Datadog.logger.debug { "Kafka each_batch: DSM enabled for topic #{batch.topic}" }
 
-                    processor = Datadog.configuration.tracing.data_streams.processor
+                                    processor = Datadog.configuration.tracing.data_streams.processor
 
-                    # For batch processing, we don't have individual message headers
-                    # so we create a consume checkpoint without pathway context
-                    processor.set_consume_checkpoint('kafka', batch.topic)
+                                    # For batch processing, we don't have individual message headers
+                                    # so we create a consume checkpoint without pathway context
+                                    processor.set_consume_checkpoint('kafka', batch.topic)
 
-                    # Call the original block if provided
-                    block.call(batch) if block_given?
-                  end
-                else
-                  block
-                end
+                                    # Call the original block if provided
+                                    yield(batch) if block
+                                  end
+                                else
+                                  block
+                                end
 
                 # Call the original method with wrapped block
                 super(**kwargs, &wrapped_block)

lib/datadog/tracing/data_streams.rb

@@ -0,0 +1,18 @@
+module Datadog
+  module Tracing
+    module DataStreams
+      def self.set_produce_checkpoint(type, target, manual_checkpoint: true, &block)
+        Datadog.configuration.tracing.data_streams.processor.set_produce_checkpoint(type, target, &block)
+      end
+
+      def self.set_consume_checkpoint(type, source, manual_checkpoint: true, &block)
+        Datadog.configuration.tracing.data_streams.processor.set_consume_checkpoint(
+          type,
+          source,
+          manual_checkpoint: manual_checkpoint,
+          &block
+        )
+      end
+    end
+  end
+end