feat(tracer): allow disabling result capture for decorators and middleware (#1065)

* feat(tracer): use decorator options to disable exception and result capture

* docs: small inline docs improvements

* chore: remove the captureError option

* chore: remove all captureLambdaHandler decorator changes

* Revert "chore: remove all captureLambdaHandler decorator changes"

This reverts commit ae72d52.

* feat: add middy middleware options

* test: add middy options e2e

* test: add decorator e2e

* refactor: initial review changes

* refactor: apply import change from review

Co-authored-by: Andrea Amorosi <[email protected]>

Co-authored-by: Andrea Amorosi <[email protected]>

Author: misterjoshua

Diff for: docs/core/tracer.md

@@ -256,7 +256,7 @@ You can trace other Class methods using the `captureMethod` decorator or any arb
     }
      
     const handlerClass = new Lambda();
-    export const handler = myFunction.handler.bind(handlerClass); // (1)
+    export const handler = handlerClass.handler.bind(handlerClass); // (1)
     ```
 
     1. Binding your handler method allows your handler to access `this`.
@@ -412,6 +412,69 @@ Use **`POWERTOOLS_TRACER_CAPTURE_RESPONSE=false`** environment variable to instr
     2. You might manipulate **streaming objects that can be read only once**; this prevents subsequent calls from being empty
     3. You might return **more than 64K** of data _e.g., `message too long` error_
 
+Alternatively, use the `captureResponse: false` option in both `tracer.captureLambdaHandler()` and `tracer.captureMethod()` decorators, or use the same option in the Middy `captureLambdaHander` middleware to instruct Tracer **not** to serialize function responses as metadata.
+
+=== "method.ts"
+
+    ```typescript hl_lines="6"
+    import { Tracer } from '@aws-lambda-powertools/tracer';
+
+    const tracer = new Tracer({ serviceName: 'serverlessAirline' });
+
+    class Lambda implements LambdaInterface {
+        @tracer.captureMethod({ captureResult: false })
+        public getChargeId(): string {
+            /* ... */
+            return 'foo bar';
+        }
+
+        public async handler(_event: any, _context: any): Promise<void> {
+            /* ... */
+        }
+    }
+
+    const handlerClass = new Lambda();
+    export const handler = handlerClass.handler.bind(handlerClass);
+    ```
+
+=== "handler.ts"
+
+    ```typescript hl_lines="7"
+    import { Tracer } from '@aws-lambda-powertools/tracer';
+    import { LambdaInterface } from '@aws-lambda-powertools/commons';
+
+    const tracer = new Tracer({ serviceName: 'serverlessAirline' });
+
+    class Lambda implements LambdaInterface {
+        @tracer.captureLambdaHandler({ captureResponse: false })
+        async handler(_event: any, _context: any): Promise<void> {
+            /* ... */
+        }
+    }
+
+    const handlerClass = new Lambda();
+    export const handler = handlerClass.handler.bind(handlerClass);
+    ```
+
+=== "middy.ts"
+
+    ```typescript hl_lines="14"
+    import { Tracer, captureLambdaHandler } from '@aws-lambda-powertools/tracer';
+    import middy from '@middy/core';
+
+    const tracer = new Tracer({ serviceName: 'serverlessAirline' });
+
+    const lambdaHandler = async (_event: any, _context: any): Promise<void> => {
+        /* ... */
+    };
+
+    // Wrap the handler with middy
+    export const handler = middy(lambdaHandler)
+        // Use the middleware by passing the Tracer instance as a parameter,
+        // but specify the captureResponse option as false.
+        .use(captureLambdaHandler(tracer, { captureResponse: false }));
+    ```
+
 ### Disabling exception auto-capture
 
 Use **`POWERTOOLS_TRACER_CAPTURE_ERROR=false`** environment variable to instruct Tracer **not** to serialize exceptions as metadata.

Diff for: packages/tracer/src/Tracer.ts

@@ -2,7 +2,7 @@ import { Handler } from 'aws-lambda';
 import { AsyncHandler, SyncHandler, Utility } from '@aws-lambda-powertools/commons';
 import { TracerInterface } from '.';
 import { ConfigServiceInterface, EnvironmentVariablesService } from './config';
-import { HandlerMethodDecorator, TracerOptions, MethodDecorator } from './types';
+import { HandlerMethodDecorator, TracerOptions, HandlerOptions, MethodDecorator } from './types';
 import { ProviderService, ProviderServiceInterface } from './provider';
 import { Segment, Subsegment } from 'aws-xray-sdk-core';
 
@@ -339,7 +339,7 @@ class Tracer extends Utility implements TracerInterface {
    * 
    * @decorator Class
    */
-  public captureLambdaHandler(): HandlerMethodDecorator {
+  public captureLambdaHandler(options?: HandlerOptions): HandlerMethodDecorator {
     return (_target, _propertyKey, descriptor) => {
       /**
        * The descriptor.value is the method this decorator decorates, it cannot be undefined.
@@ -365,7 +365,10 @@ class Tracer extends Utility implements TracerInterface {
           let result: unknown;
           try {
             result = await originalMethod.apply(handlerRef, [ event, context, callback ]);
-            tracerRef.addResponseAsMetadata(result, process.env._HANDLER);
+            if (options?.captureResponse ?? true) {
+              tracerRef.addResponseAsMetadata(result, process.env._HANDLER);
+            }
+
           } catch (error) {
             tracerRef.addErrorAsMetadata(error as Error);
             throw error;
@@ -416,7 +419,7 @@ class Tracer extends Utility implements TracerInterface {
    * 
    * @decorator Class
    */
-  public captureMethod(): MethodDecorator {
+  public captureMethod(options?: HandlerOptions): MethodDecorator {
     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
@@ -435,7 +438,9 @@ class Tracer extends Utility implements TracerInterface {
           let result;
           try {
             result = await originalMethod.apply(this, [...args]);
-            tracerRef.addResponseAsMetadata(result, originalMethod.name);
+            if (options?.captureResponse ?? true) {
+              tracerRef.addResponseAsMetadata(result, originalMethod.name);
+            }
           } catch (error) {
             tracerRef.addErrorAsMetadata(error as Error);
 

Diff for: packages/tracer/src/middleware/middy.ts

@@ -1,9 +1,10 @@
 import type middy from '@middy/core';
 import type { Tracer } from '../Tracer';
 import type { Segment, Subsegment } from 'aws-xray-sdk-core';
+import type { HandlerOptions } from '../types';
 
 /**
- * A middy middleware automating capture of metadata and annotations on segments or subsegments ofr a Lambda Handler.
+ * A middy middleware automating capture of metadata and annotations on segments or subsegments for a Lambda Handler.
  * 
  * Using this middleware on your handler function will automatically:
  * * handle the subsegment lifecycle 
@@ -26,7 +27,7 @@ import type { Segment, Subsegment } from 'aws-xray-sdk-core';
  * @param target - The Tracer instance to use for tracing
  * @returns middleware object - The middy middleware object
  */
-const captureLambdaHandler = (target: Tracer): middy.MiddlewareObj => {
+const captureLambdaHandler = (target: Tracer, options?: HandlerOptions): middy.MiddlewareObj => {
   let lambdaSegment: Subsegment | Segment;
 
   const open = (): void => {
@@ -51,7 +52,9 @@ const captureLambdaHandler = (target: Tracer): middy.MiddlewareObj => {
 
   const captureLambdaHandlerAfter = async (request: middy.Request): Promise<void> => {
     if (target.isTracingEnabled()) {
-      target.addResponseAsMetadata(request.response, process.env._HANDLER);
+      if (options?.captureResponse ?? true) {
+        target.addResponseAsMetadata(request.response, process.env._HANDLER);
+      }
       close();
     }
   };

Diff for: packages/tracer/src/types/Tracer.ts

@@ -26,6 +26,27 @@ type TracerOptions = {
   customConfigService?: ConfigServiceInterface
 };
 
+/**
+ * Options for handler decorators and middleware.
+ *
+ * Usage:
+ * @example
+ * ```typescript
+ * const tracer = new Tracer();
+ *
+ * class Lambda implements LambdaInterface {
+ *   @tracer.captureLambdaHandler({ captureResponse: false })
+ *   async handler(_event: any, _context: any): Promise<void> {}
+ * }
+ *
+ * const handlerClass = new Lambda();
+ * export const handler = handlerClass.handler.bind(handlerClass);
+ * ```
+ */
+type HandlerOptions = {
+  captureResponse?: boolean
+};
+
 type HandlerMethodDecorator = (
   target: LambdaInterface,
   propertyKey: string | symbol,
@@ -38,6 +59,7 @@ type MethodDecorator = (target: any, propertyKey: string | symbol, descriptor: T
 
 export {
   TracerOptions,
+  HandlerOptions,
   HandlerMethodDecorator,
   MethodDecorator
 };

Diff for: packages/tracer/tests/e2e/allFeatures.decorator.test.functionCode.ts

@@ -35,14 +35,13 @@ const refreshAWSSDKImport = (): void => {
 const tracer = new Tracer({ serviceName: serviceName });
 const dynamoDBv3 = tracer.captureAWSv3Client(new DynamoDBClient({}));
 
-export class MyFunctionWithDecorator {  
+export class MyFunctionBase {
   private readonly returnValue: string;
 
   public constructor() {
     this.returnValue = customResponseValue;
   }
 
-  @tracer.captureLambdaHandler()
   // eslint-disable-next-line @typescript-eslint/ban-ts-comment
   // @ts-ignore
   public handler(event: CustomEvent, _context: Context, _callback: Callback<unknown>): void | Promise<unknown> {
@@ -79,13 +78,47 @@ export class MyFunctionWithDecorator {
       });
   }
 
-  @tracer.captureMethod()
   // eslint-disable-next-line @typescript-eslint/ban-ts-comment
   // @ts-ignore
   public myMethod(): string {
     return this.returnValue;
   }
 }
 
+class MyFunctionWithDecorator extends MyFunctionBase {
+  @tracer.captureLambdaHandler()
+  // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+  // @ts-ignore
+  public handler(event: CustomEvent, _context: Context, _callback: Callback<unknown>): void | Promise<unknown> {
+    return super.handler(event, _context, _callback);
+  }
+
+  @tracer.captureMethod()
+  // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+  // @ts-ignore
+  public myMethod(): string {
+    return super.myMethod();
+  }
+}
+
 const handlerClass = new MyFunctionWithDecorator();
-export const handler = handlerClass.handler.bind(handlerClass);
+export const handler = handlerClass.handler.bind(handlerClass);
+
+class MyFunctionWithDecoratorCaptureResponseFalse extends MyFunctionBase {
+  @tracer.captureLambdaHandler({ captureResponse: false })
+  // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+  // @ts-ignore
+  public handler(event: CustomEvent, _context: Context, _callback: Callback<unknown>): void | Promise<unknown> {
+    return super.handler(event, _context, _callback);
+  }
+
+  @tracer.captureMethod({ captureResponse: false })
+  // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+  // @ts-ignore
+  public myMethod(): string {
+    return super.myMethod();
+  }
+}
+
+const handlerWithCaptureResponseFalseClass = new MyFunctionWithDecoratorCaptureResponseFalse();
+export const handlerWithCaptureResponseFalse = handlerClass.handler.bind(handlerWithCaptureResponseFalseClass);

Diff for: packages/tracer/tests/e2e/allFeatures.decorator.test.ts

@@ -78,6 +78,13 @@ const uuidFunction3 = v4();
 const functionNameWithTracerDisabled = generateUniqueName(RESOURCE_NAME_PREFIX, uuidFunction3, runtime, 'AllFeatures-Decorator-TracerDisabled');
 const serviceNameWithTracerDisabled = functionNameWithNoCaptureErrorOrResponse; 
 
+/**
+ * Function #4 disables tracer
+ */
+const uuidFunction4 = v4();
+const functionNameWithCaptureResponseFalse = generateUniqueName(RESOURCE_NAME_PREFIX, uuidFunction4, runtime, 'AllFeatures-Decorator-CaptureResponseFalse');
+const serviceNameWithCaptureResponseFalse = functionNameWithCaptureResponseFalse;
+
 const xray = new AWS.XRay();
 const invocations = 3;
 
@@ -149,13 +156,30 @@ describe(`Tracer E2E tests, all features with decorator instantiation for runtim
     });
     ddbTable.grantWriteData(functionWithTracerDisabled);
 
+    const functionWithCaptureResponseFalse = createTracerTestFunction({
+      stack,
+      functionName: functionNameWithCaptureResponseFalse,
+      handler: 'handlerWithCaptureResponseFalse',
+      entry,
+      expectedServiceName: serviceNameWithCaptureResponseFalse,
+      environmentParams: {
+        TEST_TABLE_NAME: ddbTableName,
+        POWERTOOLS_TRACER_CAPTURE_RESPONSE: 'true',
+        POWERTOOLS_TRACER_CAPTURE_ERROR: 'true',
+        POWERTOOLS_TRACE_ENABLED: 'true',
+      },
+      runtime
+    });
+    ddbTable.grantWriteData(functionWithCaptureResponseFalse);
+
     await deployStack(integTestApp, stack);
 
     // Act
     await Promise.all([
       invokeAllTestCases(functionNameWithAllFlagsEnabled),
       invokeAllTestCases(functionNameWithNoCaptureErrorOrResponse),
       invokeAllTestCases(functionNameWithTracerDisabled),
+      invokeAllTestCases(functionNameWithCaptureResponseFalse),
     ]);
 
   }, SETUP_TIMEOUT);
@@ -303,6 +327,62 @@ describe(`Tracer E2E tests, all features with decorator instantiation for runtim
 
   }, TEST_CASE_TIMEOUT);
 
+  it('should not capture response when the decorator\'s captureResponse is set to false', async () => {
+
+    const tracesWithCaptureResponseFalse = await getTraces(xray, startTime, await getFunctionArn(functionNameWithCaptureResponseFalse), invocations, 5);
+
+    expect(tracesWithCaptureResponseFalse.length).toBe(invocations);
+
+    // Assess
+    for (let i = 0; i < invocations; i++) {
+      const trace = tracesWithCaptureResponseFalse[i];
+
+      /**
+       * Expect the trace to have 5 segments:
+       * 1. Lambda Context (AWS::Lambda)
+       * 2. Lambda Function (AWS::Lambda::Function)
+       * 3. DynamoDB (AWS::DynamoDB)
+       * 4. DynamoDB Table (AWS::DynamoDB::Table)
+       * 5. Remote call (httpbin.org)
+       */
+      expect(trace.Segments.length).toBe(5);
+      const invocationSubsegment = getInvocationSubsegment(trace);
+
+      /**
+       * Invocation subsegment should have a subsegment '## index.handler' (default behavior for PowerTool tracer)
+       * '## index.handler' subsegment should have 4 subsegments
+       * 1. DynamoDB (PutItem on the table)
+       * 2. DynamoDB (PutItem overhead)
+       * 3. httpbin.org (Remote call)
+       * 4. '### myMethod' (method decorator)
+       */
+      const handlerSubsegment = getFirstSubsegment(invocationSubsegment);
+      expect(handlerSubsegment.name).toBe('## index.handlerWithCaptureResponseFalse');
+      expect(handlerSubsegment?.subsegments).toHaveLength(4);
+
+      if (!handlerSubsegment.subsegments) {
+        fail('"## index.handlerWithCaptureResponseFalse" subsegment should have subsegments');
+      }
+      const subsegments = splitSegmentsByName(handlerSubsegment.subsegments, [ 'DynamoDB', 'httpbin.org', '### myMethod' ]);
+      expect(subsegments.get('DynamoDB')?.length).toBe(2);
+      expect(subsegments.get('httpbin.org')?.length).toBe(1);
+      expect(subsegments.get('### myMethod')?.length).toBe(1);
+      expect(subsegments.get('other')?.length).toBe(0);
+
+      // No metadata because capturing the response was disabled and that's
+      // the only metadata that could be in the subsegment for the test.
+      const myMethodSegment = subsegments.get('### myMethod')?.[0];
+      expect(myMethodSegment).toBeDefined();
+      expect(myMethodSegment).not.toHaveProperty('metadata');
+
+      const shouldThrowAnError = (i === (invocations - 1));
+      if (shouldThrowAnError) {
+        assertErrorAndFault(invocationSubsegment, expectedCustomErrorMessage);
+      }
+    }
+
+  }, TEST_CASE_TIMEOUT);
+
   it('should not capture any custom traces when disabled', async () => {
     const expectedNoOfTraces = 2;
     const tracesWithTracerDisabled = await getTraces(xray, startTime, await getFunctionArn(functionNameWithTracerDisabled), invocations, expectedNoOfTraces);

Diff for: packages/tracer/tests/e2e/allFeatures.middy.test.functionCode.ts

@@ -36,7 +36,7 @@ const refreshAWSSDKImport = (): void => {
 const tracer = new Tracer({ serviceName: serviceName });
 const dynamoDBv3 = tracer.captureAWSv3Client(new DynamoDBClient({}));
 
-export const handler = middy(async (event: CustomEvent, _context: Context): Promise<void> => {
+const testHandler = async (event: CustomEvent, _context: Context): Promise<void> => {
   tracer.putAnnotation('invocation', event.invocation);
   tracer.putAnnotation(customAnnotationKey, customAnnotationValue);
   tracer.putMetadata(customMetadataKey, customMetadataValue);
@@ -63,4 +63,8 @@ export const handler = middy(async (event: CustomEvent, _context: Context): Prom
   } catch (err) {
     throw err;
   }
-}).use(captureLambdaHandler(tracer));
+};
+
+export const handler = middy(testHandler).use(captureLambdaHandler(tracer));
+
+export const handlerWithNoCaptureResponseViaMiddlewareOption = middy(testHandler).use(captureLambdaHandler(tracer, { captureResponse: false }));