feat(logger): set correlation ID in Logger

docs/core/logger.md

@@ -207,6 +207,72 @@ 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 gest started, install `@aws-lambda-powertools/jmespath` package and configure the search function for the correlationId:
+
+=== "Setup the Logger to use JMESPath search"
+
+    ```typescript
+    --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
+    --8<-- "examples/snippets/logger/correlationIdManual.ts"
+    ```
+
+=== "Middy.js Middleware"
+
+    ```typescript
+    --8<-- "examples/snippets/logger/correlationIdMiddy.ts"
+    ```
+
+=== "Decorator"
+
+    ```typescript
+    --8<-- "examples/snippets/logger/correlationIdDecorator.ts"
+    ```
+
+=== "payload.json"
+
+    ```typescript
+    --8<-- "examples/snippets/logger/correlationIdPayload.json"
+    ```
+
+=== "log-output.json"
+
+    ```json hl_lines="6"
+    --8<-- "examples/snippets/logger/correlationIdOutput.json"
+    ```
+
+To ease routine tasks like extracting correlation ID from popular event sources, we provide built-in JMESPath expressions.
+
+=== "Decorator"
+
+    ```typescript
+    --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_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({
+    ccorrelationIdPath: '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,13 @@
+import { Logger } from '@aws-lambda-powertools/logger';
+import type { APIGatewayProxyEvent, Context } from 'aws-lambda';
+
+const logger = new Logger();
+
+export const handler = async (
+  event: APIGatewayProxyEvent,
+  context: Context
+): Promise<void> => {
+  logger.setCorrelationId(event.requestContext.requestId);
+
+  logger.info('This is an INFO log with some context');
+};

examples/snippets/logger/correlationIdMiddy.ts

@@ -0,0 +1,21 @@
+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,
+});
+
+const lambdaHandler = async (
+  _event: unknown,
+  _context: unknown
+): Promise<void> => {
+  logger.info('This is an INFO log with some context');
+};
+
+export const handler = middy(lambdaHandler).use(
+  injectLambdaContext(logger, {
+    correlationIdPath: 'headers.my_request_id_header',
+  })
+);

examples/snippets/logger/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/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({
+    ccorrelationIdPath: 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/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.
    */
@@ -480,6 +485,9 @@ class Logger extends Utility implements LoggerInterface {
         loggerRef.refreshSampleRateCalculation();
         loggerRef.addContext(context);
         loggerRef.logEventIfEnabled(event, options?.logEvent);
+        if (options?.correlationIdPath) {
+          loggerRef.setCorrelationIdFromPath(options?.correlationIdPath, event);
+        }
 
         try {
           return await originalMethod.apply(this, [event, context, callback]);
@@ -1261,6 +1269,7 @@ class Logger extends Utility implements LoggerInterface {
       jsonReplacerFn,
       logRecordOrder,
       logBufferOptions,
+      correlationIdSearchFn,
     } = options;
 
     if (persistentLogAttributes && persistentKeys) {
@@ -1287,6 +1296,7 @@ class Logger extends Utility implements LoggerInterface {
     this.setLogIndentation();
     this.#jsonReplacerFn = jsonReplacerFn;
     this.#setLogBuffering(logBufferOptions);
+    this.#correlationIdSearchFn = correlationIdSearchFn;
 
     return this;
   }
@@ -1439,6 +1449,44 @@ class Logger extends Utility implements LoggerInterface {
       logLevel <= this.#bufferConfig.bufferAtVerbosity
     );
   }
+
+  /**
+   * Extract the correlation ID from the event using the provided path and append it to the log attributes.
+   * @param correlationIdPath - The path to the correlation ID in the event object
+   * @param event - The event object
+   */
+  protected setCorrelationIdFromPath(
+    correlationIdPath: string,
+    event: unknown
+  ): void {
+    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.'
+      );
+    } else {
+      const correlationId = this.#correlationIdSearchFn(
+        correlationIdPath,
+        event
+      );
+      if (correlationId) this.setCorrelationId(correlationId);
+    }
+  }
+
+  /**
+   * Set the correlation ID in the log attributes.
+   * @param value - The value to set as the correlation ID
+   */
+  public setCorrelationId(value: unknown): void {
+    this.appendKeys({ correlation_id: value });
+  }
+
+  /**
+   * Get the correlation ID from the log attributes.
+   * @returns 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,23 @@
+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';
+
+const search = (expression: string, data: JSONObject) => {
+  return JMESPathSearch(expression, data, {
+    customFunctions: new PowertoolsFunctions(),
+  });
+};
+
+const correlationPaths = {
+  API_GATEWAY_REST: 'requestContext.requestId',
+  API_GATEWAY_HTTP: 'requestContext.requestId',
+  APPSYNC_AUTHORIZER: 'requestContext.requestId',
+  APPSYNC_RESOLVER: 'request.headers."x-amzn-trace-id"',
+  APPLICATION_LOAD_BALANCER: 'headers."x-amzn-trace-id"',
+  EVENT_BRIDGE: 'id',
+  LAMBDA_FUNCTION_URL: 'requestContext.requestId',
+  S3_OBJECT_LAMBDA: 'xAmzRequestId',
+  VPC_LATTICE: 'headers."x-amzn-trace-id"',
+};
+
+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,13 @@ const injectLambdaContext = (
         request.context,
         options
       );
+
+      if (options?.correlationIdPath) {
+        logger.setCorrelationIdFromPath(
+          options.correlationIdPath,
+          request.event
+        );
+      }
     }
   };
 

packages/logger/src/types/Logger.ts

@@ -1,4 +1,7 @@
-import type { HandlerMethodDecorator } from '@aws-lambda-powertools/commons/types';
+import type {
+  HandlerMethodDecorator,
+  JSONObject,
+} from '@aws-lambda-powertools/commons/types';
 import type { Context } from 'aws-lambda';
 import type { LogLevel as LogLevelList } from '../constants.js';
 import type { LogFormatter } from '../formatter/LogFormatter.js';
@@ -54,6 +57,11 @@ type InjectLambdaContextOptions = {
    * @default `false`
    */
   flushBufferOnUncaughtError?: boolean;
+
+  /**
+   * The path to the correlation ID in the event object, used to extract the correlation ID from the event object and add it to the log attributes.
+   */
+  correlationIdPath?: string;
 };
 
 /**
@@ -217,13 +225,21 @@ type LogBufferOption = {
   };
 };
 
+type CorrelationIdOption = {
+  /**
+   * The search function for the correlation ID.
+   */
+  correlationIdSearchFn?: (expression: string, data: JSONObject) => unknown;
+};
+
 /**
  * Options to configure the Logger.
  */
 type ConstructorOptions = BaseConstructorOptions &
   (PersistentKeysOption | DeprecatedPersistentKeysOption) &
   (LogFormatterOption | LogRecordOrderOption) &
-  LogBufferOption;
+  LogBufferOption &
+  CorrelationIdOption;
 
 type LogItemMessage = string | LogAttributesWithMessage;
 type LogItemExtraInput = [Error | string] | LogAttributes[];
@@ -260,14 +276,14 @@ type LoggerInterface = {
 };
 
 export type {
+  ConstructorOptions,
+  CustomJsonReplacerFn,
   Environment,
+  InjectLambdaContextOptions,
   LogAttributes,
-  LogLevel,
   LogFunction,
   LoggerInterface,
-  LogItemMessage,
   LogItemExtraInput,
-  ConstructorOptions,
-  InjectLambdaContextOptions,
-  CustomJsonReplacerFn,
+  LogItemMessage,
+  LogLevel,
 };