feat(logger): add clearState() method (#2408)

Co-authored-by: Andrea Amorosi <[email protected]>

Author: 2 people

Diff for: packages/logger/src/Logger.ts

@@ -172,6 +172,10 @@ class Logger extends Utility implements LoggerInterface {
    * Standard attributes managed by Powertools that will be logged in all log items.
    */
   private powertoolsLogData: PowertoolsLogData = <PowertoolsLogData>{};
+  /**
+   * Temporary log attributes that can be appended with `appendKeys()` method.
+   */
+  private temporaryLogAttributes: LogAttributes = {};
   /**
    * Buffer used to store logs until the logger is initialized.
    *
@@ -183,6 +187,13 @@ class Logger extends Utility implements LoggerInterface {
    * Flag used to determine if the logger is initialized.
    */
   #isInitialized = false;
+  /**
+   * Map used to hold the list of keys and their type.
+   *
+   * Because keys of different types can be overwritten, we keep a list of keys that were added and their last
+   * type. We then use this map at log preparation time to pick the last one.
+   */
+  #keys: Map<string, 'temp' | 'persistent'> = new Map();
 
   /**
    * Log level used by the current instance of Logger.
@@ -234,23 +245,40 @@ class Logger extends Utility implements LoggerInterface {
   }
 
   /**
-   * It adds the given attributes (key-value pairs) to all log items generated by this Logger instance.
+   * It adds the given persistent attributes (key-value pairs) to all log items generated by this Logger instance.
+   *
+   * @deprecated This method is deprecated and will be removed in the future major versions, please use {@link appendPersistentKeys()} instead.
    *
    * @param {LogAttributes} attributes
    * @returns {void}
    */
-  public addPersistentLogAttributes(attributes?: LogAttributes): void {
-    merge(this.persistentLogAttributes, attributes);
+  public addPersistentLogAttributes(attributes: LogAttributes): void {
+    this.appendPersistentKeys(attributes);
   }
 
   /**
-   * Alias for addPersistentLogAttributes.
+   * It adds the given temporary attributes (key-value pairs) to all log items generated by this Logger instance.
    *
    * @param {LogAttributes} attributes
    * @returns {void}
    */
-  public appendKeys(attributes?: LogAttributes): void {
-    this.addPersistentLogAttributes(attributes);
+  public appendKeys(attributes: LogAttributes): void {
+    for (const attributeKey of Object.keys(attributes)) {
+      this.#keys.set(attributeKey, 'temp');
+    }
+    merge(this.temporaryLogAttributes, attributes);
+  }
+
+  /**
+   * It adds the given persistent attributes (key-value pairs) to all log items generated by this Logger instance.
+   *
+   * @param attributes - The attributes to add to all log items.
+   */
+  public appendPersistentKeys(attributes: LogAttributes): void {
+    for (const attributeKey of Object.keys(attributes)) {
+      this.#keys.set(attributeKey, 'persistent');
+    }
+    merge(this.persistentLogAttributes, attributes);
   }
 
   /**
@@ -274,6 +302,7 @@ class Logger extends Utility implements LoggerInterface {
           customConfigService: this.getCustomConfigService(),
           environment: this.powertoolsLogData.environment,
           persistentLogAttributes: this.persistentLogAttributes,
+          temporaryLogAttributes: this.temporaryLogAttributes,
         },
         options
       )
@@ -417,13 +446,6 @@ class Logger extends Utility implements LoggerInterface {
         context,
         callback
       ) {
-        let initialPersistentAttributes = {};
-        if (options && options.clearState === true) {
-          initialPersistentAttributes = {
-            ...loggerRef.getPersistentLogAttributes(),
-          };
-        }
-
         Logger.injectLambdaContextBefore(loggerRef, event, context, options);
 
         let result: unknown;
@@ -432,25 +454,25 @@ class Logger extends Utility implements LoggerInterface {
         } catch (error) {
           throw error;
         } finally {
-          Logger.injectLambdaContextAfterOrOnError(
-            loggerRef,
-            initialPersistentAttributes,
-            options
-          );
+          if (options?.clearState) loggerRef.resetState();
         }
 
         return result;
       };
     };
   }
 
+  /**
+   * @deprecated This method is deprecated and will be removed in the future major versions. Use {@link resetState()} instead.
+   */
+  /* istanbul ignore next */
   public static injectLambdaContextAfterOrOnError(
     logger: Logger,
-    initialPersistentAttributes: LogAttributes,
+    _persistentAttributes: LogAttributes,
     options?: InjectLambdaContextOptions
   ): void {
     if (options && options.clearState === true) {
-      logger.setPersistentLogAttributes(initialPersistentAttributes);
+      logger.resetState();
     }
   }
 
@@ -493,27 +515,53 @@ class Logger extends Utility implements LoggerInterface {
   }
 
   /**
-   * Alias for removePersistentLogAttributes.
+   * It removes temporary attributes based on provided keys to all log items generated by this Logger instance.
    *
    * @param {string[]} keys
    * @returns {void}
    */
   public removeKeys(keys: string[]): void {
-    this.removePersistentLogAttributes(keys);
+    for (const key of keys) {
+      this.temporaryLogAttributes[key] = undefined;
+
+      if (this.persistentLogAttributes[key]) {
+        this.#keys.set(key, 'persistent');
+      } else {
+        this.#keys.delete(key);
+      }
+    }
   }
 
   /**
-   * It removes attributes based on provided keys to all log items generated by this Logger instance.
+   * It removes persistent attributes based on provided keys to all log items generated by this Logger instance.
    *
    * @param {string[]} keys
    * @returns {void}
    */
   public removePersistentLogAttributes(keys: string[]): void {
     for (const key of keys) {
-      if (this.persistentLogAttributes && key in this.persistentLogAttributes) {
-        delete this.persistentLogAttributes[key];
+      this.persistentLogAttributes[key] = undefined;
+
+      if (this.temporaryLogAttributes[key]) {
+        this.#keys.set(key, 'temp');
+      } else {
+        this.#keys.delete(key);
+      }
+    }
+  }
+
+  /**
+   * It resets the state, by removing all temporary log attributes added with `appendKeys()` method.
+   */
+  public resetState(): void {
+    for (const key of Object.keys(this.temporaryLogAttributes)) {
+      if (this.persistentLogAttributes[key]) {
+        this.#keys.set(key, 'persistent');
+      } else {
+        this.#keys.delete(key);
       }
     }
+    this.temporaryLogAttributes = {};
   }
 
   /**
@@ -537,6 +585,8 @@ class Logger extends Utility implements LoggerInterface {
    * It sets the given attributes (key-value pairs) to all log items generated by this Logger instance.
    * Note: this replaces the pre-existing value.
    *
+   * @deprecated This method is deprecated and will be removed in the future major versions, please use {@link appendPersistentKeys()} instead.
+   *
    * @param {LogAttributes} attributes
    * @returns {void}
    */
@@ -665,10 +715,18 @@ class Logger extends Utility implements LoggerInterface {
       ...this.getPowertoolsLogData(),
     };
 
-    // gradually merge additional attributes starting from customer-provided persistent attributes
-    let additionalLogAttributes = { ...this.getPersistentLogAttributes() };
+    const additionalAttributes: LogAttributes = {};
+    // gradually add additional attributes picking only the last added for each key
+    for (const [key, type] of this.#keys) {
+      if (type === 'persistent') {
+        additionalAttributes[key] = this.persistentLogAttributes[key];
+      } else {
+        additionalAttributes[key] = this.temporaryLogAttributes[key];
+      }
+    }
+
     // if the main input is not a string, then it's an object with additional attributes, so we merge it
-    additionalLogAttributes = merge(additionalLogAttributes, otherInput);
+    merge(additionalAttributes, otherInput);
     // then we merge the extra input attributes (if any)
     for (const item of extraInput) {
       const attributes: LogAttributes =
@@ -678,12 +736,12 @@ class Logger extends Utility implements LoggerInterface {
             ? { extra: item }
             : item;
 
-      additionalLogAttributes = merge(additionalLogAttributes, attributes);
+      merge(additionalAttributes, attributes);
     }
 
     return this.getLogFormatter().formatAttributes(
       unformattedBaseAttributes,
-      additionalLogAttributes
+      additionalAttributes
     );
   }
 
@@ -1023,13 +1081,23 @@ class Logger extends Utility implements LoggerInterface {
       serviceName,
       sampleRateValue,
       logFormatter,
-      persistentLogAttributes,
+      persistentKeys,
+      persistentLogAttributes, // deprecated in favor of persistentKeys
       environment,
     } = options;
 
+    if (persistentLogAttributes && persistentKeys) {
+      this.warn(
+        'Both persistentLogAttributes and persistentKeys options were provided. Using persistentKeys as persistentLogAttributes is deprecated and will be removed in future releases'
+      );
+    }
+
     // configurations that affect log content
-    this.setPowertoolsLogData(serviceName, environment);
-    this.addPersistentLogAttributes(persistentLogAttributes);
+    this.setPowertoolsLogData(
+      serviceName,
+      environment,
+      persistentKeys || persistentLogAttributes
+    );
 
     // configurations that affect Logger behavior
     this.setLogEvent();
@@ -1070,7 +1138,7 @@ class Logger extends Utility implements LoggerInterface {
         this.getEnvVarsService().getServiceName() ||
         this.getDefaultServiceName(),
     });
-    this.addPersistentLogAttributes(persistentLogAttributes);
+    this.appendPersistentKeys(persistentLogAttributes);
   }
 }
 

Diff for: packages/logger/src/middleware/middy.ts

@@ -1,5 +1,4 @@
 import { Logger } from '../Logger.js';
-import type { LogAttributes } from '../types/Log.js';
 import type { InjectLambdaContextOptions } from '../types/Logger.js';
 import { LOGGER_KEY } from '@aws-lambda-powertools/commons';
 import type {
@@ -37,7 +36,6 @@ const injectLambdaContext = (
   options?: InjectLambdaContextOptions
 ): MiddlewareLikeObj => {
   const loggers = target instanceof Array ? target : [target];
-  const persistentAttributes: LogAttributes[] = [];
   const isClearState = options && options.clearState === true;
 
   /**
@@ -55,12 +53,8 @@ const injectLambdaContext = (
   const injectLambdaContextBefore = async (
     request: MiddyLikeRequest
   ): Promise<void> => {
-    loggers.forEach((logger: Logger, index: number) => {
+    loggers.forEach((logger: Logger) => {
       if (isClearState) {
-        persistentAttributes[index] = {
-          ...logger.getPersistentLogAttributes(),
-        };
-
         setCleanupFunction(request);
       }
       Logger.injectLambdaContextBefore(
@@ -74,12 +68,8 @@ const injectLambdaContext = (
 
   const injectLambdaContextAfterOrOnError = async (): Promise<void> => {
     if (isClearState) {
-      loggers.forEach((logger: Logger, index: number) => {
-        Logger.injectLambdaContextAfterOrOnError(
-          logger,
-          persistentAttributes[index],
-          options
-        );
+      loggers.forEach((logger: Logger) => {
+        logger.resetState();
       });
     }
   };

Diff for: packages/logger/src/types/Logger.ts

@@ -21,16 +21,41 @@ type InjectLambdaContextOptions = {
   clearState?: boolean;
 };
 
-type ConstructorOptions = {
+type BaseConstructorOptions = {
   logLevel?: LogLevel;
   serviceName?: string;
   sampleRateValue?: number;
   logFormatter?: LogFormatterInterface;
   customConfigService?: ConfigServiceInterface;
-  persistentLogAttributes?: LogAttributes;
   environment?: Environment;
 };
 
+type PersistentKeysOption = {
+  persistentKeys?: LogAttributes;
+  persistentLogAttributes?: never;
+};
+
+type DeprecatedOption = {
+  persistentLogAttributes?: LogAttributes;
+  persistentKeys?: never;
+};
+
+/**
+ * Options for the Logger class constructor.
+ *
+ * @type {Object} ConstructorOptions
+ * @property {LogLevel} [logLevel] - The log level.
+ * @property {string} [serviceName] - The service name.
+ * @property {number} [sampleRateValue] - The sample rate value.
+ * @property {LogFormatterInterface} [logFormatter] - The custom log formatter.
+ * @property {ConfigServiceInterface} [customConfigService] - The custom config service.
+ * @property {Environment} [environment] - The environment.
+ * @property {LogAttributes} [persistentKeys] - The keys that will be persisted in all log items.
+ * @property {LogAttributes} [persistentLogAttributes] - **Deprecated!** Use `persistentKeys`.
+ */
+type ConstructorOptions = BaseConstructorOptions &
+  (PersistentKeysOption | DeprecatedOption);
+
 type LambdaFunctionContext = Pick<
   Context,
   | 'functionName'