From b2739fe2745181350a9e3b44806f94042dec850a Mon Sep 17 00:00:00 2001
From: Eric Koslow <ekoslow@gmail.com>
Date: Tue, 9 Nov 2021 10:02:36 -0800
Subject: [PATCH] Fix aws-sdk distributed tracing when using .promise()

When using `.promise()` there is no callback for `dd-trace-js` to wrap.
Before this change, that meant we never called tracer.extract on any
incoming requests.

This changes the functionality by creating a "fake" callback for us to
wrap and then immediately calling it.

Fixes #1680
---
 packages/datadog-plugin-aws-sdk/src/index.js  | 19 +++++++++-
 .../datadog-plugin-aws-sdk/test/sqs.spec.js   | 37 +++++++++++++++++++
 2 files changed, 55 insertions(+), 1 deletion(-)

diff --git a/packages/datadog-plugin-aws-sdk/src/index.js b/packages/datadog-plugin-aws-sdk/src/index.js
index 72a3b63d137..8240a6e6904 100644
--- a/packages/datadog-plugin-aws-sdk/src/index.js
+++ b/packages/datadog-plugin-aws-sdk/src/index.js
@@ -49,7 +49,24 @@ function createWrapRequest (tracer, config) {
         if (typeof cb === 'function') {
           arguments[0] = awsHelpers.wrapCb(cb, serviceIdentifier, tags, request, tracer, childOf)
         }
-        return send.apply(this, arguments)
+        const resp = send.apply(this, arguments)
+        // If using the .promise() function, create a fake callback to wrap
+        if (typeof resp.then === 'function') {
+          return resp.then(response => {
+            return new Promise(resolve => {
+              const wrappedCb = awsHelpers.wrapCb(
+                () => { resolve(response) },
+                serviceIdentifier,
+                tags,
+                request,
+                tracer,
+                childOf
+              )
+              wrappedCb(null, response)
+            })
+          })
+        }
+        return resp
       })
     }
   }
diff --git a/packages/datadog-plugin-aws-sdk/test/sqs.spec.js b/packages/datadog-plugin-aws-sdk/test/sqs.spec.js
index a1f8035f4d0..ca291459d3a 100644
--- a/packages/datadog-plugin-aws-sdk/test/sqs.spec.js
+++ b/packages/datadog-plugin-aws-sdk/test/sqs.spec.js
@@ -87,6 +87,43 @@ describe('Plugin', () => {
           })
         })
 
+        it('should propagate the tracing context from the producer to the consumer using .promise()', (done) => {
+          // No need to test if version does not support .promise();
+          if (typeof AWS.Request.prototype.promise !== 'function') {
+            done()
+            return
+          }
+          let parentId
+          let traceId
+
+          agent.use(traces => {
+            const span = traces[0][0]
+
+            expect(span.resource.startsWith('sendMessage')).to.equal(true)
+
+            parentId = span.span_id.toString()
+            traceId = span.trace_id.toString()
+          })
+
+          agent.use(traces => {
+            const span = traces[0][0]
+
+            expect(parentId).to.be.a('string')
+            expect(span.parent_id.toString()).to.equal(parentId)
+            expect(span.trace_id.toString()).to.equal(traceId)
+          }).then(done, done)
+
+          sqs.sendMessage({
+            MessageBody: 'test body',
+            QueueUrl
+          }).promise().then(() => {
+            return sqs.receiveMessage({
+              QueueUrl,
+              MessageAttributeNames: ['.*']
+            }).promise()
+          }).catch(done)
+        })
+
         it('should run the consumer in the context of its span', (done) => {
           sqs.sendMessage({
             MessageBody: 'test body',