diff --git a/packages/logger/src/Logger.ts b/packages/logger/src/Logger.ts index f6e05cdf63..ece63b3bd5 100644 --- a/packages/logger/src/Logger.ts +++ b/packages/logger/src/Logger.ts @@ -275,8 +275,6 @@ class Logger extends Utility implements ClassThatLogs { * @returns {HandlerMethodDecorator} */ public injectLambdaContext(options?: HandlerOptions): HandlerMethodDecorator { - const isClearStateEnabled = options && options.clearState === true; - return (target, _propertyKey, descriptor) => { /** * The descriptor.value is the method this decorator decorates, it cannot be undefined. @@ -286,23 +284,21 @@ class Logger extends Utility implements ClassThatLogs { descriptor.value = (event, context, callback) => { - let initialPersistentAttributes: LogAttributes = {}; - if (isClearStateEnabled) { + let initialPersistentAttributes = {}; + if (options && options.clearState === true) { initialPersistentAttributes = { ...this.getPersistentLogAttributes() }; } - this.addContext(context); - let shouldLogEvent = undefined; - if ( options && options.hasOwnProperty('logEvent') ) { - shouldLogEvent = options.logEvent; - } - this.logEventIfEnabled(event, shouldLogEvent); + Logger.injectLambdaContextBefore(this, event, context, options); /* eslint-disable @typescript-eslint/no-non-null-assertion */ - const result = originalMethod!.apply(target, [ event, context, callback ]); - - if (isClearStateEnabled) { - this.setPersistentLogAttributes(initialPersistentAttributes); + let result: unknown; + try { + result = originalMethod!.apply(target, [ event, context, callback ]); + } catch (error) { + throw error; + } finally { + Logger.injectLambdaContextAfterOrOnError(this, initialPersistentAttributes, options); } return result; @@ -310,6 +306,22 @@ class Logger extends Utility implements ClassThatLogs { }; } + public static injectLambdaContextAfterOrOnError(logger: Logger, initialPersistentAttributes: LogAttributes, options?: HandlerOptions): void { + if (options && options.clearState === true) { + logger.setPersistentLogAttributes(initialPersistentAttributes); + } + } + + public static injectLambdaContextBefore(logger: Logger, event: unknown, context: Context, options?: HandlerOptions): void { + logger.addContext(context); + + let shouldLogEvent = undefined; + if (options && options.hasOwnProperty('logEvent')) { + shouldLogEvent = options.logEvent; + } + logger.logEventIfEnabled(event, shouldLogEvent); + } + /** * Logs a Lambda invocation event, if it *should*. * diff --git a/packages/logger/src/middleware/middy.ts b/packages/logger/src/middleware/middy.ts index 6800686d18..85466c5de1 100644 --- a/packages/logger/src/middleware/middy.ts +++ b/packages/logger/src/middleware/middy.ts @@ -1,4 +1,4 @@ -import type { Logger } from '../Logger'; +import { Logger } from '../Logger'; import type middy from '@middy/core'; import { HandlerOptions, LogAttributes } from '../types'; @@ -34,30 +34,25 @@ const injectLambdaContext = (target: Logger | Logger[], options?: HandlerOptions const injectLambdaContextBefore = async (request: middy.Request): Promise => { loggers.forEach((logger: Logger) => { - logger.addContext(request.context); if (options && options.clearState === true) { persistentAttributes.push({ ...logger.getPersistentLogAttributes() }); } - - let shouldLogEvent = undefined; - if ( options && options.hasOwnProperty('logEvent') ) { - shouldLogEvent = options.logEvent; - } - logger.logEventIfEnabled(request.event, shouldLogEvent); + Logger.injectLambdaContextBefore(logger, request.event, request.context, options); }); }; - const injectLambdaContextAfter = async (): Promise => { + const injectLambdaContextAfterOrOnError = async (): Promise => { if (options && options.clearState === true) { loggers.forEach((logger: Logger, index: number) => { - logger.setPersistentLogAttributes(persistentAttributes[index]); + Logger.injectLambdaContextAfterOrOnError(logger, persistentAttributes[index], options); }); } }; return { before: injectLambdaContextBefore, - after: injectLambdaContextAfter + after: injectLambdaContextAfterOrOnError, + onError: injectLambdaContextAfterOrOnError }; }; diff --git a/packages/logger/tests/unit/Logger.test.ts b/packages/logger/tests/unit/Logger.test.ts index 88b628453e..43cf6fc0c1 100644 --- a/packages/logger/tests/unit/Logger.test.ts +++ b/packages/logger/tests/unit/Logger.test.ts @@ -982,6 +982,52 @@ describe('Class: Logger', () => { }); + test('when used as decorator with the clear state flag enabled and the handler throws an error, the persistent log attributes added in the handler are removed after the handler\'s code is executed', async () => { + + // Prepare + const logger = new Logger({ + logLevel: 'DEBUG', + persistentLogAttributes: { + foo: 'bar', + biz: 'baz' + } + }); + + type CustomEvent = { user_id: string }; + + class LambdaFunction implements LambdaInterface { + + @logger.injectLambdaContext({ clearState: true }) + // eslint-disable-next-line @typescript-eslint/ban-ts-comment + // @ts-ignore + public handler(event: CustomEvent, _context: Context, _callback: Callback): void | Promise { + // Only add these persistent for the scope of this lambda handler + logger.appendKeys({ + details: { user_id: event['user_id'] } + }); + logger.debug('This is a DEBUG log with the user_id'); + logger.debug('This is another DEBUG log with the user_id'); + + throw new Error('Unexpected error occurred!'); + } + } + + const persistentAttribs = { ...logger.getPersistentLogAttributes() }; + + // Act & Assess + const executeLambdaHandler = async (): Promise => { + await new LambdaFunction().handler({ user_id: '123456' }, dummyContext, () => console.log('Lambda invoked!')); + }; + await expect(executeLambdaHandler()).rejects.toThrow('Unexpected error occurred!'); + const persistentAttribsAfterInvocation = { ...logger.getPersistentLogAttributes() }; + expect(persistentAttribs).toEqual({ + foo: 'bar', + biz: 'baz' + }); + expect(persistentAttribsAfterInvocation).toEqual(persistentAttribs); + + }); + test('when used as decorator with the log event flag enabled, it logs the event', async () => { // Prepare diff --git a/packages/logger/tests/unit/middleware/middy.test.ts b/packages/logger/tests/unit/middleware/middy.test.ts index d6eb08b0f6..e734e61e61 100644 --- a/packages/logger/tests/unit/middleware/middy.test.ts +++ b/packages/logger/tests/unit/middleware/middy.test.ts @@ -207,6 +207,59 @@ describe('Middy middleware', () => { }); + test('when enabled and the handler throws an error, the persistent log attributes added within the handler scope are removed after the invocation ends', async () => { + + // Prepare + const logger = new Logger({ + logLevel: 'DEBUG', + persistentLogAttributes: { + foo: 'bar', + biz: 'baz' + } + }); + const context = { + callbackWaitsForEmptyEventLoop: true, + functionVersion: '$LATEST', + functionName: 'foo-bar-function', + memoryLimitInMB: '128', + logGroupName: '/aws/lambda/foo-bar-function', + logStreamName: '2021/03/09/[$LATEST]abcdef123456abcdef123456abcdef123456', + invokedFunctionArn: 'arn:aws:lambda:eu-west-1:123456789012:function:foo-bar-function', + awsRequestId: 'abcdef123456abcdef123456', + getRemainingTimeInMillis: () => 1234, + done: () => console.log('Done!'), + fail: () => console.log('Failed!'), + succeed: () => console.log('Succeeded!'), + }; + + const lambdaHandler = (event: { user_id: string }): void => { + // Only add these persistent for the scope of this lambda handler + logger.appendKeys({ + details: { user_id: event['user_id'] } + }); + logger.debug('This is a DEBUG log with the user_id'); + logger.debug('This is another DEBUG log with the user_id'); + + throw new Error('Unexpected error occurred!'); + }; + + const persistentAttribs = { ...logger.getPersistentLogAttributes() }; + const handler = middy(lambdaHandler).use(injectLambdaContext(logger, { clearState: true })); + + // Act & Assess + const executeLambdaHandler = async (): Promise => { + await handler({ user_id: '123456' }, context, () => console.log('Lambda invoked!')); + }; + await expect(executeLambdaHandler()).rejects.toThrow('Unexpected error occurred!'); + const persistentAttribsAfterInvocation = { ...logger.getPersistentLogAttributes() }; + expect(persistentAttribs).toEqual({ + foo: 'bar', + biz: 'baz' + }); + expect(persistentAttribsAfterInvocation).toEqual(persistentAttribs); + + }); + }); describe('Feature: log event', () => {