fix(core): capture command data for SDK v3 clients (#611)

danno-s · SergeyPoluektov · web-flow · commit 5eb97ef43f71 · 2023-09-18T10:40:02.000-07:00
* fix(core): capture additional command data with SDK v3 * adds unit tests for abandoned PR #527 * adds output mapping to test * removes leftover util import * requested changes on #611 --------- Co-authored-by: SergeiPoluektov <sergei@hellohakuna.com>
diff --git a/packages/core/lib/patchers/aws3_p.ts b/packages/core/lib/patchers/aws3_p.ts
@@ -39,6 +39,7 @@ const buildAttributesFromMetadata = async (
   service: string,
   operation: string,
   region: string,
+  commandInput: any,
   res: any | null,
   error: SdkError | null,
 ): Promise<[ServiceSegment, HttpResponse]> => {
@@ -49,8 +50,10 @@ const buildAttributesFromMetadata = async (
       extendedRequestId,
       requestId,
       retryCount: attempts,
+      data: res?.output,
       request: {
         operation,
+        params: commandInput,
         httpRequest: {
           region,
           statusCode,
@@ -93,11 +96,13 @@ function addFlags(http: HttpResponse, subsegment: Subsegment, err?: SdkError): v
 const getXRayMiddleware = (config: RegionResolvedConfig, manualSegment?: SegmentLike): BuildMiddleware<any, any> => (next: any, context: any) => async (args: any) => {
   const segment = contextUtils.isAutomaticMode() ? contextUtils.resolveSegment() : manualSegment;
   const {clientName, commandName} = context;
-  const operation: string = commandName.slice(0, -7); // Strip trailing "Command" string
+  const commandInput = args?.input ?? {};
+  const commandOperation: string = commandName.slice(0, -7); // Strip trailing "Command" string
+  const operation: string = commandOperation.charAt(0).toLowerCase() + commandOperation.slice(1);
   const service: string = clientName.slice(0, -6);    // Strip trailing "Client" string
 
   if (!segment) {
-    const output = service + '.' + operation.charAt(0).toLowerCase() + operation.slice(1);
+    const output = service + '.' + operation;
 
     if (!contextUtils.isAutomaticMode()) {
       logger.getLogger().info('Call ' + output + ' requires a segment object' +
@@ -148,6 +153,7 @@ const getXRayMiddleware = (config: RegionResolvedConfig, manualSegment?: Segment
       service,
       operation,
       await config.region(),
+      commandInput,
       res,
       null,
     );
@@ -164,6 +170,7 @@ const getXRayMiddleware = (config: RegionResolvedConfig, manualSegment?: Segment
         service,
         operation,
         await config.region(),
+        commandInput,
         null,
         err,
       );
diff --git a/packages/core/test/unit/patchers/aws3_p.test.js b/packages/core/test/unit/patchers/aws3_p.test.js
@@ -18,38 +18,38 @@ chai.use(sinonChai);
 
 var traceId = '1-57fbe041-2c7ad569f5d6ff149137be86';
 
-describe('AWS v3 patcher', function() {
-  describe('#captureAWSClient', function() {
+describe('AWS v3 patcher', function () {
+  describe('#captureAWSClient', function () {
     var sandbox, useMiddleware;
 
     var awsClient = {
-      send: function() {},
+      send: function () { },
       config: {
         serviceId: 's3',
       },
       middlewareStack: constructStack(),
     };
 
-    beforeEach(function() {
+    beforeEach(function () {
       sandbox = sinon.createSandbox();
       useMiddleware = sandbox.stub(awsClient.middlewareStack, 'use');
     });
 
-    afterEach(function() {
+    afterEach(function () {
       sandbox.restore();
     });
 
-    it('should call middlewareStack.use and return the service', function() {
+    it('should call middlewareStack.use and return the service', function () {
       const patched = awsPatcher.captureAWSClient(awsClient);
       useMiddleware.should.have.been.calledOnce;
       assert.equal(patched, awsClient);
     });
   });
 
-  describe('#captureAWSRequest', function() {
-    var awsClient, awsRequest, sandbox, segment, stubResolve, addNewSubsegmentStub, sub;
+  describe('#captureAWSRequest', function () {
+    let awsClient, awsRequest, ddbClient, ddbGetRequest, sandbox, segment, stubResolve, addNewSubsegmentStub, sub;
 
-    before(function() {
+    before(function () {
       awsClient = {
         send: async (req) => {
           const context = {
@@ -74,9 +74,34 @@ describe('AWS v3 patcher', function() {
         },
         middlewareStack: constructStack(),
       };
+
+      ddbClient = {
+        send: async (req) => {
+          const middlewareContext = {
+            clientName: 'DynamoDBClient',
+            commandName: 'GetItemCommand',
+          };
+          const handler = awsClient.middlewareStack.resolve((args) => {
+            const error = req.response.error;
+            if (error) {
+              const err = new Error(error.message);
+              err.name = error.code;
+              err.$metadata = req.response.$metadata;
+              throw err;
+            }
+            return args;
+          }, middlewareContext);
+          await handler(req);
+          return req.response;
+        },
+        config: {
+          region: async () => 'us-east-1',
+        },
+        middlewareStack: constructStack(),
+      };
     });
 
-    beforeEach(function() {
+    beforeEach(function () {
       sandbox = sinon.createSandbox();
 
       awsRequest = new (class ListBucketsCommand {
@@ -99,14 +124,39 @@ describe('AWS v3 patcher', function() {
         }
       })();
 
+      ddbGetRequest = new (class GetItemCommand {
+        constructor() {
+          this.request = {
+            method: 'GET',
+            url: '/',
+            connection: {
+              remoteAddress: 'localhost'
+            },
+            headers: {},
+          };
+          this.input = {
+            TableName: 'TestTableName',
+            ConsistentRead: true,
+          };
+          this.response = {};
+          this.output = {
+            ConsumedCapacity: 10,
+            $metadata: {
+              requestId: '123',
+              extendedRequestId: '456',
+            },
+          };
+        }
+      })();
+
       segment = new Segment('testSegment', traceId);
-      segment.additionalTraceData = {'Foo': 'bar'};
+      segment.additionalTraceData = { 'Foo': 'bar' };
       sub = segment.addNewSubsegment('subseg');
       stubResolve = sandbox.stub(contextUtils, 'resolveSegment').returns(segment);
       addNewSubsegmentStub = sandbox.stub(segment, 'addNewSubsegment').returns(sub);
     });
 
-    afterEach(function() {
+    afterEach(function () {
       sandbox.restore();
     });
 
@@ -125,13 +175,13 @@ describe('AWS v3 patcher', function() {
 
         awsClient.send(awsRequest);
 
-        setTimeout(function() {
+        setTimeout(function () {
           logStub.should.have.been.calledOnce;
           done();
         }, 50);
       });
 
-      it('should inject the tracing headers', async function() {
+      it('should inject the tracing headers', async function () {
         await awsClient.send(awsRequest);
         assert.isTrue(addNewSubsegmentStub.calledWith('S3'));
 
@@ -140,7 +190,7 @@ describe('AWS v3 patcher', function() {
         assert.match(awsRequest.request.headers['X-Amzn-Trace-Id'], expected);
       });
 
-      it('should close on complete with no errors when code 200 is seen', async function() {
+      it('should close on complete with no errors when code 200 is seen', async function () {
         const closeStub = sandbox.stub(sub, 'close').returns();
         sandbox.stub(sub, 'addAttribute').returns();
         sandbox.stub(Aws.prototype, 'init').returns();
@@ -154,7 +204,7 @@ describe('AWS v3 patcher', function() {
         closeStub.should.have.been.calledWithExactly();
       });
 
-      it('should mark the subsegment as throttled and error if code 429 is seen', async function() {
+      it('should mark the subsegment as throttled and error if code 429 is seen', async function () {
         const throttleStub = sandbox.stub(sub, 'addThrottleFlag').returns();
 
         sandbox.stub(sub, 'addAttribute').returns();
@@ -171,7 +221,7 @@ describe('AWS v3 patcher', function() {
         assert.isTrue(sub.error);
       });
 
-      it('should mark the subsegment as throttled and error if code service.throttledError returns true, regardless of status code', async function() {
+      it('should mark the subsegment as throttled and error if code service.throttledError returns true, regardless of status code', async function () {
         const throttleStub = sandbox.stub(sub, 'addThrottleFlag').returns();
 
         sandbox.stub(sub, 'addAttribute').returns();
@@ -188,7 +238,7 @@ describe('AWS v3 patcher', function() {
         assert.isTrue(sub.error);
       });
 
-      it('should capture an error on the response and mark exception as remote', async function() {
+      it('should capture an error on the response and mark exception as remote', async function () {
         const closeStub = sandbox.stub(sub, 'close').returns();
         const getCauseStub = sandbox.stub(Utils, 'getCauseTypeFromHttpStatus').returns();
 
@@ -205,7 +255,15 @@ describe('AWS v3 patcher', function() {
         await awsClient.send(awsRequest).catch(() => null);
 
         getCauseStub.should.have.been.calledWithExactly(500);
-        closeStub.should.have.been.calledWithExactly(sinon.match({ message: error.message, name: error.code}), true);
+        closeStub.should.have.been.calledWithExactly(sinon.match({ message: error.message, name: error.code }), true);
+      });
+
+      it('should pass params into aws subsegment', async function () {
+        await ddbClient.send(ddbGetRequest);
+
+        assert.isTrue(addNewSubsegmentStub.calledWith('DynamoDB'));
+        addNewSubsegmentStub.returnValues[0].should.have.property('aws');
+        addNewSubsegmentStub.returnValues[0].aws.should.include({ consistent_read: true, table_name: 'TestTableName', consumed_capacity: 10 });
       });
     });
 
@@ -220,7 +278,7 @@ describe('AWS v3 patcher', function() {
 
         awsClient.send(awsRequest);
 
-        setTimeout(function() {
+        setTimeout(function () {
           logStub.should.have.been.calledOnce;
           done();
         }, 50);
@@ -250,10 +308,10 @@ describe('AWS v3 patcher', function() {
   });
 
 
-  describe('#captureAWSRequest-Unsampled', function() {
+  describe('#captureAWSRequest-Unsampled', function () {
     var awsClient, awsRequest, sandbox, segment, stubResolve, addNewSubsegmentStub, sub, service, addNewServiceSubsegmentStub;
 
-    before(function() {
+    before(function () {
       awsClient = {
         send: async (req) => {
           const context = {
@@ -280,7 +338,7 @@ describe('AWS v3 patcher', function() {
       };
     });
 
-    beforeEach(function() {
+    beforeEach(function () {
       sandbox = sinon.createSandbox();
 
       awsRequest = new (class ListBucketsCommand {
@@ -304,7 +362,7 @@ describe('AWS v3 patcher', function() {
       })();
 
       segment = new Segment('testSegment', traceId);
-      segment.additionalTraceData = {'Foo': 'bar'};
+      segment.additionalTraceData = { 'Foo': 'bar' };
       sub = segment.addNewSubsegmentWithoutSampling('subseg');
       service = sub.addNewSubsegmentWithoutSampling('service');
 
@@ -314,7 +372,7 @@ describe('AWS v3 patcher', function() {
 
     });
 
-    afterEach(function() {
+    afterEach(function () {
       sandbox.restore();
     });
 
@@ -328,7 +386,7 @@ describe('AWS v3 patcher', function() {
       });
 
 
-      it('should inject the tracing headers', async function() {
+      it('should inject the tracing headers', async function () {
         await awsClient.send(awsRequest);
         assert.isTrue(addNewServiceSubsegmentStub.calledWith('S3'));
 
