feat(tracer): middy middleware (#324)

* Implemented `captureLambdaHandler` middleware

* Documentation

Author: dreamorosi

docs/core/tracer.md

@@ -42,17 +42,16 @@ You can quickly start by importing the `Tracer` class, initialize it outside the
 
 === "Middleware"
 
-    ```typescript hl_lines="1 3 6"
+    ```typescript hl_lines="1-2 4 7 9"
     import { Tracer } from '@aws-lambda-powertools/tracer';
+    import middy from '@middy/core';
 
     const tracer = Tracer(); // Sets service via env var
     // OR tracer = Tracer({ service: 'example' });
 
-    // TODO: update example once middleware has been implemented.
-
-    export const handler = async (_event: any, _context: any) => {
+    export const handler = middy(async (_event: any, _context: any) => {
         ...
-    }
+    }).use(captureLambdaHandler(tracer));
     ```
 
 === "Decorator"
@@ -76,7 +75,7 @@ You can quickly start by importing the `Tracer` class, initialize it outside the
 
 === "Manual"
 
-    ```typescript hl_lines="1-2 4 8-9 11 17 20 24"
+    ```typescript hl_lines="1-2 4 9-10 12 18 21 25"
     import { Tracer } from '@aws-lambda-powertools/tracer';
     import { Segment } from 'aws-xray-sdk-core';
 
@@ -107,8 +106,7 @@ You can quickly start by importing the `Tracer` class, initialize it outside the
     }
     ```
 
-<!-- TODO: Replace name of middleware once implemented -->
-When using thes `captureLambdaHanlder` decorator or the `TBD` middleware, Tracer performs these additional tasks to ease operations:
+When using the `captureLambdaHandler` decorator or middleware, Tracer performs these additional tasks to ease operations:
 
 * Handles the lifecycle of the subsegment
 * Creates a `ColdStart` annotation to easily filter traces that have had an initialization overhead
@@ -148,23 +146,10 @@ When using thes `captureLambdaHanlder` decorator or the `TBD` middleware, Tracer
 
 ### Methods
 
-You can trace other methods using the `captureMethod` decorator.
-
-=== "Middleware"
-
-    ```typescript hl_lines="1 3 6"
-    import { Tracer } from '@aws-lambda-powertools/tracer';
-
-    const tracer = Tracer();
+You can trace other methods using the `captureMethod` decorator or manual instrumentation.
 
-    // TODO: update example once middleware has been implemented.
-
-
-
-    export const handler = async (_event: any, _context: any) => {
-        ...
-    }
-    ```
+!!! info
+    We currently support a middleware for tracing methods, [let us know](https://github.com/awslabs/aws-lambda-powertools-typescript/issues/new?assignees=&labels=feature-request%2C+triage&template=feature_request.md&title=) if you'd like to see one!
 
 === "Decorator"
 

packages/tracing/npm-shrinkwrap.json

@@ -54,6 +54,7 @@
     "url": "https://github.com/awslabs/aws-lambda-powertools-typescript/issues"
   },
   "dependencies": {
+    "@middy/core": "^2.5.3",
     "@aws-lambda-powertools/commons": "^0.0.2",
     "aws-xray-sdk-core": "^3.3.3"
   }

packages/tracing/package.json

@@ -20,11 +20,28 @@ import { Segment, Subsegment } from 'aws-xray-sdk-core';
  * ## Usage
  * 
  * ### Functions usage with middlewares
- * TBD
+ * 
+ * If you use function-based Lambda handlers you can use the [captureLambdaHanlder()](./_aws_lambda_powertools_tracer.Tracer.html) middy middleware to automatically:
+ * * handle the subsegment lifecycle 
+ * * add the `ColdStart` annotation
+ * * add the function response as metadata
+ * * add the function error as metadata (if any)
+ * 
+ * @example
+ * ```typescript
+ * import { Tracer, captureLambdaHandler } from '@aws-lambda-powertools/tracer';
+ * import middy from '@middy/core';
+ * 
+ * const tracer = new Tracer({ serviceName: 'my-service' });
+ * 
+ * export const handler = middy(async (_event: any, _context: any) => {
+ *   ...
+ * }).use(captureLambdaHandler(tracer));
+ * ```
  * 
  * ### Object oriented usage with decorators
  * 
- * If you use TypeScript Classes to wrap your Lambda handler you can use the [@tracer.captureLambdaHanlder()](./_aws_lambda_powertools_tracer.Tracer.html#captureLambdaHanlder) decorator to automatically:
+ * If instead you use TypeScript Classes to wrap your Lambda handler you can use the [@tracer.captureLambdaHanlder()](./_aws_lambda_powertools_tracer.Tracer.html#captureLambdaHanlder) decorator to automatically:
  * * handle the subsegment lifecycle 
  * * add the `ColdStart` annotation
  * * add the function response as metadata
@@ -65,7 +82,7 @@ import { Segment, Subsegment } from 'aws-xray-sdk-core';
  *   const subsegment = new Subsegment(`## ${context.functionName}`);
  *   tracer.setSegment(subsegment);
  *   // Add the ColdStart annotation
- *   this.putAnnotation('ColdStart', tracer.coldStart);
+ *   this.putAnnotation('ColdStart', tracer.isColdStart());
  * 
  *   let res;
  *   try {
@@ -242,8 +259,7 @@ class Tracer implements TracerInterface {
             this.addResponseAsMetadata(result, context.functionName);
           } 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;
+            throw error;
           } finally {
             subsegment?.close();
           }
@@ -351,6 +367,30 @@ class Tracer implements TracerInterface {
 
     return segment;
   }
+  
+  /**
+   * Get the current value of the `captureError` property.
+   * 
+   * You can use this method during manual instrumentation to determine
+   * if tracer should be capturing errors.
+   * 
+   * @returns captureError - `true` if errors should be captured, `false` otherwise. 
+   */
+  public isCaptureErrorEnabled(): boolean {
+    return this.captureError;
+  }
+
+  /**
+   * Get the current value of the `captureResponse` property.
+   * 
+   * You can use this method during manual instrumentation to determine
+   * if tracer should be capturing function responses.
+   * 
+   * @returns captureResponse - `true` if responses should be captured, `false` otherwise. 
+   */
+  public isCaptureResponseEnabled(): boolean {
+    return this.captureResponse;
+  }
 
   /**
    * Retrieve the current value of `ColdStart`.
@@ -361,7 +401,7 @@ class Tracer implements TracerInterface {
    * 
    * @see https://docs.aws.amazon.com/lambda/latest/dg/runtimes-context.html
    * 
-   * @returns boolean - true if is cold start otherwise false
+   * @returns boolean - `true` if is cold start, otherwise `false`
    */
   public static isColdStart(): boolean {
     if (Tracer.coldStart === true) {
@@ -373,6 +413,18 @@ class Tracer implements TracerInterface {
     return false;
   }
 
+  /**
+   * Get the current value of the `tracingEnabled` property.
+   * 
+   * You can use this method during manual instrumentation to determine
+   * if tracer is currently enabled.
+   * 
+   * @returns tracingEnabled - `true` if tracing is enabled, `false` otherwise. 
+   */
+  public isTracingEnabled(): boolean {
+    return this.tracingEnabled;
+  }
+
   /**
    * Adds annotation to existing segment or subsegment.
    * 

packages/tracing/src/Tracer.ts

@@ -1,3 +1,4 @@
 export * from './helpers';
 export * from './Tracer';
-export * from './TracerInterface';
+export * from './TracerInterface';
+export * from './middleware/middy';

packages/tracing/src/index.ts

@@ -0,0 +1,76 @@
+import middy from '@middy/core';
+import { Subsegment } from 'aws-xray-sdk-core';
+import { Tracer } from '../Tracer';
+
+/**
+ * A middy middleware automating capture of metadata and annotations on segments or subsegments ofr a Lambda Handler.
+ * 
+ * Using this middleware on your handler function will automatically:
+ * * handle the subsegment lifecycle 
+ * * add the `ColdStart` annotation
+ * * add the function response as metadata
+ * * add the function error as metadata (if any)
+ * 
+ * @example
+ * ```typescript
+ * import { Tracer, captureLambdaHandler } from '@aws-lambda-powertools/tracer';
+ * import middy from '@middy/core';
+ * 
+ * const tracer = new Tracer({ serviceName: 'my-service' });
+ * 
+ * export const handler = middy(async (_event: any, _context: any) => {
+ *   ...
+ * }).use(captureLambdaHandler(tracer));
+ * ```
+ * 
+ * @param tracer - The Tracer instance to use for tracing
+ * @returns middleware object - The middy middleware object
+ */
+const captureLambdaHandler = (target: Tracer): middy.MiddlewareObj => {
+  const captureLambdaHandlerBefore = async (request: middy.Request): Promise<void> => {
+    if (target.isTracingEnabled()) {
+      const subsegment = new Subsegment(`## ${request.context.functionName}`);
+      target.setSegment(subsegment);
+  
+      if (Tracer.isColdStart()) {
+        target.putAnnotation('ColdStart', true);
+      }
+    }
+  };
+
+  const captureLambdaHandlerAfter = async (request: middy.Request): Promise<void> => {
+    if (target.isTracingEnabled()) {
+      const subsegment = target.getSegment();
+      if (request.response !== undefined && target.isCaptureResponseEnabled() === true) {
+        target.putMetadata(`${request.context.functionName} response`, request.response);
+      }
+      
+      subsegment?.close();
+    }
+  };
+
+  const captureLambdaHandlerError = async (request: middy.Request): Promise<void> => {
+    if (target.isTracingEnabled()) {
+      const subsegment = target.getSegment();
+      if (target.isCaptureErrorEnabled() === false) {  
+        subsegment?.addErrorFlag();
+      } else {
+        subsegment?.addError(request.error as Error, false);
+      }
+      // TODO: should this error be thrown?? I.e. should we stop the event flow & return?
+      // throw request.error;
+
+      subsegment?.close();
+    }
+  };
+
+  return {
+    before: captureLambdaHandlerBefore,
+    after: captureLambdaHandlerAfter,
+    onError: captureLambdaHandlerError
+  };
+};
+
+export { 
+  captureLambdaHandler,
+};

packages/tracing/src/middleware/middy.ts

@@ -451,19 +451,19 @@ describe('Class: Tracer', () => {
 
       }
 
-      // Act
-      await new Lambda().handler(dummyEvent, dummyContext, () => console.log('Lambda invoked!'));
-
-      // Assess
+      // Act & Assess
+      await expect(new Lambda().handler({}, dummyContext, () => console.log('Lambda invoked!'))).rejects.toThrowError(Error);
       expect(captureAsyncFuncSpy).toHaveBeenCalledTimes(1);
       expect(newSubsegment).toEqual(expect.objectContaining({
         name: '## foo-bar-function',
       }));
       expect('cause' in newSubsegment).toBe(false);
       expect(addErrorFlagSpy).toHaveBeenCalledTimes(1);
       expect(addErrorSpy).toHaveBeenCalledTimes(0);
+      expect.assertions(6);
 
       delete process.env.POWERTOOLS_TRACER_CAPTURE_ERROR;
+
     });
 
     test('when used as decorator and with standard config, it captures the exception correctly', async () => {
@@ -487,17 +487,16 @@ describe('Class: Tracer', () => {
 
       }
 
-      // Act
-      await new Lambda().handler(dummyEvent, dummyContext, () => console.log('Lambda invoked!'));
-
-      // Assess
+      // Act & Assess
+      await expect(new Lambda().handler({}, dummyContext, () => console.log('Lambda invoked!'))).rejects.toThrowError(Error);
       expect(captureAsyncFuncSpy).toHaveBeenCalledTimes(1);
       expect(newSubsegment).toEqual(expect.objectContaining({
         name: '## foo-bar-function',
       }));
       expect('cause' in newSubsegment).toBe(true);
       expect(addErrorSpy).toHaveBeenCalledTimes(1);
       expect(addErrorSpy).toHaveBeenCalledWith(new Error('Exception thrown!'), false);
+      expect.assertions(6);
 
     });