feat(logger): add clear state functionality

docs/core/logger.md

@@ -273,6 +273,125 @@ You can append additional persistent keys and values in the logs generated durin
 
 !!! tip "Logger will automatically ignore any key with an `undefined` value"
 
+#### Clearing all state
+
+The Logger utility is commonly initialized in the global scope, outside the handler function.
+When you attach persistent log attributes through the `persistentLogAttributes` constructor option or via the `appendKeys`, `addPersistentLogAttributes` methods, this data is attached to the Logger instance.  
+
+Due to the [Lambda Execution Context reuse](https://docs.aws.amazon.com/lambda/latest/dg/runtimes-context.html), this means those persistent log attributes may be reused across invocations.
+If you want to make sure that persistent attributes added **inside the handler function** code are not persisted across invocations, you can set the parameter `clearState` as `true`  in the `injectLambdaContext` middleware or decorator.
+
+=== "Middy Middleware"
+
+    ```typescript hl_lines="27"
+    import { Logger, injectLambdaContext } from '@aws-lambda-powertools/logger';
+    import middy from '@middy/core';
+
+    // Persistent attributes added outside the handler will be 
+    // cached across invocations
+    const logger = new Logger({
+        logLevel: 'DEBUG',
+        persistentLogAttributes: {
+            foo: "bar",
+            biz: "baz"
+        }
+    });
+
+    const lambdaHandler = async (event: { special_key: string }, _context: any): Promise<void> => {
+        // Persistent attributes added inside the handler will NOT be cached
+        // across invocations
+        if (event['special_key'] === '123456') {
+            logger.appendKeys({
+                details: { special_key: event['special_key'] }
+            });
+        }
+        logger.debug('This is a DEBUG log');
+    };
+
+    // Enable the clear state flag
+    export const handler = middy(lambdaHandler)
+        .use(injectLambdaContext(logger, { clearState: true }));
+    ```
+
+=== "Decorator"
+
+    ```typescript hl_lines="16"
+    import { Logger } from '@aws-lambda-powertools/logger';
+    import { LambdaInterface } from '@aws-lambda-powertools/commons';
+
+    // Persistent attributes added outside the handler will be 
+    // cached across invocations
+    const logger = new Logger({
+        logLevel: 'DEBUG',
+        persistentLogAttributes: {
+            foo: "bar",
+            biz: "baz"
+        }
+    });
+
+    class Lambda implements LambdaInterface {
+        // Enable the clear state flag
+        @logger.injectLambdaContext({ clearState: true })
+        public async handler(_event: any, _context: any): Promise<void> {
+            // Persistent attributes added inside the handler will NOT be cached
+            // across invocations
+            if (event['special_key'] === '123456'){
+                logger.appendKeys({
+                    details: { special_key: '123456' }
+                });
+            }
+            logger.debug('This is a DEBUG log');
+        }
+
+    }
+
+    export const myFunction = new Lambda();
+    export const handler = myFunction.handler;
+    ```
+
+In each case, the printed log will look like this:
+
+=== "First invocation"
+
+    ```json hl_lines="2 4-7"
+    {
+        "biz": "baz",
+        "cold_start": true,
+        "details": {
+            "special_key": "123456",
+        },
+        "foo": "bar",
+        "function_arn": "arn:aws:lambda:eu-west-1:123456789012:function:foo-bar-function",
+        "function_memory_size": 128,
+        "function_name": "foo-bar-function",
+        "function_request_id": "abcdef123456abcdef123456",
+        "level": "DEBUG",
+        "message": "This is a DEBUG log with the user_id",
+        "service": "hello-world",
+        "timestamp": "2021-12-12T22:32:54.670Z",
+        "xray_trace_id": "1-5759e988-bd862e3fe1be46a994272793"
+    }
+    ```
+=== "Second invocation"
+
+    ```json hl_lines="2 4"
+    {
+        "biz": "baz",
+        "cold_start": false,
+        "foo": "bar",
+        "function_arn": "arn:aws:lambda:eu-west-1:123456789012:function:foo-bar-function",
+        "function_memory_size": 128,
+        "function_name": "foo-bar-function",
+        "function_request_id": "abcdef123456abcdef123456",
+        "level": "DEBUG",
+        "message": "This is a DEBUG log with the user_id",
+        "service": "hello-world",
+        "timestamp": "2021-12-12T22:40:23.120Z",
+        "xray_trace_id": "1-5759e988-bd862e3fe1be46a994272793"
+    }
+    ```
+
+
 ### Appending additional data to a single log item
 
 You can append additional data to a single log item by passing objects as additional parameters.
@@ -431,6 +550,7 @@ The error will be logged with default key name `error`, but you can also pass yo
 !!! tip "Logging errors and log level"
     You can also log errors using the `warn`, `info`, and `debug` methods. Be aware of the log level though, you might miss those  errors when analyzing the log later depending on the log level configuration.
 
+
 ## Advanced
 
 ### Using multiple Logger instances across your code

packages/logger/src/Logger.ts

@@ -12,12 +12,13 @@ import type {
   HandlerMethodDecorator,
   LambdaFunctionContext,
   LogAttributes,
-  LoggerOptions,
+  ConstructorOptions,
   LogItemExtraInput,
   LogItemMessage,
   LogLevel,
   LogLevelThresholds,
   PowertoolLogData,
+  HandlerOptions,
 } from './types';
 
 /**
@@ -138,9 +139,9 @@ class Logger extends Utility implements ClassThatLogs {
   /**
    * It initializes the Logger class with an optional set of options (settings).
    * *
-   * @param {LoggerOptions} options
+   * @param {ConstructorOptions} options
    */
-  public constructor(options: LoggerOptions = {}) {
+  public constructor(options: ConstructorOptions = {}) {
     super();
 
     this.setOptions(options);
@@ -192,10 +193,10 @@ class Logger extends Utility implements ClassThatLogs {
    * It creates a separate Logger instance, identical to the current one
    * It's possible to overwrite the new instance options by passing them.
    *
-   * @param {LoggerOptions} options
+   * @param {ConstructorOptions} options
    * @returns {Logger}
    */
-  public createChild(options: LoggerOptions = {}): Logger {
+  public createChild(options: ConstructorOptions = {}): Logger {
     return cloneDeep(this).setOptions(options);
   }
 
@@ -230,6 +231,17 @@ class Logger extends Utility implements ClassThatLogs {
     return this.logsSampled;
   }
 
+  /**
+   * It returns the persistent log attributes, which are the attributes
+   * that will be logged in all log items.
+   *
+   * @private
+   * @returns {LogAttributes}
+   */
+  public getPersistentLogAttributes(): LogAttributes {
+    return <LogAttributes> this.persistentLogAttributes;
+  }
+
   /**
    * It prints a log item with level INFO.
    *
@@ -250,14 +262,27 @@ class Logger extends Utility implements ClassThatLogs {
    * @see https://www.typescriptlang.org/docs/handbook/decorators.html#method-decorators
    * @returns {HandlerMethodDecorator}
    */
-  public injectLambdaContext(): HandlerMethodDecorator {
+  public injectLambdaContext(options?: HandlerOptions): HandlerMethodDecorator {
+    const isClearStateEnabled = options && options.clearState === true;
+
     return (target, _propertyKey, descriptor) => {
       const originalMethod = descriptor.value;
-
       descriptor.value = (event, context, callback) => {
+
+        let initialPersistentAttributes: LogAttributes = {};
+        if (isClearStateEnabled) {
+          initialPersistentAttributes = { ...this.getPersistentLogAttributes() };
+        }
+
         this.addContext(context);
 
-        return originalMethod?.apply(target, [ event, context, callback ]);
+        const result = originalMethod?.apply(target, [ event, context, callback ]);
+
+        if (isClearStateEnabled) {
+          this.setPersistentLogAttributes(initialPersistentAttributes);
+        }
+
+        return result;
       };
     };
   }
@@ -274,6 +299,17 @@ class Logger extends Utility implements ClassThatLogs {
     this.setLogsSampled();
   }
 
+  /**
+   * It sets the given attributes (key-value pairs) to all log items generated by this Logger instance.
+   * Note: this replaces the pre-existing value.
+   *
+   * @param {LogAttributes} attributes
+   * @returns {void}
+   */
+  public setPersistentLogAttributes(attributes: LogAttributes): void {
+    this.persistentLogAttributes = attributes;
+  }
+
   /**
    * It sets the user-provided sample rate value.
    *
@@ -393,17 +429,6 @@ class Logger extends Utility implements ClassThatLogs {
     return <LogLevel> this.logLevel;
   }
 
-  /**
-   * It returns the persistent log attributes, which are the attributes
-   * that will be logged in all log items.
-   *
-   * @private
-   * @returns {LogAttributes}
-   */
-  private getPersistentLogAttributes(): LogAttributes {
-    return <LogAttributes> this.persistentLogAttributes;
-  }
-
   /**
    * It returns information that will be added in all log item by
    * this Logger instance (different from user-provided persistent attributes).
@@ -596,10 +621,10 @@ class Logger extends Utility implements ClassThatLogs {
    * and the content of all logs.
    *
    * @private
-   * @param {LoggerOptions} options
+   * @param {ConstructorOptions} options
    * @returns {Logger}
    */
-  private setOptions(options: LoggerOptions): Logger {
+  private setOptions(options: ConstructorOptions): Logger {
     const {
       logLevel,
       serviceName,

packages/logger/src/helpers.ts

@@ -1,7 +1,7 @@
 import { Logger } from '.';
-import { LoggerOptions } from './types';
+import { ConstructorOptions } from './types';
 
-const createLogger = (options: LoggerOptions = {}): Logger => new Logger(options);
+const createLogger = (options: ConstructorOptions = {}): Logger => new Logger(options);
 
 export {
   createLogger,

packages/logger/src/middleware/middy.ts

@@ -1,5 +1,6 @@
 import type { Logger } from '../Logger';
 import type middy from '@middy/core';
+import { HandlerOptions, LogAttributes } from '../types';
 
 /**
  * A middy middleware that adds the current Lambda invocation's context inside all log items.
@@ -26,16 +27,31 @@ import type middy from '@middy/core';
  * @param {Logger|Logger[]} target - The Tracer instance to use for tracing
  * @returns {middy.MiddlewareObj} - The middy middleware object
  */
-const injectLambdaContext = (target: Logger | Logger[]): middy.MiddlewareObj => {
+const injectLambdaContext = (target: Logger | Logger[], options?: HandlerOptions): middy.MiddlewareObj => {
+
+  const loggers = target instanceof Array ? target : [target];
+  const persistentAttributes: LogAttributes[] = [];
+
   const injectLambdaContextBefore = async (request: middy.Request): Promise<void> => {
-    const loggers = target instanceof Array ? target : [target];
     loggers.forEach((logger: Logger) => {
       logger.addContext(request.context);
+      if (options && options.clearState === true) {
+        persistentAttributes.push({ ...logger.getPersistentLogAttributes() });
+      }
     });
   };
+
+  const injectLambdaContextAfter = async (): Promise<void> => {
+    if (options && options.clearState === true) {
+      loggers.forEach((logger: Logger, index: number) => {
+        logger.setPersistentLogAttributes(persistentAttributes[index]);
+      });
+    }
+  };
 
   return {
     before: injectLambdaContextBefore,
+    after: injectLambdaContextAfter
   };
 };
 

packages/logger/src/types/Logger.ts

@@ -8,7 +8,11 @@ type ClassThatLogs = {
   [key in 'debug' | 'error' | 'info' | 'warn']: (input: LogItemMessage, ...extraInput: LogItemExtraInput) => void;
 };
 
-type LoggerOptions = {
+type HandlerOptions = {
+  clearState?: boolean
+};
+
+type ConstructorOptions = {
   logLevel?: LogLevel
   serviceName?: string
   sampleRateValue?: number
@@ -66,5 +70,6 @@ export {
   LambdaFunctionContext,
   UnformattedAttributes,
   PowertoolLogData,
-  LoggerOptions,
+  ConstructorOptions,
+  HandlerOptions
 };

packages/logger/tests/unit/Logger.test.ts

@@ -843,6 +843,49 @@ describe('Class: Logger', () => {
 
     });
 
+    test('when used as decorator with the clear state flag enabled, 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');
+        }
+      }
+
+      const persistentAttribs = { ...logger.getPersistentLogAttributes() };
+
+      // Act
+      await new LambdaFunction().handler({ user_id: '123456' }, dummyContext, () => console.log('Lambda invoked!'));
+      const persistentAttribsAfterInvocation = { ...logger.getPersistentLogAttributes() };
+
+      // Assess
+      expect(persistentAttribs).toEqual({
+        foo: 'bar',
+        biz: 'baz'
+      });
+      expect(persistentAttribsAfterInvocation).toEqual(persistentAttribs);
+
+    });
+
   });
 
   describe('Method: refreshSampleRateCalculation', () => {