Add context to publish batch

Author: dormesica1

plugins/node/opentelemetry-instrumentation-aws-sdk/src/services/sns.ts

@@ -25,6 +25,7 @@ import {
   NormalizedResponse,
   AwsSdkInstrumentationConfig,
 } from '../types';
+import type { SNS } from '../aws-sdk.types';
 import { injectPropagationContext } from './MessageAttributes';
 import { RequestMetadata, ServiceExtension } from './ServiceExtension';
 
@@ -67,14 +68,35 @@ export class SnsServiceExtension implements ServiceExtension {
   }
 
   requestPostSpanHook(request: NormalizedRequest): void {
-    if (request.commandName === 'Publish') {
-      const origMessageAttributes =
-        request.commandInput['MessageAttributes'] ?? {};
-      if (origMessageAttributes) {
-        request.commandInput['MessageAttributes'] = injectPropagationContext(
-          origMessageAttributes
-        );
-      }
+    switch (request.commandName) {
+      case 'Publish':
+        {
+          const origMessageAttributes =
+            request.commandInput['MessageAttributes'] ?? {};
+          if (origMessageAttributes) {
+            request.commandInput['MessageAttributes'] = injectPropagationContext(
+              origMessageAttributes
+            );
+          }
+        }
+        break;
+
+      case 'PublishBatch':
+        {
+          const entries = request.commandInput?.PublishBatchRequestEntries;
+          if (Array.isArray(entries)) {
+            entries.forEach(
+              (messageParams: {
+                MessageAttributes: SNS.MessageAttributeMap;
+              }) => {
+                messageParams.MessageAttributes = injectPropagationContext(
+                  messageParams.MessageAttributes ?? {}
+                );
+              }
+            );
+          }
+        }
+        break;
     }
   }
 

plugins/node/opentelemetry-instrumentation-aws-sdk/test/mock-responses/sns-publish-batch.xml

@@ -0,0 +1 @@
+<?xml version="1.0"?><PublishBatchResponse xmlns="http://sns.amazonaws.com/doc/2010-03-31/"><PublishBatchResult><Successful><MessageId>90d1987b-4853-54ad-a499-c2d89c4edf3a</MessageId><Successful></PublishBatchResult><ResponseMetadata><RequestId>d81e4f4f-2d70-51f3-a040-15ecf96d5a64</RequestId></ResponseMetadata></PublishBatchResponse>

plugins/node/opentelemetry-instrumentation-aws-sdk/test/sns.test.ts

@@ -47,6 +47,11 @@ const responseMockSuccess = {
 const topicName = 'topic';
 const fakeARN = `arn:aws:sns:region:000000000:${topicName}`;
 
+const hookSpy = sinon.spy(
+  (instrumentation['servicesExtensions'] as any)['services'].get('SNS'),
+  'requestPostSpanHook'
+);
+
 describe('SNS - v2', () => {
   before(() => {
     AWSv2.config.credentials = {
@@ -64,6 +69,10 @@ describe('SNS - v2', () => {
     } as AWS.SNS.Types.PublishResponse);
   });
 
+  afterEach(() => {
+    hookSpy.resetHistory();
+  })
+
   describe('publish', () => {
     it('topic arn', async () => {
       const sns = new AWSv2.SNS();
@@ -120,10 +129,6 @@ describe('SNS - v2', () => {
 
     it('inject context propagation', async () => {
       const sns = new AWSv2.SNS();
-      const hookSpy = sinon.spy(
-        (instrumentation['servicesExtensions'] as any)['services'].get('SNS'),
-        'requestPostSpanHook'
-      );
 
       await sns
         .publish({
@@ -180,16 +185,18 @@ describe('SNS - v3', () => {
         secretAccessKey: 'abcde',
       },
     });
+  });
 
-    nock('https://sns.us-east-1.amazonaws.com/')
+  describe('publish', () => {
+    beforeEach(() => {
+      nock('https://sns.us-east-1.amazonaws.com/')
       .post('/')
       .reply(
         200,
         fs.readFileSync('./test/mock-responses/sns-publish.xml', 'utf8')
       );
-  });
+    });
 
-  describe('publish', () => {
     it('topic arn', async () => {
       const topicV3Name = 'dummy-sns-v3-topic';
       const topicV3ARN = `arn:aws:sns:us-east-1:000000000:${topicV3Name}`;
@@ -236,4 +243,29 @@ describe('SNS - v3', () => {
       );
     });
   });
+
+  describe('publish batch', () => {
+    it('inject context propagation for publish batch command', async () => {
+      nock('https://sns.us-east-1.amazonaws.com/')
+        .post('/')
+        .reply(
+          200,
+          fs.readFileSync('./test/mock-responses/sns-publish-batch.xml', 'utf8')
+        );
+
+      await sns
+        .publishBatch({
+          TopicArn: fakeARN,
+          PublishBatchRequestEntries: [{ Id: '1', Message: 'sns message' }]
+        });
+
+      const publishSpans = getTestSpans().filter(
+        (s: ReadableSpan) => s.name === `${topicName} send`
+      );
+      expect(publishSpans.length).toBe(1);
+      expect(
+        hookSpy.args[0][0].commandInput.PublishBatchRequestEntries[0].MessageAttributes.traceparent
+      ).toBeDefined();
+    });
+  });
 });