Added capture method

dreamorosi · dreamorosi · commit 3c449eafa575 · 2021-12-07T04:25:38.000+01:00
diff --git a/packages/tracing/README.md b/packages/tracing/README.md
@@ -98,7 +98,7 @@ new Lambda().handler(dummyEvent, dummyContext, () => console.log('Lambda invoked
                                     "ColdStart": true
                                 },
                                 "metadata": {
-                                    "my-service": {
+                                    "hello-world": {
                                         "foo-bar-function response": {
                                             "foo": "bar"
                                         }
@@ -217,7 +217,7 @@ new Lambda().handler(dummyEvent, dummyContext, () => console.log('Lambda invoked
 ```
 </details>
 
-### Add annotation on subsegment
+### Adding annotation on subsegment
 
 ```typescript
 // Environment variables set for the Lambda
@@ -290,7 +290,7 @@ new Lambda().handler(dummyEvent, dummyContext, () => console.log('Lambda invoked
                                     "my-annotation": "my-value"
                                 },
                                 "metadata": {
-                                    "my-service": {
+                                    "hello-world": {
                                         "foo-bar-function response": {
                                             "foo": "bar"
                                         }
@@ -309,7 +309,7 @@ new Lambda().handler(dummyEvent, dummyContext, () => console.log('Lambda invoked
 ```
 </details>
 
-### Add metadata to subsegment
+### Adding metadata to subsegment
 
 ```typescript
 // Environment variables set for the Lambda
@@ -382,7 +382,7 @@ new Lambda().handler(dummyEvent, dummyContext, () => console.log('Lambda invoked
                                     "ColdStart": true
                                 },
                                 "metadata": {
-                                    "my-service": {
+                                    "hello-world": {
                                         "foo-bar-function response": {
                                             "foo": "bar"
                                         },
@@ -405,7 +405,110 @@ new Lambda().handler(dummyEvent, dummyContext, () => console.log('Lambda invoked
 ```
 </details>
 
-### Capture AWS SDK clients
+### Capturing other methods
+
+```typescript
+// Environment variables set for the Lambda
+process.env.POWERTOOLS_SERVICE_NAME = 'hello-world';
+
+const tracer = new Tracer();
+
+class Lambda implements LambdaInterface {
+  @tracer.captureMethod()
+  public async dummyMethod(some: string): Promise<string> {
+      // Some async logic
+      return new Promise((resolve, _reject) => setTimeout(() => resolve(some), 3000));
+  }
+
+  public async handler<TEvent, TResult>(_event: TEvent, _context: Context, _callback: Callback<TResult>): Promise<TResult> {
+    const result = await this.dummyMethod('bar');
+    
+    return new Promise((resolve, _reject) => resolve({
+      foo: result
+    } as unknown as TResult));
+  }
+
+}
+
+new Lambda().handler(dummyEvent, dummyContext, () => console.log('Lambda invoked!'));
+
+```
+
+<details>
+ <summary>Click to expand and see the trace</summary>
+
+```json
+{
+    "Id": "abcdef123456abcdef123456abcdef123456",
+    "Duration": 0.656,
+    "LimitExceeded": false,
+    "Segments": [
+        {
+            "Id": "1234567890abcdef0",
+            "Document": {
+                "id": "1234567890abcdef0",
+                "name": "foo-bar-function",
+                "start_time": 1638792392.764036,
+                "trace_id": "abcdef123456abcdef123456abcdef123456",
+                "end_time": 1638792392.957155,
+                "parent_id": "abcdef01234567890",
+                "aws": {
+                    "account_id": "111122223333",
+                    "function_arn": "arn:aws:lambda:us-east-1:111122223333:function:foo-bar-function",
+                    "resource_names": [
+                        "foo-bar-function"
+                    ]
+                },
+                "origin": "AWS::Lambda::Function",
+                "subsegments": [
+                    // Initialization subsegment (if any)
+                    {
+                        "id": "4be0933d48d5b52f",
+                        "name": "Invocation",
+                        "start_time": 1638792392.7642102,
+                        "end_time": 1638792392.9384046,
+                        "aws": {
+                            "function_arn": "arn:aws:lambda:eu-west-1:111122223333:function:foo-bar-function"
+                        },
+                        "subsegments": [
+                            {
+                                "id": "aae0c94a16d66abd",
+                                "name": "## foo-bar-function",
+                                "start_time": 1638792392.766,
+                                "end_time": 1638792392.836,
+                                "annotations": {
+                                    "ColdStart": true
+                                },
+                                "subsegments": [
+                                    {
+                                        "id": "b1ac78c231577476",
+                                        "name": "### dummyMethod",
+                                        "start_time": 1638845552.777,
+                                        "end_time": 1638845553.459,
+                                        "metadata": {
+                                            "hello-world": {
+                                                "dummyMethod response": "bar"
+                                            }
+                                            // Other metadata (if any)
+                                        }
+                                        // Annotations (if any)
+                                        // Other subsegments (if any)
+                                    }
+                                ]
+                            }
+                        ]
+                    },
+                    // Overhead subsegment (if any)
+                ]
+            }
+        }
+    ]
+}
+
+```
+</details>
+
+### Capturing AWS SDK clients
 
 **AWS SDK JS v3**
 ```typescript
@@ -639,7 +742,7 @@ new Lambda().handler(dummyEvent, dummyContext, () => console.log('Lambda invoked
 ```
 </details>
 
-### Disable capture response body
+### Disabling capture response body
 
 ```typescript
 // Environment variables set for the Lambda
@@ -713,7 +816,7 @@ new Lambda().handler(dummyEvent, dummyContext, () => console.log('Lambda invoked
 ```
 </details>
 
-### Disable capture error
+### Disabling capture error
 
 ```typescript
 // Environment variables set for the Lambda
diff --git a/packages/tracing/src/Tracer.ts b/packages/tracing/src/Tracer.ts
@@ -44,19 +44,19 @@ class Tracer implements ClassThatTraces {
   }
 
   public captureLambdaHanlder(): HandlerMethodDecorator {
-    return (_target, _propertyKey, descriptor) => {
+    return (target, _propertyKey, descriptor) => {
       const originalMethod = descriptor.value;
 
-      descriptor.value = (event, context, callback) => {
+      descriptor.value = (event, context, callback): any => {
         if (this.tracingEnabled === false) {
-          return originalMethod?.apply(this, [ event, context, callback ]);
+          return originalMethod?.apply(target, [ event, context, callback ]);
         }
 
-        this.provider.captureAsyncFunc(`## ${context.functionName}`, async subsegment => {
+        return this.provider.captureAsyncFunc(`## ${context.functionName}`, async subsegment => {
           this.annotateColdStart();
           let result;
           try {
-            result = await originalMethod?.apply(this, [ event, context, callback ]);
+            result = await originalMethod?.apply(target, [ event, context, callback ]);
             this.addResponseAsMetadata(result, context.functionName);
           } catch (error) {
             this.addErrorAsMetadata(error as Error);
@@ -69,19 +69,38 @@ class Tracer implements ClassThatTraces {
           return result;
         });
       };
+
+      return descriptor;
     };
   }
 
-  // TODO: Finish implementation, type definition is wrong & it doesn't work. Need help.
   public captureMethod(): MethodDecorator {
-    return (_target, _propertyKey, descriptor) => {
+    return (target, _propertyKey, descriptor) => {
       const originalMethod = descriptor.value;
-      console.debug(originalMethod);
-      /* descriptor.value = () => {
+      
+      descriptor.value = (...args: unknown[]) => {
         if (this.tracingEnabled === false) {
-          return originalMethod?.apply(this, [ ]);
+          return originalMethod?.apply(target, [...args]);
         }
-      }; */
+
+        return this.provider.captureAsyncFunc(`### ${originalMethod.name}`, async subsegment => {
+          let result;
+          try {
+            result = await originalMethod?.apply(this, [...args]);
+            this.addResponseAsMetadata(result, originalMethod.name);
+          } catch (error) {
+            this.addErrorAsMetadata(error as Error);
+            // TODO: should this error be thrown?? If thrown we get a ERR_UNHANDLED_REJECTION. If not aren't we are basically catching a Customer error?
+            // throw error;
+          } finally {
+            subsegment?.close();
+          }
+          
+          return result;
+        });
+      };
+
+      return descriptor;
     };
   }
   
diff --git a/packages/tracing/tests/unit/Tracer.test.ts b/packages/tracing/tests/unit/Tracer.test.ts
@@ -551,30 +551,114 @@ describe('Class: Tracer', () => {
 
   describe('Method: captureMethod', () => {
 
-    test('when called while tracing is disabled, it does nothing', () => {
+    test('when called while tracing is disabled, it does nothing', async () => {
 
       // Prepare
       const tracer: Tracer = new Tracer({ enabled: false });
+      const captureAsyncFuncSpy = jest.spyOn(tracer.provider, 'captureAsyncFunc');
       class Lambda implements LambdaInterface {
 
         @tracer.captureMethod()
         // eslint-disable-next-line @typescript-eslint/ban-ts-comment
         // @ts-ignore
-        public dummyMethod(): boolean {
-          return true;
+        public async dummyMethod(some: string): Promise<any> {
+          return new Promise((resolve, _reject) => resolve(some));
         }
 
-        public handler<TEvent, TResult>(_event: TEvent, _context: Context, _callback: Callback<TResult>): void | Promise<TResult> {
-          return new Promise((resolve, _reject) => resolve({} as unknown as TResult));
+        public async handler<TEvent, TResult>(_event: TEvent, _context: Context, _callback: Callback<TResult>): Promise<TResult> {
+          const result = await this.dummyMethod('foo bar');
+          return new Promise((resolve, _reject) => resolve(result as unknown as TResult));
+        }
+
+      }
+
+      // Act
+      await new Lambda().handler(dummyEvent, dummyContext, () => console.log('Lambda invoked!'));
+
+      // Assess
+      expect(captureAsyncFuncSpy).toBeCalledTimes(0);
+
+    });
+
+    test('when used as decorator and with standard config, it captures the response as metadata', async () => {
+
+      // Prepare
+      const tracer: Tracer = new Tracer();
+      const newSubsegment: Segment | Subsegment | undefined = new Subsegment('### dummyMethod');
+      jest.spyOn(tracer.provider, 'getSegment')
+        .mockImplementation(() => newSubsegment);
+      setContextMissingStrategy(() => null);
+      const captureAsyncFuncSpy = jest.spyOn(tracer.provider, 'captureAsyncFunc');
+      class Lambda implements LambdaInterface {
+
+        @tracer.captureMethod()
+        // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+        // @ts-ignore
+        public async dummyMethod(some: string): Promise<string> {
+          return new Promise((resolve, _reject) => setTimeout(() => resolve(some), 3000));
+        }
+
+        public async handler<TEvent, TResult>(_event: TEvent, _context: Context, _callback: Callback<TResult>): Promise<TResult> {
+          const result = await this.dummyMethod('foo bar');
+          return new Promise((resolve, _reject) => resolve(result as unknown as TResult));
         }
 
       }
 
       // Act
-      new Lambda().dummyMethod();
+      await new Lambda().handler(dummyEvent, dummyContext, () => console.log('Lambda invoked!'));
 
       // Assess
-      expect(console.debug).toBeCalledTimes(1);
+      expect(captureAsyncFuncSpy).toHaveBeenCalledTimes(1);
+      expect(captureAsyncFuncSpy).toHaveBeenCalledWith('### dummyMethod', expect.anything());
+      expect(newSubsegment).toEqual(expect.objectContaining({
+        name: '### dummyMethod',
+        metadata: {
+          'hello-world': {
+            'dummyMethod response': 'foo bar',
+          },
+        }
+      }));
+
+    });
+
+    test('when used as decorator and with standard config, it captures the exception correctly', async () => {
+
+      // Prepare
+      const tracer: Tracer = new Tracer();
+      const newSubsegment: Segment | Subsegment | undefined = new Subsegment('### dummyMethod');
+      jest.spyOn(tracer.provider, 'getSegment')
+        .mockImplementation(() => newSubsegment);
+      setContextMissingStrategy(() => null);
+      const captureAsyncFuncSpy = jest.spyOn(tracer.provider, 'captureAsyncFunc');
+      const addErrorSpy = jest.spyOn(newSubsegment, 'addError');
+      class Lambda implements LambdaInterface {
+
+        @tracer.captureMethod()
+        // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+        // @ts-ignore
+        public async dummyMethod(_some: string): Promise<string> {
+          throw new Error('Exception thrown!');
+        }
+
+        public async handler<TEvent, TResult>(_event: TEvent, _context: Context, _callback: Callback<TResult>): Promise<TResult> {
+          const result = await this.dummyMethod('foo bar');
+          return new Promise((resolve, _reject) => resolve(result as unknown as TResult));
+        }
+
+      }
+
+      // Act
+      await new Lambda().handler(dummyEvent, dummyContext, () => console.log('Lambda invoked!'));
+
+      // Assess
+      expect(captureAsyncFuncSpy).toHaveBeenCalledTimes(1);
+      expect(newSubsegment).toEqual(expect.objectContaining({
+        name: '### dummyMethod',
+      }));
+      expect('cause' in newSubsegment).toBe(true);
+      expect(addErrorSpy).toHaveBeenCalledTimes(1);
+      expect(addErrorSpy).toHaveBeenCalledWith(new Error('Exception thrown!'), false);
 
     });
 
diff --git a/packages/tracing/types/Tracer.ts b/packages/tracing/types/Tracer.ts
@@ -21,10 +21,10 @@ type TracerOptions = {
   customConfigService?: ConfigServiceInterface
 };
 
+// TODO: Revisit type below, it doesn't allow to define async handlers.
 type HandlerMethodDecorator = (target: LambdaInterface, propertyKey: string | symbol, descriptor: TypedPropertyDescriptor<Handler>) => TypedPropertyDescriptor<Handler> | void;
 
-// TODO: Revisit these types that don't work.
-type MethodDecorator = <T>(target: LambdaInterface, propertyKey: string | symbol, descriptor: TypedPropertyDescriptor<T>) => TypedPropertyDescriptor<T> | void;
+type MethodDecorator = (target: any, propertyKey: string | symbol, descriptor: TypedPropertyDescriptor<any>) => any;
 
 export {
   ClassThatTraces,
