fix(tracer): properly return DynamoDB.DocumentClient (#528)

* fix: properly return DynamoDB.DocumentClient

* added comment to explain fix

* updated e2e tests

* fix: applied proper generic

Author: dreamorosi

packages/tracing/src/Tracer.ts

@@ -261,7 +261,12 @@ class Tracer implements TracerInterface {
       return this.provider.captureAWSClient(service);
     } catch (error) {
       try {
-        return this.provider.captureAWSClient((service as unknown as T & { service: T }).service);
+        // This is needed because some aws-sdk clients like AWS.DynamoDB.DocumentDB don't comply with the same
+        // instrumentation contract like most base clients. 
+        // For detailed explanation see: https://github.com/awslabs/aws-lambda-powertools-typescript/issues/524#issuecomment-1024493662
+        this.provider.captureAWSClient((service as T & { service: T }).service);
+        
+        return service;
       } catch {
         throw error;
       }

packages/tracing/tests/e2e/tracer.test.Decorator.ts

@@ -1,6 +1,6 @@
 import { Tracer } from '../../src';
 import { Callback, Context } from 'aws-lambda';
-import { DynamoDBClient, ScanCommand } from '@aws-sdk/client-dynamodb';
+import { DynamoDBClient, PutItemCommand } from '@aws-sdk/client-dynamodb';
 // eslint-disable-next-line @typescript-eslint/no-var-requires
 let AWS = require('aws-sdk');
 
@@ -52,8 +52,8 @@ export class MyFunctionWithDecorator {
     }
 
     return Promise.all([
-      dynamoDBv2.scan({ TableName: testTableName }).promise(),
-      dynamoDBv3.send(new ScanCommand({ TableName: testTableName })),
+      dynamoDBv2.put({ TableName: testTableName, Item: { id: `${serviceName}-${event.invocation}-sdkv2` } }).promise(),
+      dynamoDBv3.send(new PutItemCommand({ TableName: testTableName, Item: { id: { 'S': `${serviceName}-${event.invocation}-sdkv3` } } })),
       new Promise((resolve, reject) => {
         setTimeout(() => {
           const res = this.myMethod();

packages/tracing/tests/e2e/tracer.test.DecoratorWithAsyncHandler.ts

@@ -1,6 +1,6 @@
 import { Tracer } from '../../src';
 import { Context } from 'aws-lambda';
-import { DynamoDBClient, ScanCommand } from '@aws-sdk/client-dynamodb';
+import { DynamoDBClient, PutItemCommand } from '@aws-sdk/client-dynamodb';
 // eslint-disable-next-line @typescript-eslint/no-var-requires
 let AWS = require('aws-sdk');
 
@@ -52,13 +52,13 @@ export class MyFunctionWithDecorator {
     }
 
     try {
-      await dynamoDBv2.scan({ TableName: testTableName }).promise();
+      await dynamoDBv2.put({ TableName: testTableName, Item: { id: `${serviceName}-${event.invocation}-sdkv2` } }).promise();
     } catch (err) {
       console.error(err);
     }
 
     try {
-      await dynamoDBv3.send(new ScanCommand({ TableName: testTableName }));
+      await dynamoDBv3.send(new PutItemCommand({ TableName: testTableName, Item: { id: { 'S': `${serviceName}-${event.invocation}-sdkv3` } } }));
     } catch (err) {
       console.error(err);
     }

packages/tracing/tests/e2e/tracer.test.Manual.ts

@@ -1,6 +1,6 @@
 import { Tracer } from '../../src';
 import { Context } from 'aws-lambda';
-import { DynamoDBClient, ScanCommand } from '@aws-sdk/client-dynamodb';
+import { DynamoDBClient, PutItemCommand } from '@aws-sdk/client-dynamodb';
 // eslint-disable-next-line @typescript-eslint/no-var-requires
 let AWS = require('aws-sdk');
 
@@ -55,13 +55,13 @@ export const handler = async (event: CustomEvent, _context: Context): Promise<vo
     dynamoDBv2 = new AWS.DynamoDB.DocumentClient();
   }
   try {
-    await dynamoDBv2.scan({ TableName: testTableName }).promise();
+    await dynamoDBv2.put({ TableName: testTableName, Item: { id: `${serviceName}-${event.invocation}-sdkv2` } }).promise();
   } catch (err) {
     console.error(err);
   }
 
   try {
-    await dynamoDBv3.send(new ScanCommand({ TableName: testTableName }));
+    await dynamoDBv3.send(new PutItemCommand({ TableName: testTableName, Item: { id: { 'S': `${serviceName}-${event.invocation}-sdkv3` } } }));
   } catch (err) {
     console.error(err);
   }

packages/tracing/tests/e2e/tracer.test.Middleware.ts

@@ -1,7 +1,7 @@
 import middy from '@middy/core';
 import { captureLambdaHandler, Tracer } from '../../src';
 import { Context } from 'aws-lambda';
-import { DynamoDBClient, ScanCommand } from '@aws-sdk/client-dynamodb';
+import { DynamoDBClient, PutItemCommand } from '@aws-sdk/client-dynamodb';
 // eslint-disable-next-line @typescript-eslint/no-var-requires
 let AWS = require('aws-sdk');
 
@@ -49,13 +49,13 @@ export const handler = middy(async (event: CustomEvent, _context: Context): Prom
     dynamoDBv2 = new AWS.DynamoDB.DocumentClient();
   }
   try {
-    await dynamoDBv2.scan({ TableName: testTableName }).promise();
+    await dynamoDBv2.put({ TableName: testTableName, Item: { id: `${serviceName}-${event.invocation}-sdkv2` } }).promise();
   } catch (err) {
     console.error(err);
   }
 
   try {
-    await dynamoDBv3.send(new ScanCommand({ TableName: testTableName }));
+    await dynamoDBv3.send(new PutItemCommand({ TableName: testTableName, Item: { id: { 'S': `${serviceName}-${event.invocation}-sdkv3` } } }));
   } catch (err) {
     console.error(err);
   }

packages/tracing/tests/e2e/tracer.test.ts

@@ -92,7 +92,7 @@ describe('Tracer integration tests', () => {
         },
         timeout: Duration.seconds(30),
       });
-      table.grantReadData(fn);
+      table.grantWriteData(fn);
       invocationsMap[functionName] = {
         serviceName: expectedServiceName,
         resourceArn: `arn:aws:lambda:${region}:${account}:function:${functionName}`, // ARN is still a token at this point, so we construct the ARN manually

packages/tracing/tests/unit/Tracer.test.ts

@@ -7,6 +7,7 @@
 import { Tracer } from '../../src';
 import { Callback, Context, Handler } from 'aws-lambda/handler';
 import { Segment, setContextMissingStrategy, Subsegment } from 'aws-xray-sdk-core';
+import { DynamoDB } from 'aws-sdk';
 
 interface LambdaInterface {
   handler: Handler
@@ -1057,29 +1058,27 @@ describe('Class: Tracer', () => {
       const captureAWSClientSpy = jest.spyOn(tracer.provider, 'captureAWSClient');
 
       // Act
-      tracer.captureAWSClient({});
+      const client = tracer.captureAWSClient(new DynamoDB());
 
       // Assess
       expect(captureAWSClientSpy).toBeCalledTimes(0);
+      expect(client).toBeInstanceOf(DynamoDB);
 
     });
 
-    test('when called with a simple AWS SDK v2 client, it returns it back instrumented', () => {
+    test('when called with a base AWS SDK v2 client, it returns it back instrumented', () => {
 
       // Prepare
       const tracer: Tracer = new Tracer();
       const captureAWSClientSpy = jest.spyOn(tracer.provider, 'captureAWSClient');
-      // Minimum shape required for a regular AWS v2 client (i.e. AWS.S3) to be instrumented
-      const dummyClient = {
-        customizeRequests: () => null,
-      };
 
       // Act
-      tracer.captureAWSClient(dummyClient);
+      const client = tracer.captureAWSClient(new DynamoDB());
 
       // Assess
       expect(captureAWSClientSpy).toBeCalledTimes(1);
-      expect(captureAWSClientSpy).toBeCalledWith(dummyClient);
+      expect(captureAWSClientSpy).toBeCalledWith(client);
+      expect(client).toBeInstanceOf(DynamoDB);
 
     });
 
@@ -1088,20 +1087,15 @@ describe('Class: Tracer', () => {
       // Prepare
       const tracer: Tracer = new Tracer();
       const captureAWSClientSpy = jest.spyOn(tracer.provider, 'captureAWSClient');
-      // Minimum shape required for a complex AWS v2 client (i.e. AWS.DocumentClient) to be instrumented
-      const dummyClient = {
-        service: {
-          customizeRequests: () => null,
-        }
-      };
 
       // Act
-      tracer.captureAWSClient(dummyClient);
+      const client = tracer.captureAWSClient(new DynamoDB.DocumentClient());
 
       // Assess
       expect(captureAWSClientSpy).toBeCalledTimes(2);
-      expect(captureAWSClientSpy).toHaveBeenNthCalledWith(1, dummyClient);
-      expect(captureAWSClientSpy).toHaveBeenNthCalledWith(2, dummyClient.service);
+      expect(captureAWSClientSpy).toHaveBeenNthCalledWith(1, client);
+      expect(captureAWSClientSpy).toHaveBeenNthCalledWith(2, (client as unknown as DynamoDB & { service: DynamoDB }).service);
+      expect(client).toBeInstanceOf(DynamoDB.DocumentClient);
 
     });