feat(tracer): specify subsegment name when capturing class method (#1092)

* feat: specify subsegment name when capturing class method

* chore: update key for e2e test cases

* chore: add tips to the docs about setting custom name

* Update docs/core/tracer.md

Co-authored-by: Josh Kellendonk <misterjoshua@users.noreply.github.com>

* Update packages/tracer/src/types/Tracer.ts

* fix merge conflicts

Co-authored-by: Josh Kellendonk <misterjoshua@users.noreply.github.com>

Author: dreamorosi

docs/core/tracer.md

@@ -244,7 +244,7 @@ You can trace other Class methods using the `captureMethod` decorator or any arb
 
     class Lambda implements LambdaInterface {
         // Decorate your class method
-        @tracer.captureMethod()
+        @tracer.captureMethod() // (1)
         public getChargeId(): string {
             /* ... */
             return 'foo bar';
@@ -256,10 +256,11 @@ You can trace other Class methods using the `captureMethod` decorator or any arb
     }
      
     const handlerClass = new Lambda();
-    export const handler = handlerClass.handler.bind(handlerClass); // (1)
+    export const handler = handlerClass.handler.bind(handlerClass); // (2)
     ```
 
-    1. Binding your handler method allows your handler to access `this`.
+    1. You can set a custom name for the subsegment by passing `subSegmentName` to the decorator, like: `@tracer.captureMethod({ subSegmentName: '### myCustomMethod' })`.
+    2. Binding your handler method allows your handler to access `this`.
 
 === "Manual"
 

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, HandlerOptions, MethodDecorator } from './types';
+import { HandlerMethodDecorator, TracerOptions, MethodDecorator, CaptureLambdaHandlerOptions, CaptureMethodOptions } from './types';
 import { ProviderService, ProviderServiceInterface } from './provider';
 import { Segment, Subsegment } from 'aws-xray-sdk-core';
 
@@ -338,8 +338,9 @@ class Tracer extends Utility implements TracerInterface {
    * ```
    * 
    * @decorator Class
+   * @param options - (_optional_) Options for the decorator
    */
-  public captureLambdaHandler(options?: HandlerOptions): HandlerMethodDecorator {
+  public captureLambdaHandler(options?: CaptureLambdaHandlerOptions): HandlerMethodDecorator {
     return (_target, _propertyKey, descriptor) => {
       /**
        * The descriptor.value is the method this decorator decorates, it cannot be undefined.
@@ -418,8 +419,9 @@ class Tracer extends Utility implements TracerInterface {
    * ```
    * 
    * @decorator Class
+   * @param options - (_optional_) Options for the decorator
    */
-  public captureMethod(options?: HandlerOptions): MethodDecorator {
+  public captureMethod(options?: CaptureMethodOptions): 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,8 +437,9 @@ class Tracer extends Utility implements TracerInterface {
         }
 
         const methodName = String(propertyKey);
+        const subsegmentName = options?.subSegmentName ? options.subSegmentName : `### ${methodName}`;
 
-        return tracerRef.provider.captureAsyncFunc(`### ${methodName}`, async subsegment => {
+        return tracerRef.provider.captureAsyncFunc(subsegmentName, async subsegment => {
           let result;
           try {
             result = await originalMethod.apply(this, [...args]);

packages/tracer/src/TracerInterface.ts

@@ -1,4 +1,4 @@
-import { HandlerMethodDecorator, MethodDecorator } from './types';
+import { CaptureLambdaHandlerOptions, CaptureMethodOptions, HandlerMethodDecorator, MethodDecorator } from './types';
 import { Segment, Subsegment } from 'aws-xray-sdk-core';
 
 interface TracerInterface {
@@ -9,8 +9,8 @@ interface TracerInterface {
   captureAWS<T>(aws: T): void | T
   captureAWSv3Client<T>(service: T): void | T
   captureAWSClient<T>(service: T): void | T
-  captureLambdaHandler(): HandlerMethodDecorator
-  captureMethod(): MethodDecorator
+  captureLambdaHandler(options?: CaptureLambdaHandlerOptions): HandlerMethodDecorator
+  captureMethod(options?: CaptureMethodOptions): MethodDecorator
   getSegment(): Segment | Subsegment
   isTracingEnabled(): boolean
   putAnnotation: (key: string, value: string | number | boolean) => void

packages/tracer/src/middleware/middy.ts

@@ -1,7 +1,7 @@
 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';
+import type { CaptureLambdaHandlerOptions } from '../types';
 
 /**
  * A middy middleware automating capture of metadata and annotations on segments or subsegments for a Lambda Handler.
@@ -25,9 +25,10 @@ import type { HandlerOptions } from '../types';
  * ```
  * 
  * @param target - The Tracer instance to use for tracing
+ * @param options - (_optional_) Options for the middleware
  * @returns middleware object - The middy middleware object
  */
-const captureLambdaHandler = (target: Tracer, options?: HandlerOptions): middy.MiddlewareObj => {
+const captureLambdaHandler = (target: Tracer, options?: CaptureLambdaHandlerOptions): middy.MiddlewareObj => {
   let lambdaSegment: Subsegment | Segment;
 
   const open = (): void => {

packages/tracer/src/types/Tracer.ts

@@ -29,21 +29,70 @@ type TracerOptions = {
 /**
  * Options for handler decorators and middleware.
  *
+ * Options supported:
+ * * `captureResponse` - (_optional_) - Disable response serialization as subsegment metadata
+ * 
+ * Middleware usage:
+ * @example
+ * ```typescript
+ * import middy from '@middy/core';
+ * 
+ * const tracer = new Tracer();
+ * 
+ * const lambdaHandler = async (_event: any, _context: any): Promise<void> => {};
+ * 
+ * export const handler = middy(lambdaHandler)
+ *  .use(captureLambdaHandler(tracer, { captureResponse: false }));
+ * ```
+ * 
+ * Decorator usage:
+ * @example
+ * ```typescript
+ * const tracer = new Tracer();
+ *
+ * class Lambda implements LambdaInterface {
+ *   @tracer.captureLambdaHandler({ captureResponse: false })
+ *   public async handler(_event: any, _context: any): Promise<void> {}
+ * }
+ *
+ * const handlerClass = new Lambda();
+ * export const handler = handlerClass.handler.bind(handlerClass);
+ * ```
+ */
+type CaptureLambdaHandlerOptions = {
+  captureResponse?: boolean
+};
+
+/**
+ * Options for method decorators.
+ * 
+ * Options supported:
+ * * `subSegmentName` - (_optional_) - Set a custom name for the subsegment
+ * * `captureResponse` - (_optional_) - Disable response serialization as subsegment metadata
+ *
  * Usage:
  * @example
  * ```typescript
  * const tracer = new Tracer();
  *
  * class Lambda implements LambdaInterface {
+ *   @tracer.captureMethod({ subSegmentName: 'gettingChargeId', captureResponse: false })
+ *   private getChargeId(): string {
+ *     return 'foo bar';
+ *   }
+ *
  *   @tracer.captureLambdaHandler({ captureResponse: false })
- *   async handler(_event: any, _context: any): Promise<void> {}
+ *   public async handler(_event: any, _context: any): Promise<void> {
+ *     this.getChargeId();
+ *   }
  * }
  *
  * const handlerClass = new Lambda();
  * export const handler = handlerClass.handler.bind(handlerClass);
  * ```
  */
-type HandlerOptions = {
+type CaptureMethodOptions = {
+  subSegmentName?: string
   captureResponse?: boolean
 };
 
@@ -59,7 +108,8 @@ type MethodDecorator = (target: any, propertyKey: string | symbol, descriptor: T
 
 export {
   TracerOptions,
-  HandlerOptions,
+  CaptureLambdaHandlerOptions,
+  CaptureMethodOptions,
   HandlerMethodDecorator,
   MethodDecorator
 };

packages/tracer/tests/e2e/allFeatures.decorator.test.functionCode.ts

@@ -9,8 +9,8 @@ const serviceName = process.env.EXPECTED_SERVICE_NAME ?? 'MyFunctionWithStandard
 const customAnnotationKey = process.env.EXPECTED_CUSTOM_ANNOTATION_KEY ?? 'myAnnotation';
 const customAnnotationValue = process.env.EXPECTED_CUSTOM_ANNOTATION_VALUE ?? 'myValue';
 const customMetadataKey = process.env.EXPECTED_CUSTOM_METADATA_KEY ?? 'myMetadata';
-const customMetadataValue = JSON.parse(process.env.EXPECTED_CUSTOM_METADATA_VALUE) ?? { bar: 'baz' };
-const customResponseValue = JSON.parse(process.env.EXPECTED_CUSTOM_RESPONSE_VALUE) ?? { foo: 'bar' };
+const customMetadataValue = process.env.EXPECTED_CUSTOM_METADATA_VALUE ? JSON.parse(process.env.EXPECTED_CUSTOM_METADATA_VALUE) : { bar: 'baz' };
+const customResponseValue = process.env.EXPECTED_CUSTOM_RESPONSE_VALUE ? JSON.parse(process.env.EXPECTED_CUSTOM_RESPONSE_VALUE) : { foo: 'bar' };
 const customErrorMessage = process.env.EXPECTED_CUSTOM_ERROR_MESSAGE ?? 'An error has occurred';
 const testTableName = process.env.TEST_TABLE_NAME ?? 'TestTable';
 
@@ -42,8 +42,6 @@ export class MyFunctionBase {
     this.returnValue = customResponseValue;
   }
 
-  // eslint-disable-next-line @typescript-eslint/ban-ts-comment
-  // @ts-ignore
   public handler(event: CustomEvent, _context: Context, _callback: Callback<unknown>): void | Promise<unknown> {
     tracer.putAnnotation(customAnnotationKey, customAnnotationValue);
     tracer.putMetadata(customMetadataKey, customMetadataValue);
@@ -78,8 +76,6 @@ export class MyFunctionBase {
       });
   }
 
-  // eslint-disable-next-line @typescript-eslint/ban-ts-comment
-  // @ts-ignore
   public myMethod(): string {
     return this.returnValue;
   }
@@ -121,4 +117,4 @@ class MyFunctionWithDecoratorCaptureResponseFalse extends MyFunctionBase {
 }
 
 const handlerWithCaptureResponseFalseClass = new MyFunctionWithDecoratorCaptureResponseFalse();
-export const handlerWithCaptureResponseFalse = handlerClass.handler.bind(handlerWithCaptureResponseFalseClass);
+export const handlerWithCaptureResponseFalse = handlerWithCaptureResponseFalseClass.handler.bind(handlerWithCaptureResponseFalseClass);

packages/tracer/tests/e2e/allFeatures.decorator.test.ts

@@ -62,7 +62,7 @@ let startTime: Date;
  * Function #1 is with all flags enabled.
  */
 const uuidFunction1 = v4();
-const functionNameWithAllFlagsEnabled = generateUniqueName(RESOURCE_NAME_PREFIX, uuidFunction1, runtime, 'AllFeatures-Decoratory-AllFlagsEnabled');
+const functionNameWithAllFlagsEnabled = generateUniqueName(RESOURCE_NAME_PREFIX, uuidFunction1, runtime, 'AllFeatures-Decorator-AllFlagsEnabled');
 const serviceNameWithAllFlagsEnabled = functionNameWithAllFlagsEnabled; 
 
 /**
@@ -79,7 +79,7 @@ const functionNameWithTracerDisabled = generateUniqueName(RESOURCE_NAME_PREFIX,
 const serviceNameWithTracerDisabled = functionNameWithNoCaptureErrorOrResponse; 
 
 /**
- * Function #4 disables tracer
+ * Function #4 disables capture response via decorator options
  */
 const uuidFunction4 = v4();
 const functionNameWithCaptureResponseFalse = generateUniqueName(RESOURCE_NAME_PREFIX, uuidFunction4, runtime, 'AllFeatures-Decorator-CaptureResponseFalse');

packages/tracer/tests/e2e/asyncHandler.decorator.test.functionCode.ts

@@ -9,10 +9,11 @@ const serviceName = process.env.EXPECTED_SERVICE_NAME ?? 'MyFunctionWithStandard
 const customAnnotationKey = process.env.EXPECTED_CUSTOM_ANNOTATION_KEY ?? 'myAnnotation';
 const customAnnotationValue = process.env.EXPECTED_CUSTOM_ANNOTATION_VALUE ?? 'myValue';
 const customMetadataKey = process.env.EXPECTED_CUSTOM_METADATA_KEY ?? 'myMetadata';
-const customMetadataValue = JSON.parse(process.env.EXPECTED_CUSTOM_METADATA_VALUE) ?? { bar: 'baz' };
-const customResponseValue = JSON.parse(process.env.EXPECTED_CUSTOM_RESPONSE_VALUE) ?? { foo: 'bar' };
+const customMetadataValue = process.env.EXPECTED_CUSTOM_METADATA_VALUE ? JSON.parse(process.env.EXPECTED_CUSTOM_METADATA_VALUE) : { bar: 'baz' };
+const customResponseValue = process.env.EXPECTED_CUSTOM_RESPONSE_VALUE ? JSON.parse(process.env.EXPECTED_CUSTOM_RESPONSE_VALUE) : { foo: 'bar' };
 const customErrorMessage = process.env.EXPECTED_CUSTOM_ERROR_MESSAGE ?? 'An error has occurred';
 const testTableName = process.env.TEST_TABLE_NAME ?? 'TestTable';
+const customSubSegmentName = process.env.EXPECTED_CUSTOM_SUBSEGMENT_NAME ?? 'mySubsegment';
 
 interface CustomEvent {
   throw: boolean
@@ -35,16 +36,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 async handler(event: CustomEvent, _context: Context): Promise<unknown> {
     tracer.putAnnotation(customAnnotationKey, customAnnotationValue);
     tracer.putMetadata(customMetadataKey, customMetadataValue);
@@ -74,13 +72,45 @@ export class MyFunctionWithDecorator {
     }
   }
 
+  public myMethod(): string {
+    return this.returnValue;
+  }
+}
+
+class MyFunctionWithDecorator extends MyFunctionBase {
+  @tracer.captureLambdaHandler()
+  // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+  // @ts-ignore
+  public async handler(event: CustomEvent, _context: Context, _callback: Callback<unknown>): void | Promise<unknown> {
+    return super.handler(event, _context);
+  }
+
   @tracer.captureMethod()
   // eslint-disable-next-line @typescript-eslint/ban-ts-comment
   // @ts-ignore
   public myMethod(): string {
-    return this.returnValue;
+    return super.myMethod();
   }
 }
 
 const handlerClass = new MyFunctionWithDecorator();
-export const handler = handlerClass.handler.bind(handlerClass);
+export const handler = handlerClass.handler.bind(handlerClass);
+
+export class MyFunctionWithDecoratorAndCustomNamedSubSegmentForMethod extends MyFunctionBase {
+  @tracer.captureLambdaHandler()
+  // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+  // @ts-ignore
+  public async handler(event: CustomEvent, _context: Context, _callback: Callback<unknown>): void | Promise<unknown> {
+    return super.handler(event, _context);
+  }
+
+  @tracer.captureMethod({ subSegmentName: customSubSegmentName })
+  // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+  // @ts-ignore
+  public myMethod(): string {
+    return super.myMethod();
+  }
+}
+
+const handlerWithCustomSubsegmentNameInMethodClass = new MyFunctionWithDecoratorAndCustomNamedSubSegmentForMethod();
+export const handlerWithCustomSubsegmentNameInMethod = handlerClass.handler.bind(handlerWithCustomSubsegmentNameInMethodClass);