Skip to content

Commit

Permalink
feat(tracer): allow disabling result capture for decorators and middl…
Browse files Browse the repository at this point in the history
…eware (#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 <dreamorosi@gmail.com>

Co-authored-by: Andrea Amorosi <dreamorosi@gmail.com>
  • Loading branch information
misterjoshua and dreamorosi authored Aug 23, 2022
1 parent 8b8b25c commit c3b9a37
Show file tree
Hide file tree
Showing 11 changed files with 513 additions and 16 deletions.
65 changes: 64 additions & 1 deletion docs/core/tracer.md
Original file line number Diff line number Diff line change
Expand Up @@ -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`.
Expand Down Expand Up @@ -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.
Expand Down
15 changes: 10 additions & 5 deletions packages/tracer/src/Tracer.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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';

Expand Down Expand Up @@ -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.
Expand All @@ -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;
Expand Down Expand Up @@ -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
Expand All @@ -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);

Expand Down
9 changes: 6 additions & 3 deletions packages/tracer/src/middleware/middy.ts
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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 => {
Expand All @@ -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();
}
};
Expand Down
22 changes: 22 additions & 0 deletions packages/tracer/src/types/Tracer.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -38,6 +59,7 @@ type MethodDecorator = (target: any, propertyKey: string | symbol, descriptor: T

export {
TracerOptions,
HandlerOptions,
HandlerMethodDecorator,
MethodDecorator
};
Original file line number Diff line number Diff line change
Expand Up @@ -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> {
Expand Down Expand Up @@ -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);
80 changes: 80 additions & 0 deletions packages/tracer/tests/e2e/allFeatures.decorator.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand All @@ -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 }));
Loading

0 comments on commit c3b9a37

Please sign in to comment.