wip

lib/datadog/tracing/contrib/aws/configuration/settings.rb

@@ -68,6 +68,7 @@ class Settings < Contrib::Configuration::Settings
               o.env Ext::ENV_TRACE_PARENTAGE_STYLE
               o.default 'distributed'
             end
+          end
         end
       end
     end

lib/datadog/tracing/contrib/aws/service/base.rb

@@ -20,8 +20,8 @@ def extract_propagation!(context)
               Tracing.continue_trace!(Contrib.extract(datadog))
             end
 
-            def inject_propagation(trace, context, data_type)
-              message_attributes = (context.params[:message_attributes] ||= {})
+            def inject_propagation(trace, params, data_type)
+              message_attributes = (params[:message_attributes] ||= {})
               return if message_attributes.size >= MESSAGE_ATTRIBUTES_LIMIT
 
               data = {}

lib/datadog/tracing/contrib/aws/service/sns.rb

@@ -13,18 +13,14 @@ module Service
           # SNS tag handlers.
           class SNS < Base
             def process(config, trace, context)
-              inject_propagation(trace, context, 'Binary') if config[:propagation]
-
               return unless config[:propagation]
 
               case context.operation
               when :publish
-                inject_propagation(trace, context, 'Binary')
+                inject_propagation(trace, context.params, 'Binary')
               when :publish_batch
-                if config[:batch_propagation]
-                  inject_propagation(trace, context, 'Binary')
-                else
-                  inject_propagation(trace, context, 'Binary')
+                context.params[:publish_batch_request_entries].each do |entry|
+                  inject_propagation(trace, entry, 'Binary')
                 end
               end
             end

spec/datadog/tracing/contrib/aws/instrumentation_spec.rb

@@ -475,8 +475,59 @@
     end
 
     describe '#publish_batch' do
-      # TODO: Implement this test
-      # Should tag all messages, not just one.
+      subject!(:publish_batch) do
+        client.publish_batch(
+          topic_arn: 'arn:aws:sns:us-west-2:123456789012:my-topic-name',
+          publish_batch_request_entries: [
+            { id: 'id1', message: 'body1' },
+            { id: 'id2', message: 'body2' },
+          ]
+        )
+      end
+
+      let(:configuration_options) { super().merge(propagation: true) }
+
+      let(:responses) do
+        {
+          publish_batch: {
+            successful: [{ id: 'id1' }],
+            failed: [{ id: 'id2', code: 'error_code', sender_fault: true }]
+          }
+        }
+      end
+
+      it_behaves_like 'schema version span'
+      it_behaves_like 'environment service name', 'DD_TRACE_AWS_SERVICE_NAME'
+      it_behaves_like 'configured peer service span', 'DD_TRACE_AWS_PEER_SERVICE'
+      it_behaves_like 'a peer service span' do
+        let(:peer_service_val) { 'my-topic-name' }
+        let(:peer_service_source) { 'topicname' }
+      end
+
+      it 'generates a span' do
+        expect(span.name).to eq('aws.command')
+        expect(span.service).to eq('aws')
+        expect(span.type).to eq('http')
+        expect(span.resource).to eq('sns.publish_batch')
+
+        expect(span.get_tag('aws.agent')).to eq('aws-sdk-ruby')
+        expect(span.get_tag('aws.operation')).to eq('publish_batch')
+        expect(span.get_tag('region')).to eq('us-stubbed-1')
+        expect(span.get_tag('aws_service')).to eq('sns')
+        expect(span.get_tag('aws_account')).to eq('123456789012')
+        expect(span.get_tag('topicname')).to eq('my-topic-name')
+        expect(span.get_tag('path')).to eq('')
+        expect(span.get_tag('host')).to eq('sns.us-stubbed-1.amazonaws.com')
+        expect(span.get_tag('http.method')).to eq('POST')
+        expect(span.get_tag('http.status_code')).to eq('200')
+        expect(span.get_tag('span.kind')).to eq('client')
+
+        expect(span.get_tag(Datadog::Tracing::Metadata::Ext::TAG_COMPONENT)).to eq('aws')
+        expect(span.get_tag(Datadog::Tracing::Metadata::Ext::TAG_OPERATION))
+          .to eq('command')
+        expect(span.get_tag(Datadog::Tracing::Metadata::Ext::TAG_PEER_HOSTNAME))
+          .to eq('sns.us-stubbed-1.amazonaws.com')
+      end
     end
 
     describe '#create_topic' do

spec/datadog/tracing/contrib/aws/service/sns_spec.rb

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require 'datadog/tracing/contrib/aws/service/sns'
+require 'datadog/tracing/contrib/aws/parsed_context'
 
 RSpec.describe Datadog::Tracing::Contrib::Aws::Service::SNS do
   let(:span) { instance_double('Span') }
@@ -43,8 +44,9 @@
 
     let(:config) { { propagation: true } }
     let(:trace) { Datadog::Tracing::TraceOperation.new(id: trace_id, parent_span_id: span_id) }
-    let(:context) { instance_double('Context', params: params) }
+    let(:context) { instance_double(Datadog::Tracing::Contrib::Aws::ParsedContext, params: params, operation: operation) }
     let(:params) { {} }
+    let(:operation) { :publish }
     let(:trace_id) { 1 }
     let(:span_id) { 2 }
 
@@ -94,6 +96,44 @@
         expect { inject_propagation }.to_not(change { params })
       end
     end
+
+    context 'with batch publishing' do
+      let(:operation) { :publish_batch }
+
+      # resp = client.publish_batch({
+      #                               topic_arn: "topicARN", # required
+      #                               publish_batch_request_entries: [ # required
+      #                                 {
+      #                                   id: "String", # required
+      #                                   message: "message", # required
+      #                                   subject: "subject",
+      #                                   message_structure: "messageStructure",
+      #                                   message_attributes: {
+      #                                     "String" => {
+      #                                       data_type: "String", # required
+      #                                       string_value: "String",
+      #                                       binary_value: "data",
+      #                                     },
+      #                                   },
+      #                                   message_deduplication_id: "String",
+      #                                   message_group_id: "String",
+      #                                 },
+      #                               ],
+      #                             })
+
+      it 'adds a propagation attribute' do
+        inject_propagation
+        expect(params[:message_attributes]).to eq(
+                                                 '_datadog' => {
+                                                   binary_value:
+                                                     '{"x-datadog-trace-id":"1","x-datadog-parent-id":"2",' \
+                                                       '"traceparent":"00-00000000000000000000000000000001-0000000000000002-00",' \
+                                                       '"tracestate":"dd=p:0000000000000002"}',
+                                                   data_type: 'Binary'
+                                                 }
+                                               )
+      end
+    end
   end
 
   it_behaves_like 'injects attribute propagation' do