From 81b51aaa9108750f0c1f0723f498d7b704bcc73f Mon Sep 17 00:00:00 2001 From: Dor Mesica Date: Wed, 12 Feb 2025 13:16:41 +0200 Subject: [PATCH 1/2] Add context to publish batch --- .../src/services/sns.ts | 38 +++++++++++---- .../test/mock-responses/sns-publish-batch.xml | 1 + .../test/sns.test.ts | 46 ++++++++++++++++--- 3 files changed, 70 insertions(+), 15 deletions(-) create mode 100644 plugins/node/opentelemetry-instrumentation-aws-sdk/test/mock-responses/sns-publish-batch.xml diff --git a/plugins/node/opentelemetry-instrumentation-aws-sdk/src/services/sns.ts b/plugins/node/opentelemetry-instrumentation-aws-sdk/src/services/sns.ts index f6f6182253..15f6fbbc12 100644 --- a/plugins/node/opentelemetry-instrumentation-aws-sdk/src/services/sns.ts +++ b/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; } } diff --git a/plugins/node/opentelemetry-instrumentation-aws-sdk/test/mock-responses/sns-publish-batch.xml b/plugins/node/opentelemetry-instrumentation-aws-sdk/test/mock-responses/sns-publish-batch.xml new file mode 100644 index 0000000000..6422466882 --- /dev/null +++ b/plugins/node/opentelemetry-instrumentation-aws-sdk/test/mock-responses/sns-publish-batch.xml @@ -0,0 +1 @@ +90d1987b-4853-54ad-a499-c2d89c4edf3ad81e4f4f-2d70-51f3-a040-15ecf96d5a64 diff --git a/plugins/node/opentelemetry-instrumentation-aws-sdk/test/sns.test.ts b/plugins/node/opentelemetry-instrumentation-aws-sdk/test/sns.test.ts index f8d4dd777f..3a4bba58d4 100644 --- a/plugins/node/opentelemetry-instrumentation-aws-sdk/test/sns.test.ts +++ b/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(); + }); + }); }); From eb3db00a37a4d4f7d84b37e4e65f4c3980fdc47d Mon Sep 17 00:00:00 2001 From: Dor Mesica Date: Wed, 12 Feb 2025 13:22:16 +0200 Subject: [PATCH 2/2] Fix --- .../test/sns.test.ts | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/plugins/node/opentelemetry-instrumentation-aws-sdk/test/sns.test.ts b/plugins/node/opentelemetry-instrumentation-aws-sdk/test/sns.test.ts index 3a4bba58d4..f393be40d5 100644 --- a/plugins/node/opentelemetry-instrumentation-aws-sdk/test/sns.test.ts +++ b/plugins/node/opentelemetry-instrumentation-aws-sdk/test/sns.test.ts @@ -71,7 +71,7 @@ describe('SNS - v2', () => { afterEach(() => { hookSpy.resetHistory(); - }) + }); describe('publish', () => { it('topic arn', async () => { @@ -187,6 +187,10 @@ describe('SNS - v3', () => { }); }); + afterEach(() => { + hookSpy.resetHistory(); + }); + describe('publish', () => { beforeEach(() => { nock('https://sns.us-east-1.amazonaws.com/') @@ -245,14 +249,16 @@ describe('SNS - v3', () => { }); describe('publish batch', () => { - it('inject context propagation for publish batch command', async () => { + beforeEach(() => { nock('https://sns.us-east-1.amazonaws.com/') .post('/') .reply( 200, fs.readFileSync('./test/mock-responses/sns-publish-batch.xml', 'utf8') ); + }); + it('inject context propagation for publish batch command', async () => { await sns .publishBatch({ TopicArn: fakeARN,