feat(logger): set correlation ID in Logger

docs/core/logger.md

@@ -207,6 +207,76 @@ When debugging in non-production environments, you can log the incoming event us
 
 Use `POWERTOOLS_LOGGER_LOG_EVENT` environment variable to enable or disable (`true`/`false`) this feature. When using Middy.js middleware or class method decorator, the `logEvent` option will take precedence over the environment variable.
 
+### Setting a Correlation ID
+
+To get started, install the `@aws-lambda-powertools/jmespath` package, and pass the search function using the `correlationIdSearchFn` constructor parameter:
+
+=== "Setup the Logger to use JMESPath search"
+
+    ```typescript hl_lines="5"
+    --8<-- "examples/snippets/logger/correlationIdLogger.ts"
+    ```
+
+???+ tip
+    You can retrieve correlation IDs via `getCorrelationId` method.
+
+You can set a correlation ID using `correlationIdPath` parameter by passing a JMESPath expression, including our custom JMESPath functions or set it manually by calling `setCorrelationId` function.
+
+=== "Setting correlation ID manually"
+
+    ```typescript hl_lines="7"
+    --8<-- "examples/snippets/logger/correlationIdManual.ts"
+    ```
+
+    1. Alternatively, if the payload is more complex you can use a JMESPath expression as second parameter when prividing a search function in the constructor.
+
+=== "Middy.js"
+
+    ```typescript hl_lines="13"
+    --8<-- "examples/snippets/logger/correlationIdMiddy.ts"
+    ```
+
+=== "Decorator"
+
+    ```typescript hl_lines="11"
+    --8<-- "examples/snippets/logger/correlationIdDecorator.ts"
+    ```
+
+=== "payload.json"
+
+    ```typescript
+    --8<-- "examples/snippets/logger/samples/correlationIdPayload.json"
+    ```
+
+=== "log-output.json"
+
+    ```json hl_lines="6"
+    --8<-- "examples/snippets/logger/samples/correlationIdOutput.json"
+    ```
+
+To ease routine tasks like extracting correlation ID from popular event sources, we provide built-in JMESPath expressions.
+
+=== "Decorator"
+
+    ```typescript hl_lines="4 14" 
+    --8<-- "examples/snippets/logger/correlationIdPaths.ts"
+    ```
+
+???+ note "Note: Any object key named with `-` must be escaped"
+    For example, **`request.headers."x-amzn-trace-id"`**.
+
+| Name                          | Expression                            | Description                     |
+| ----------------------------- | ------------------------------------- | ------------------------------- |
+| **API_GATEWAY_REST**          | `'requestContext.requestId'`          | API Gateway REST API request ID |
+| **API_GATEWAY_HTTP**          | `'requestContext.requestId'`          | API Gateway HTTP API request ID |
+| **APPSYNC_AUTHORIZER**        | `'requestContext.requestId'`          | AppSync resolver request ID     |
+| **APPSYNC_RESOLVER**          | `'request.headers."x-amzn-trace-id"'` | AppSync X-Ray Trace ID          |
+| **APPLICATION_LOAD_BALANCER** | `'headers."x-amzn-trace-id"'`         | ALB X-Ray Trace ID              |
+| **EVENT_BRIDGE**              | `'id'`                                | EventBridge Event ID            |
+| **LAMBDA_FUNCTION_URL**       | `'requestContext.requestId'`          | Lambda Function URL request ID  |
+| **S3_OBJECT_LAMBDA**          | `'xAmzRequestId'`                     | S3 Object trigger request ID    |
+| **VPC_LATTICE**               | `'headers."x-amzn-trace-id'`          | VPC Lattice X-Ray Trace ID      |
+
 ### Appending additional keys
 
 You can append additional keys using either mechanism:

examples/snippets/logger/correlationIdDecorator.ts

@@ -0,0 +1,19 @@
+import type { LambdaInterface } from '@aws-lambda-powertools/commons/types';
+import { Logger } from '@aws-lambda-powertools/logger';
+import { search } from '@aws-lambda-powertools/logger/correlationId';
+
+const logger = new Logger({
+  correlationIdSearchFn: search,
+});
+
+class Lambda implements LambdaInterface {
+  @logger.injectLambdaContext({
+    correlationIdPath: 'headers.my_request_id_header',
+  })
+  public async handler(_event: unknown, _context: unknown): Promise<void> {
+    logger.info('This is an INFO log with some context');
+  }
+}
+
+const myFunction = new Lambda();
+export const handler = myFunction.handler.bind(myFunction);

examples/snippets/logger/correlationIdLogger.ts

@@ -0,0 +1,6 @@
+import { Logger } from '@aws-lambda-powertools/logger';
+import { search } from '@aws-lambda-powertools/logger/correlationId';
+
+const logger = new Logger({
+  correlationIdSearchFn: search,
+});

examples/snippets/logger/correlationIdManual.ts

@@ -0,0 +1,10 @@
+import { Logger } from '@aws-lambda-powertools/logger';
+import type { APIGatewayProxyEvent } from 'aws-lambda';
+
+const logger = new Logger();
+
+export const handler = async (event: APIGatewayProxyEvent) => {
+  logger.setCorrelationId(event.requestContext.requestId); // (1)!
+
+  logger.info('log with correlation_id');
+};

examples/snippets/logger/correlationIdMiddy.ts

@@ -0,0 +1,18 @@
+import { Logger } from '@aws-lambda-powertools/logger';
+import { search } from '@aws-lambda-powertools/logger/correlationId';
+import { injectLambdaContext } from '@aws-lambda-powertools/logger/middleware';
+import middy from '@middy/core';
+
+const logger = new Logger({
+  correlationIdSearchFn: search,
+});
+
+export const handler = middy()
+  .use(
+    injectLambdaContext(logger, {
+      correlationIdPath: 'headers.my_request_id_header',
+    })
+  )
+  .handler(async () => {
+    logger.info('log with correlation_id');
+  });

examples/snippets/logger/correlationIdPaths.ts

@@ -0,0 +1,22 @@
+import type { LambdaInterface } from '@aws-lambda-powertools/commons/types';
+import { Logger } from '@aws-lambda-powertools/logger';
+import {
+  correlationPaths,
+  search,
+} from '@aws-lambda-powertools/logger/correlationId';
+
+const logger = new Logger({
+  correlationIdSearchFn: search,
+});
+
+class Lambda implements LambdaInterface {
+  @logger.injectLambdaContext({
+    correlationIdPath: correlationPaths.API_GATEWAY_REST,
+  })
+  public async handler(_event: unknown, _context: unknown): Promise<void> {
+    logger.info('This is an INFO log with some context');
+  }
+}
+
+const myFunction = new Lambda();
+export const handler = myFunction.handler.bind(myFunction);

examples/snippets/logger/samples/correlationIdOutput.json

@@ -0,0 +1,7 @@
+{
+  "level": "INFO",
+  "message": "This is an INFO log with some context",
+  "timestamp": "2021-05-03 11:47:12,494+0000",
+  "service": "payment",
+  "correlation_id": "correlation_id_value"
+}

examples/snippets/logger/samples/correlationIdPayload.json

@@ -0,0 +1,5 @@
+{
+  "headers": {
+    "my_request_id_header": "correlation_id_value"
+  }
+}

packages/logger/package.json

@@ -47,6 +47,10 @@
     "./types": {
       "import": "./lib/esm/types/index.js",
       "require": "./lib/cjs/types/index.js"
+    },
+    "./correlationId": {
+      "import": "./lib/esm/correlationId.js",
+      "require": "./lib/cjs/correlationId.js"
     }
   },
   "typesVersions": {
@@ -68,16 +72,18 @@
     "@types/lodash.merge": "^4.6.9"
   },
   "peerDependencies": {
-    "@middy/core": "4.x || 5.x || 6.x"
+    "@middy/core": "4.x || 5.x || 6.x",
+    "@aws-lambda-powertools/jmespath": "2.x"
   },
   "peerDependenciesMeta": {
     "@middy/core": {
       "optional": true
+    },
+    "@aws-lambda-powertools/jmespath": {
+      "optional": true
     }
   },
-  "files": [
-    "lib"
-  ],
+  "files": ["lib"],
   "repository": {
     "type": "git",
     "url": "git+https://github.com/aws-powertools/powertools-lambda-typescript.git"

packages/logger/src/Logger.ts

@@ -215,6 +215,11 @@ class Logger extends Utility implements LoggerInterface {
    */
   #buffer?: CircularMap<string>;
 
+  /**
+   * Search function for the correlation ID.
+   */
+  #correlationIdSearchFn?: (expression: string, data: unknown) => unknown;
+
   /**
    * The debug sampling rate configuration.
    */
@@ -326,6 +331,7 @@ class Logger extends Utility implements LoggerInterface {
           environment: this.powertoolsLogData.environment,
           persistentLogAttributes: this.persistentLogAttributes,
           jsonReplacerFn: this.#jsonReplacerFn,
+          correlationIdSearchFn: this.#correlationIdSearchFn,
           ...(this.#bufferConfig.enabled && {
             logBufferOptions: {
               maxBytes: this.#bufferConfig.maxBytes,
@@ -480,6 +486,9 @@ class Logger extends Utility implements LoggerInterface {
         loggerRef.refreshSampleRateCalculation();
         loggerRef.addContext(context);
         loggerRef.logEventIfEnabled(event, options?.logEvent);
+        if (options?.correlationIdPath) {
+          loggerRef.setCorrelationId(event, options?.correlationIdPath);
+        }
 
         try {
           return await originalMethod.apply(this, [event, context, callback]);
@@ -1261,6 +1270,7 @@ class Logger extends Utility implements LoggerInterface {
       jsonReplacerFn,
       logRecordOrder,
       logBufferOptions,
+      correlationIdSearchFn,
     } = options;
 
     if (persistentLogAttributes && persistentKeys) {
@@ -1287,6 +1297,7 @@ class Logger extends Utility implements LoggerInterface {
     this.setLogIndentation();
     this.#jsonReplacerFn = jsonReplacerFn;
     this.#setLogBuffering(logBufferOptions);
+    this.#correlationIdSearchFn = correlationIdSearchFn;
 
     return this;
   }
@@ -1439,6 +1450,56 @@ class Logger extends Utility implements LoggerInterface {
       logLevel <= this.#bufferConfig.bufferAtVerbosity
     );
   }
+
+  /**
+   * Set the correlation ID for the log item.
+   * This method can be used to set the correlation ID for the log item or to search for the correlation ID in the event.
+   *
+   * @example
+   * ```typescript
+   * import { Logger } from '@aws-lambda-powertools/logger';
+   *
+   * const logger = new Logger();
+   * logger.setCorrelationId('my-correlation-id'); // sets the correlation ID directly with the first argument as value
+   * ```
+   *
+   * ```typescript
+   * import { Logger } from '@aws-lambda-powertools/logger';
+   * import { search } from '@aws-lambda-powertools/logger/correlationId';
+   *
+   * const logger = new Logger({ correlationIdSearchFn: search });
+   * logger.setCorrelationId(event, 'requestContext.requestId'); // sets the correlation ID from the event using JMSPath expression
+   * ```
+   *
+   * @param value - The value to set as the correlation ID or the event to search for the correlation ID
+   * @param correlationIdPath - Optional JMESPath expression to extract the correlation ID for the payload
+   */
+  public setCorrelationId(value: unknown, correlationIdPath?: string): void {
+    if (typeof correlationIdPath === 'string') {
+      if (!this.#correlationIdSearchFn) {
+        this.warn(
+          'correlationIdPath is set but no search function was provided. The correlation ID will not be added to the log attributes.'
+        );
+        return;
+      }
+      const correlationId = this.#correlationIdSearchFn(
+        correlationIdPath,
+        value
+      );
+      if (correlationId) this.appendKeys({ correlation_id: correlationId });
+      return;
+    }
+
+    // If no correlationIdPath is provided, set the correlation ID directly
+    this.appendKeys({ correlation_id: value });
+  }
+
+  /**
+   * Get the correlation ID from the log attributes.
+   */
+  public getCorrelationId(): unknown {
+    return this.temporaryLogAttributes.correlation_id;
+  }
 }
 
 export { Logger };

packages/logger/src/correlationId.ts

@@ -0,0 +1,60 @@
+import type { JSONObject } from '@aws-lambda-powertools/commons/types';
+import { search as JMESPathSearch } from '@aws-lambda-powertools/jmespath';
+import { PowertoolsFunctions } from '@aws-lambda-powertools/jmespath/functions';
+
+/**
+ * This function is used to search for a correlation ID in the event data and is a wrapper
+ * around the JMESPath search function. It allows you to specify a JMESPath expression
+ * to extract the correlation ID from the event data.
+ * @param expression - The JMESPath expression to use for searching the correlation ID.
+ * @param data - The event data to search in.
+ */
+const search = (expression: string, data: unknown) => {
+  return JMESPathSearch(expression, data as JSONObject, {
+    customFunctions: new PowertoolsFunctions(),
+  });
+};
+
+/**
+ * The correlationPaths object contains the JMESPath expressions for extracting the correlation ID for various AWS services.
+ */
+const correlationPaths = {
+  /**
+   * API Gateway REST API request ID
+   */
+  API_GATEWAY_REST: 'requestContext.requestId',
+  /**
+   * API Gateway HTTP API request ID
+   */
+  API_GATEWAY_HTTP: 'requestContext.requestId',
+  /**
+   * AppSync API request ID
+   */
+  APPSYNC_AUTHORIZER: 'requestContext.requestId',
+  /**
+   * AppSync resolver X-Ray trace ID
+   */
+  APPSYNC_RESOLVER: 'request.headers."x-amzn-trace-id"',
+  /**
+   * ALB X-Ray trace ID
+   */
+  APPLICATION_LOAD_BALANCER: 'headers."x-amzn-trace-id"',
+  /**
+   * EventBridge event ID
+   */
+  EVENT_BRIDGE: 'id',
+  /**
+   * Lambda Function URL request ID
+   */
+  LAMBDA_FUNCTION_URL: 'requestContext.requestId',
+  /**
+   * S3 Object trigger request ID
+   */
+  S3_OBJECT_LAMBDA: 'xAmzRequestId',
+  /**
+   * VPC Lattice X-Ray trace ID
+   */
+  VPC_LATTICE: 'headers."x-amzn-trace-id"',
+} as const;
+
+export { correlationPaths, search };

packages/logger/src/index.ts

@@ -1,4 +1,4 @@
-export { Logger } from './Logger.js';
+export { LogLevel, LogLevelThreshold } from './constants.js';
 export { LogFormatter } from './formatter/LogFormatter.js';
 export { LogItem } from './formatter/LogItem.js';
-export { LogLevel, LogLevelThreshold } from './constants.js';
+export { Logger } from './Logger.js';

packages/logger/src/middleware/middy.ts

@@ -95,6 +95,10 @@ const injectLambdaContext = (
         request.context,
         options
       );
+
+      if (options?.correlationIdPath) {
+        logger.setCorrelationId(request.event, options.correlationIdPath);
+      }
     }
   };