fix(tracer): decorated class methods cannot access this (#1055)

* fix(tracer): use instance scope in capture decorators * docs: add more inline comments * docs: small grammatical inline docs correction * refactor: apply review changes * docs: .bind changes * test: use member variable in e2e tests * fix: bind the handler in the e2e tests
aws-powertools · Aug 16, 2022 · 107fa04 · 107fa04
1 parent 9ce180a
commit 107fa04
Show file tree

Hide file tree

Showing 5 changed files with 126 additions and 28 deletions.
diff --git a/docs/core/tracer.md b/docs/core/tracer.md
@@ -143,8 +143,11 @@ You can quickly start by importing the `Tracer` class, initialize it outside the
     }
      
     export const handlerClass = new Lambda();
-    export const handler = handlerClass.handler;
+    export const handler = handlerClass.handler.bind(handlerClass); // (1)
     ```
+
+    1. Binding your handler method allows your handler to access `this`.
+
 === "Manual"
 
     ```typescript hl_lines="6 8-9 12-13 19 22 26 28"
@@ -253,8 +256,11 @@ You can trace other Class methods using the `captureMethod` decorator or any arb
     }
      
     export const myFunction = new Lambda();
-    export const handler = myFunction.handler; 
+    export const handler = myFunction.handler.bind(myFunction); // (1)
     ```
+
+    1. Binding your handler method allows your handler to access `this`.
+
 === "Manual"
 
     ```typescript hl_lines="6 8-9 15 18 23 25"

diff --git a/packages/tracer/src/Tracer.ts b/packages/tracer/src/Tracer.ts
@@ -1,5 +1,5 @@
 import { Handler } from 'aws-lambda';
-import { Utility } from '@aws-lambda-powertools/commons';
+import { AsyncHandler, SyncHandler, Utility } from '@aws-lambda-powertools/commons';
 import { TracerInterface } from '.';
 import { ConfigServiceInterface, EnvironmentVariablesService } from './config';
 import { HandlerMethodDecorator, TracerOptions, MethodDecorator } from './types';
@@ -67,7 +67,7 @@ import { Segment, Subsegment } from 'aws-xray-sdk-core';
  * }
  * 
  * export const handlerClass = new Lambda();
- * export const handler = handlerClass.handler; 
+ * export const handler = handlerClass.handler.bind(handlerClass);
  * ```
  * 
  * ### Functions usage with manual instrumentation
@@ -334,33 +334,40 @@ class Tracer extends Utility implements TracerInterface {
    * }
    * 
    * export const handlerClass = new Lambda();
-   * export const handler = handlerClass.handler; 
+   * export const handler = handlerClass.handler.bind(handlerClass);
    * ```
    * 
    * @decorator Class
    */
   public captureLambdaHandler(): HandlerMethodDecorator {
-    return (target, _propertyKey, descriptor) => {
+    return (_target, _propertyKey, descriptor) => {
       /**
        * The descriptor.value is the method this decorator decorates, it cannot be undefined.
        */ 
       // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
       const originalMethod = descriptor.value!;
 
-      descriptor.value = ((event, context, callback) => {
-        if (!this.isTracingEnabled()) {
-          return originalMethod.apply(target, [ event, context, callback ]);
+      // eslint-disable-next-line @typescript-eslint/no-this-alias
+      const tracerRef = this;
+      // Use a function() {} instead of an () => {} arrow function so that we can
+      // access `myClass` as `this` in a decorated `myClass.myMethod()`.
+      descriptor.value = (function (this: Handler, event, context, callback) {
+        // eslint-disable-next-line @typescript-eslint/no-this-alias
+        const handlerRef: Handler = this;
+
+        if (!tracerRef.isTracingEnabled()) {
+          return originalMethod.apply(handlerRef, [ event, context, callback ]);
         }
 
-        return this.provider.captureAsyncFunc(`## ${process.env._HANDLER}`, async subsegment => {
-          this.annotateColdStart();
-          this.addServiceNameAnnotation();
+        return tracerRef.provider.captureAsyncFunc(`## ${process.env._HANDLER}`, async subsegment => {
+          tracerRef.annotateColdStart();
+          tracerRef.addServiceNameAnnotation();
           let result: unknown;
           try {
-            result = await originalMethod.apply(target, [ event, context, callback ]);
-            this.addResponseAsMetadata(result, process.env._HANDLER);
+            result = await originalMethod.apply(handlerRef, [ event, context, callback ]);
+            tracerRef.addResponseAsMetadata(result, process.env._HANDLER);
           } catch (error) {
-            this.addErrorAsMetadata(error as Error);
+            tracerRef.addErrorAsMetadata(error as Error);
             throw error;
           } finally {
             subsegment?.close();
@@ -369,7 +376,7 @@ class Tracer extends Utility implements TracerInterface {
 
           return result;
         });
-      }) as Handler;
+      }) as SyncHandler<Handler> | AsyncHandler<Handler>;
 
       return descriptor;
     };
@@ -411,23 +418,27 @@ class Tracer extends Utility implements TracerInterface {
    * @decorator Class
    */
   public captureMethod(): MethodDecorator {
-    return (target, _propertyKey, descriptor) => {
+    return (_target, _propertyKey, descriptor) => {
       // The descriptor.value is the method this decorator decorates, it cannot be undefined.
       // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
       const originalMethod = descriptor.value!;
 
-      descriptor.value = (...args: unknown[]) => {
-        if (!this.isTracingEnabled()) {
-          return originalMethod.apply(target, [...args]);
+      // eslint-disable-next-line @typescript-eslint/no-this-alias
+      const tracerRef = this;
+      // Use a function() {} instead of an () => {} arrow function so that we can
+      // access `myClass` as `this` in a decorated `myClass.myMethod()`.
+      descriptor.value = function (...args: unknown[]) {
+        if (!tracerRef.isTracingEnabled()) {
+          return originalMethod.apply(this, [...args]);
         }
 
-        return this.provider.captureAsyncFunc(`### ${originalMethod.name}`, async subsegment => {          
+        return tracerRef.provider.captureAsyncFunc(`### ${originalMethod.name}`, async subsegment => {
           let result;
           try {
-            result = await originalMethod.apply(target, [...args]);
-            this.addResponseAsMetadata(result, originalMethod.name);
+            result = await originalMethod.apply(this, [...args]);
+            tracerRef.addResponseAsMetadata(result, originalMethod.name);
           } catch (error) {
-            this.addErrorAsMetadata(error as Error); 
+            tracerRef.addErrorAsMetadata(error as Error);
 
             throw error;
           } finally {

diff --git a/packages/tracer/tests/e2e/allFeatures.decorator.test.functionCode.ts b/packages/tracer/tests/e2e/allFeatures.decorator.test.functionCode.ts
@@ -36,6 +36,12 @@ const tracer = new Tracer({ serviceName: serviceName });
 const dynamoDBv3 = tracer.captureAWSv3Client(new DynamoDBClient({}));
 
 export class MyFunctionWithDecorator {  
+  private readonly returnValue: string;
+
+  public constructor() {
+    this.returnValue = customResponseValue;
+  }
+
   @tracer.captureLambdaHandler()
   // eslint-disable-next-line @typescript-eslint/ban-ts-comment
   // @ts-ignore
@@ -77,9 +83,9 @@ export class MyFunctionWithDecorator {
   // eslint-disable-next-line @typescript-eslint/ban-ts-comment
   // @ts-ignore
   public myMethod(): string {
-    return customResponseValue;
+    return this.returnValue;
   }
 }
 
 export const handlerClass = new MyFunctionWithDecorator();
-export const handler = handlerClass.handler;
+export const handler = handlerClass.handler.bind(handlerClass);
diff --git a/packages/tracer/tests/e2e/asyncHandler.decorator.test.functionCode.ts b/packages/tracer/tests/e2e/asyncHandler.decorator.test.functionCode.ts
@@ -36,6 +36,12 @@ const tracer = new Tracer({ serviceName: serviceName });
 const dynamoDBv3 = tracer.captureAWSv3Client(new DynamoDBClient({}));
 
 export class MyFunctionWithDecorator {  
+  private readonly returnValue: string;
+
+  public constructor() {
+    this.returnValue = customResponseValue;
+  }
+
   @tracer.captureLambdaHandler()
   // eslint-disable-next-line @typescript-eslint/ban-ts-comment
   // @ts-ignore
@@ -72,9 +78,9 @@ export class MyFunctionWithDecorator {
   // eslint-disable-next-line @typescript-eslint/ban-ts-comment
   // @ts-ignore
   public myMethod(): string {
-    return customResponseValue;
+    return this.returnValue;
   }
 }
 
 export const handlerClass = new MyFunctionWithDecorator();
-export const handler = handlerClass.handler;
+export const handler = handlerClass.handler.bind(handlerClass);
diff --git a/packages/tracer/tests/unit/Tracer.test.ts b/packages/tracer/tests/unit/Tracer.test.ts
@@ -854,6 +854,37 @@ describe('Class: Tracer', () => {
 
     });
 
+    test('when used as decorator and when calling the handler, it has access to member variables', async () => {
+
+      // Prepare
+      const tracer: Tracer = new Tracer();
+      const newSubsegment: Segment | Subsegment | undefined = new Subsegment('### dummyMethod');
+      jest.spyOn(tracer.provider, 'getSegment')
+        .mockImplementation(() => newSubsegment);
+      setContextMissingStrategy(() => null);
+
+      class Lambda implements LambdaInterface {
+        private readonly memberVariable: string;
+
+        public constructor(memberVariable: string) {
+          this.memberVariable = memberVariable;
+        }
+
+        @tracer.captureLambdaHandler()
+        // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+        // @ts-ignore
+        public async handler<TEvent, TResult>(_event: TEvent, _context: Context, _callback: Callback<TResult>): void | Promise<TResult> {
+          return <TResult>(`memberVariable:${this.memberVariable}` as unknown);
+        }
+
+      }
+
+      // Act / Assess
+      const lambda = new Lambda('someValue');
+      const handler = lambda.handler.bind(lambda);
+      expect(await handler({}, context, () => console.log('Lambda invoked!'))).toEqual('memberVariable:someValue');
+
+    });
   });
 
   describe('Method: captureMethod', () => {
@@ -1009,6 +1040,44 @@ describe('Class: Tracer', () => {
 
     });
 
+    test('when used as decorator and when calling a method in the class, it has access to member variables', async () => {
+
+      // Prepare
+      const tracer: Tracer = new Tracer();
+      const newSubsegment: Segment | Subsegment | undefined = new Subsegment('### dummyMethod');
+      jest.spyOn(tracer.provider, 'getSegment')
+        .mockImplementation(() => newSubsegment);
+      setContextMissingStrategy(() => null);
+
+      class Lambda implements LambdaInterface {
+        private readonly memberVariable: string;
+
+        public constructor(memberVariable: string) {
+          this.memberVariable = memberVariable;
+        }
+
+        @tracer.captureMethod()
+        // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+        // @ts-ignore
+        public async dummyMethod(): Promise<string> {
+          return `memberVariable:${this.memberVariable}`;
+        }
+
+        @tracer.captureLambdaHandler()
+        // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+        // @ts-ignore
+        public async handler<TEvent, TResult>(_event: TEvent, _context: Context, _callback: Callback<TResult>): void | Promise<TResult> {
+          return <TResult>(await this.dummyMethod() as unknown);
+        }
+
+      }
+
+      // Act / Assess
+      const lambda = new Lambda('someValue');
+      expect(await lambda.dummyMethod()).toEqual('memberVariable:someValue');
+
+    });
+
   });
 
   describe('Method: captureAWS', () => {