chore(logger): removed unused var

Author: saragerion

Diff for: 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,30 +284,43 @@ class Logger extends Utility implements ClassThatLogs {
 
       descriptor.value = (event, context, callback) => {
 
-        let initialPersistentAttributes: LogAttributes = {};
-        if (isClearStateEnabled) {
-          initialPersistentAttributes = { ...this.getPersistentLogAttributes() };
-        }
-
-        this.addContext(context);
-        let shouldLogEvent = undefined;
-        if ( options && options.hasOwnProperty('logEvent') ) {
-          shouldLogEvent = options.logEvent;
-        }
-        this.logEventIfEnabled(event, shouldLogEvent);
+        const initialPersistentAttributes = 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;
       };
     };
   }
 
+  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): LogAttributes | undefined {
+    logger.addContext(context);
+
+    let shouldLogEvent = undefined;
+    if (options && options.hasOwnProperty('logEvent')) {
+      shouldLogEvent = options.logEvent;
+    }
+    logger.logEventIfEnabled(event, shouldLogEvent);
+
+    if (options && options.clearState === true) {
+      return { ...logger.getPersistentLogAttributes() };
+    }
+  }
+
   /**
    * Logs a Lambda invocation event, if it *should*.
    *

Diff for: 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';
 
@@ -30,34 +30,29 @@ import { HandlerOptions, LogAttributes } from '../types';
 const injectLambdaContext = (target: Logger | Logger[], options?: HandlerOptions): middy.MiddlewareObj => {
 
   const loggers = target instanceof Array ? target : [target];
-  const persistentAttributes: LogAttributes[] = [];
+  const persistentAttributes: (LogAttributes | undefined)[] = [];
 
   const injectLambdaContextBefore = async (request: middy.Request): Promise<void> => {
     loggers.forEach((logger: Logger) => {
-      logger.addContext(request.context);
+      const initialPersistentAttributes = Logger.injectLambdaContextBefore(logger, request.event, request.context, options);
       if (options && options.clearState === true) {
-        persistentAttributes.push({ ...logger.getPersistentLogAttributes() });
+        persistentAttributes.push(initialPersistentAttributes);
       }
-
-      let shouldLogEvent = undefined;
-      if ( options && options.hasOwnProperty('logEvent') ) {
-        shouldLogEvent = options.logEvent;
-      }
-      logger.logEventIfEnabled(request.event, shouldLogEvent);
     });
   };
 
-  const injectLambdaContextAfter = async (): Promise<void> => {
+  const injectLambdaContextAfterOrOnError = async (): Promise<void> => {
     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 for: packages/logger/tests/unit/Logger.test.ts

@@ -982,6 +982,55 @@ describe('Class: Logger', () => {
 
     });
 
+    test('when used as decorator with the clear state flag enabled and the handler throws, 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<TResult>(event: CustomEvent, _context: Context, _callback: Callback<TResult>): void | Promise<TResult> {
+          // 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
+      try {
+        await new LambdaFunction().handler({ user_id: '123456' }, dummyContext, () => console.log('Lambda invoked!'));
+      } catch (error) {
+        // Do nothing
+      }
+      const persistentAttribsAfterInvocation = { ...logger.getPersistentLogAttributes() };
+
+      // Assess
+      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 for: packages/logger/tests/unit/middleware/middy.test.ts

@@ -207,6 +207,62 @@ describe('Middy middleware', () => {
 
     });
 
+    test('when enabled and the handler throws, 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 handler = middy(lambdaHandler).use(injectLambdaContext(logger, { clearState: true }));
+      const persistentAttribs = { ...logger.getPersistentLogAttributes() };
+
+      // Act
+      try {
+        await handler({ user_id: '123456' }, context, () => console.log('Lambda invoked!'));
+      } catch (error) {
+        // Do nothing
+      }
+      const persistentAttribsAfterInvocation = { ...logger.getPersistentLogAttributes() };
+
+      // Assess
+      expect(persistentAttribs).toEqual({
+        foo: 'bar',
+        biz: 'baz'
+      });
+      expect(persistentAttribsAfterInvocation).toEqual(persistentAttribs);
+
+    });
+
   });
 
   describe('Feature: log event', () => {