Skip to content

Commit df2fc25

Browse files
dreamorosidreamorosi
committed
WIP: captureLambdaHandler middleware
1 parent a0b9c6d commit df2fc25

File tree

8 files changed

+242
-10
lines changed

8 files changed

+242
-10
lines changed

‎packages/tracing/npm-shrinkwrap.json

Lines changed: 17 additions & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎packages/tracing/package-lock.json

Lines changed: 0 additions & 6 deletions
This file was deleted.

‎packages/tracing/package.json

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@
5454
"url": "https://github.com/awslabs/aws-lambda-powertools-typescript/issues"
5555
},
5656
"dependencies": {
57+
"@middy/core": "^2.5.3",
5758
"aws-xray-sdk-core": "^3.3.3"
5859
}
59-
}
60+
}

‎packages/tracing/src/Tracer.ts

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -351,6 +351,14 @@ class Tracer implements TracerInterface {
351351

352352
return segment;
353353
}
354+
355+
public isCaptureErrorEnabled(): boolean {
356+
return this.captureError;
357+
}
358+
359+
public isCaptureResponseEnabled(): boolean {
360+
return this.captureResponse;
361+
}
354362

355363
/**
356364
* Retrieve the current value of `ColdStart`.
@@ -373,6 +381,10 @@ class Tracer implements TracerInterface {
373381
return false;
374382
}
375383

384+
public isTracingEnabled(): boolean {
385+
return this.tracingEnabled;
386+
}
387+
376388
/**
377389
* Adds annotation to existing segment or subsegment.
378390
*

‎packages/tracing/src/middleware/middy.ts

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
import middy from '@middy/core';
2+
import { Subsegment } from 'aws-xray-sdk-core';
3+
import { Tracer } from '../Tracer';
4+
5+
const captureLambdaHandler = (target: Tracer): middy.MiddlewareObj => {
6+
let subsegment: Subsegment | undefined;
7+
const captureLambdaHandlerBefore = async (request: middy.Request): Promise<void> => {
8+
if (target.isTracingEnabled()) {
9+
subsegment = new Subsegment(`## ${request.context.functionName}`);
10+
target.setSegment(subsegment);
11+
12+
if (Tracer.coldStart) {
13+
target.putAnnotation('ColdStart', true);
14+
}
15+
}
16+
};
17+
18+
const captureLambdaHandlerAfter = async (request: middy.Request): Promise<void> => {
19+
if (target.isTracingEnabled()) {
20+
if (request.error) {
21+
if (target.isCaptureErrorEnabled() === false) {
22+
subsegment?.addErrorFlag();
23+
} else {
24+
subsegment?.addError(request.error, false);
25+
}
26+
// TODO: should this error be thrown??
27+
// throw error;
28+
} else {
29+
if (request.response !== undefined && target.isCaptureResponseEnabled() === true) {
30+
target.putMetadata(`${request.context.functionName} response`, request.response);
31+
}
32+
}
33+
34+
subsegment?.close();
35+
}
36+
};
37+
38+
return {
39+
before: captureLambdaHandlerBefore,
40+
after: captureLambdaHandlerAfter,
41+
};
42+
};
43+
44+
export {
45+
captureLambdaHandler,
46+
};

‎packages/tracing/tests/unit/Tracer.test.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,9 @@ import { context as dummyContext } from '../../../../tests/resources/contexts/he
22
// eslint-disable-next-line @typescript-eslint/ban-ts-comment
33
// @ts-ignore
44
import * as dummyEvent from '../../../../tests/resources/events/custom/hello-world.json';
5+
// import { captureLambdaHandler } from '../../src/middleware/middy';
56
import { LambdaInterface } from '../../examples/utils/lambda';
7+
// import middy from '@middy/core';
68
import { Tracer } from '../../src';
79
import { Callback, Context } from 'aws-lambda/handler';
810
import { Segment, setContextMissingStrategy, Subsegment } from 'aws-xray-sdk-core';

‎packages/tracing/tests/unit/middy.test.ts

Lines changed: 160 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,160 @@
1+
import { captureLambdaHandler } from '../../src/middleware/middy';
2+
import middy from '@middy/core';
3+
// import { EnvironmentVariablesService } from '../../../src/config';
4+
import { Tracer } from './../../src';
5+
import type { Context, Handler } from 'aws-lambda/handler';
6+
import { Segment, setContextMissingStrategy, Subsegment } from 'aws-xray-sdk-core';
7+
8+
jest.spyOn(console, 'debug').mockImplementation(() => null);
9+
jest.spyOn(console, 'warn').mockImplementation(() => null);
10+
jest.spyOn(console, 'error').mockImplementation(() => null);
11+
12+
describe('Middy middlewares', () => {
13+
const ENVIRONMENT_VARIABLES = process.env;
14+
15+
const mockContext: Context = {
16+
callbackWaitsForEmptyEventLoop: true,
17+
functionVersion: '$LATEST',
18+
functionName: 'foo-bar-function',
19+
memoryLimitInMB: '128',
20+
logGroupName: '/aws/lambda/foo-bar-function-123456abcdef',
21+
logStreamName: '2021/03/09/[$LATEST]abcdef123456abcdef123456abcdef123456',
22+
invokedFunctionArn: 'arn:aws:lambda:eu-central-1:123456789012:function:Example',
23+
awsRequestId: Math.floor(Math.random() * 1000000000).toString(),
24+
getRemainingTimeInMillis: () => 1234,
25+
done: () => console.log('Done!'),
26+
fail: () => console.log('Failed!'),
27+
succeed: () => console.log('Succeeded!'),
28+
};
29+
30+
beforeEach(() => {
31+
Tracer.coldStart = true;
32+
jest.clearAllMocks();
33+
jest.resetModules();
34+
process.env = { ...ENVIRONMENT_VARIABLES };
35+
});
36+
37+
afterAll(() => {
38+
process.env = ENVIRONMENT_VARIABLES;
39+
});
40+
describe('Middleware: captureLambdaHandler', () => {
41+
42+
test('when used while tracing is disabled, it does nothing', async () => {
43+
44+
// Prepare
45+
const tracer: Tracer = new Tracer({ enabled: false });
46+
const setSegmentSpy = jest.spyOn(tracer.provider, 'setSegment').mockImplementation();
47+
const getSegmentSpy = jest.spyOn(tracer.provider, 'getSegment')
48+
.mockImplementationOnce(() => new Segment('facade', process.env._X_AMZN_TRACE_ID || null))
49+
.mockImplementationOnce(() => new Subsegment('## foo-bar-function'));
50+
const lambdaHandler: Handler = async (_event: unknown, _context: Context) => ({
51+
foo: 'bar'
52+
});
53+
const handler = middy(lambdaHandler).use(captureLambdaHandler(tracer));
54+
const context = Object.assign({}, mockContext);
55+
56+
// Act
57+
await handler({}, context, () => console.log('Lambda invoked!'));
58+
59+
// Assess
60+
expect(setSegmentSpy).toHaveBeenCalledTimes(0);
61+
expect(getSegmentSpy).toHaveBeenCalledTimes(0);
62+
63+
});
64+
65+
test('when used as decorator while POWERTOOLS_TRACER_CAPTURE_RESPONSE is set to false, it does not capture the response as metadata', async () => {
66+
67+
// Prepare
68+
process.env.POWERTOOLS_TRACER_CAPTURE_RESPONSE = 'false';
69+
const tracer: Tracer = new Tracer();
70+
const newSubsegment: Segment | Subsegment | undefined = new Subsegment('## foo-bar-function');
71+
const setSegmentSpy = jest.spyOn(tracer.provider, 'setSegment').mockImplementation();
72+
jest.spyOn(tracer.provider, 'getSegment').mockImplementation(() => newSubsegment);
73+
setContextMissingStrategy(() => null);
74+
const lambdaHandler: Handler = async (_event: unknown, _context: Context) => ({
75+
foo: 'bar'
76+
});
77+
const handler = middy(lambdaHandler).use(captureLambdaHandler(tracer));
78+
const context = Object.assign({}, mockContext);
79+
80+
// Act
81+
await handler({}, context, () => console.log('Lambda invoked!'));
82+
83+
// Assess
84+
expect(setSegmentSpy).toHaveBeenCalledTimes(1);
85+
expect('metadata' in newSubsegment).toBe(false);
86+
delete process.env.POWERTOOLS_TRACER_CAPTURE_RESPONSE;
87+
88+
});
89+
90+
test('when used as decorator and with standard config, it captures the response as metadata', async () => {
91+
92+
// Prepare
93+
const tracer: Tracer = new Tracer();
94+
const newSubsegment: Segment | Subsegment | undefined = new Subsegment('## foo-bar-function');
95+
const setSegmentSpy = jest.spyOn(tracer.provider, 'setSegment').mockImplementation();
96+
jest.spyOn(tracer.provider, 'getSegment').mockImplementation(() => newSubsegment);
97+
setContextMissingStrategy(() => null);
98+
const lambdaHandler: Handler = async (_event: unknown, _context: Context) => ({
99+
foo: 'bar'
100+
});
101+
const handler = middy(lambdaHandler).use(captureLambdaHandler(tracer));
102+
const context = Object.assign({}, mockContext);
103+
104+
// Act
105+
await handler({}, context, () => console.log('Lambda invoked!'));
106+
107+
// Assess
108+
expect(setSegmentSpy).toHaveBeenCalledTimes(1);
109+
expect(setSegmentSpy).toHaveBeenCalledWith(expect.objectContaining({
110+
name: '## foo-bar-function',
111+
}));
112+
expect(newSubsegment).toEqual(expect.objectContaining({
113+
name: '## foo-bar-function',
114+
metadata: {
115+
'hello-world': {
116+
'foo-bar-function response': {
117+
foo: 'bar',
118+
},
119+
},
120+
}
121+
}));
122+
123+
});
124+
125+
test('when used as decorator while POWERTOOLS_TRACER_CAPTURE_ERROR is set to false, it does not capture the exceptions', async () => {
126+
127+
// Prepare
128+
process.env.POWERTOOLS_TRACER_CAPTURE_ERROR = 'false';
129+
const tracer: Tracer = new Tracer();
130+
const newSubsegment: Segment | Subsegment | undefined = new Subsegment('## foo-bar-function');
131+
const setSegmentSpy = jest.spyOn(tracer.provider, 'setSegment').mockImplementation();
132+
jest.spyOn(tracer.provider, 'getSegment').mockImplementation(() => newSubsegment);
133+
setContextMissingStrategy(() => null);
134+
const addErrorSpy = jest.spyOn(newSubsegment, 'addError');
135+
const addErrorFlagSpy = jest.spyOn(newSubsegment, 'addErrorFlag');
136+
const lambdaHandler: Handler = async (_event: unknown, _context: Context) => {
137+
throw new Error('Exception thrown!');
138+
};
139+
const handler = middy(lambdaHandler).use(captureLambdaHandler(tracer));
140+
const context = Object.assign({}, mockContext);
141+
142+
// Act
143+
await handler({}, context, () => console.log('Lambda invoked!'));
144+
145+
// Assess
146+
expect(setSegmentSpy).toHaveBeenCalledTimes(1);
147+
expect(setSegmentSpy).toHaveBeenCalledWith(expect.objectContaining({
148+
name: '## foo-bar-function',
149+
}));
150+
expect('cause' in newSubsegment).toBe(false);
151+
expect(addErrorFlagSpy).toHaveBeenCalledTimes(1);
152+
expect(addErrorSpy).toHaveBeenCalledTimes(0);
153+
154+
delete process.env.POWERTOOLS_TRACER_CAPTURE_ERROR;
155+
156+
});
157+
158+
});
159+
160+
});

‎packages/tracing/tsconfig.json

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,9 +14,10 @@
1414
"resolveJsonModule": true,
1515
"pretty": true,
1616
"baseUrl": "src/",
17-
"rootDirs": [ "src/" ]
17+
"rootDirs": [ "src/" ],
18+
"esModuleInterop": true
1819
},
19-
"include": [ "src/**/*", "examples/**/*", "tests/**/*" ],
20+
"include": [ "src/**/*", "examples/**/*", "**/tests/**/*" ],
2021
"exclude": [ "./node_modules"],
2122
"watchOptions": {
2223
"watchFile": "useFsEvents",

0 commit comments

Comments
 (0)