feat(idempotency): manipulate idempotent response via response hook (#3071)

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

Author: 3 people

docs/utilities/idempotency.md

@@ -373,6 +373,40 @@ sequenceDiagram
 <i>Idempotent successful request cached</i>
 </center>
 
+#### Successful request with responseHook configured
+
+<center>
+```mermaid
+sequenceDiagram
+    participant Client
+    participant Lambda
+    participant Response hook
+    participant Persistence Layer
+    alt initial request
+        Client->>Lambda: Invoke (event)
+        Lambda->>Persistence Layer: Get or set idempotency_key=hash(payload)
+        activate Persistence Layer
+        Note over Lambda,Persistence Layer: Set record status to INPROGRESS. <br> Prevents concurrent invocations <br> with the same payload
+        Lambda-->>Lambda: Call your function
+        Lambda->>Persistence Layer: Update record with result
+        deactivate Persistence Layer
+        Persistence Layer-->>Persistence Layer: Update record
+        Note over Lambda,Persistence Layer: Set record status to COMPLETE. <br> New invocations with the same payload <br> now return the same result
+        Lambda-->>Client: Response sent to client
+    else retried request
+        Client->>Lambda: Invoke (event)
+        Lambda->>Persistence Layer: Get or set idempotency_key=hash(payload)
+        activate Persistence Layer
+        Persistence Layer-->>Response hook: Already exists in persistence layer.
+        deactivate Persistence Layer
+        Note over Response hook,Persistence Layer: Record status is COMPLETE and not expired
+        Response hook->>Lambda: Response hook invoked
+        Lambda-->>Client: Manipulated idempotent response sent to client
+    end
+```
+<i>Successful idempotent request with a response hook</i>
+</center>
+
 #### Expired idempotency records
 
 <center>
@@ -544,6 +578,7 @@ Idempotent decorator can be further configured with **`IdempotencyConfig`** as s
 | **useLocalCache**             | `false` | Whether to locally cache idempotency results                                                                                                                                               |
 | **localCacheMaxItems**        | 256     | Max number of items to store in local cache                                                                                                                                                |
 | **hashFunction**              | `md5`   | Function to use for calculating hashes, as provided by the [crypto](https://nodejs.org/api/crypto.html#cryptocreatehashalgorithm-options){target="_blank"} module in the standard library. |
+| **responseHook**               | `undefined`  | Function to use for processing the stored Idempotent response. This function hook is called when an existing idempotent response is found. See [Manipulating The Idempotent Response](idempotency.md#manipulating-the-idempotent-response) |
 
 ### Handling concurrent executions with the same payload
 
@@ -744,6 +779,42 @@ Below an example implementation of a custom persistence layer backed by a generi
 
     For example, the `_putRecord()` method needs to throw an error if a non-expired record already exists in the data store with a matching key.
 
+### Manipulating the Idempotent Response
+
+You can set up a `responseHook` in the `IdempotentConfig` class to manipulate the returned data when an operation is idempotent. The hook function will be called with the current deserialized response object and the Idempotency record.
+
+=== "Using an Idempotent Response Hook"
+
+    ```typescript hl_lines="16 19 27 56"
+    --8<-- "examples/snippets/idempotency/workingWithResponseHook.ts"
+    ```
+
+=== "Sample event"
+
+    ```json
+    --8<-- "examples/snippets/idempotency/samples/workingWithResponseHookSampleEvent.json"
+    ```
+
+=== "Sample Idempotent response"
+
+    ```json hl_lines="6"
+    --8<-- "examples/snippets/idempotency/samples/workingWithResponseHookIdempotentResponse.json"
+    ```
+
+???+ info "Info: Using custom de-serialization?"
+
+    The responseHook is called after the custom de-serialization so the payload you process will be the de-serialized version.
+
+#### Being a good citizen
+
+When using response hooks to manipulate returned data from idempotent operations, it's important to follow best practices to avoid introducing complexity or issues. Keep these guidelines in mind:
+
+1. **Response hook works exclusively when operations are idempotent.** The hook will not be called when an operation is not idempotent, or when the idempotent logic fails.
+
+2. **Catch and Handle Exceptions.** Your response hook code should catch and handle any exceptions that may arise from your logic. Unhandled exceptions will cause the Lambda function to fail unexpectedly.
+
+3. **Keep Hook Logic Simple** Response hooks should consist of minimal and straightforward logic for manipulating response data. Avoid complex conditional branching and aim for hooks that are easy to reason about.
+
 ## Testing your code
 
 The idempotency utility provides several routes to test your code.

examples/snippets/idempotency/samples/workingWithResponseHookIdempotentResponse.json

@@ -0,0 +1,8 @@
+{
+    "message": "success",
+    "paymentId": "31a964eb-7477-4fe1-99fe-7f8a6a351a7e",
+    "statusCode": 200,
+    "headers": {
+      "x-idempotency-key": "function-name#mHfGv2vJ8h+ZvLIr/qGBbQ=="
+    }
+  }

examples/snippets/idempotency/samples/workingWithResponseHookSampleEvent.json

@@ -0,0 +1,4 @@
+{
+    "user": "John Doe",
+    "productId": "123456"
+}

examples/snippets/idempotency/workingWithResponseHook.ts

@@ -0,0 +1,58 @@
+import { randomUUID } from 'node:crypto';
+import type { JSONValue } from '@aws-lambda-powertools/commons/types';
+import {
+  IdempotencyConfig,
+  makeIdempotent,
+} from '@aws-lambda-powertools/idempotency';
+import { DynamoDBPersistenceLayer } from '@aws-lambda-powertools/idempotency/dynamodb';
+import type { IdempotencyRecord } from '@aws-lambda-powertools/idempotency/persistence';
+import type { Context } from 'aws-lambda';
+import type { Request, Response, SubscriptionResult } from './types.js';
+
+const persistenceStore = new DynamoDBPersistenceLayer({
+  tableName: 'idempotencyTableName',
+});
+
+const responseHook = (response: JSONValue, record: IdempotencyRecord) => {
+  // Return inserted Header data into the Idempotent Response
+  (response as Response).headers = {
+    'x-idempotency-key': record.idempotencyKey,
+  };
+
+  // Must return the response here
+  return response as JSONValue;
+};
+
+const config = new IdempotencyConfig({
+  responseHook,
+});
+
+const createSubscriptionPayment = async (
+  event: Request
+): Promise<SubscriptionResult> => {
+  // ... create payment
+  return {
+    id: randomUUID(),
+    productId: event.productId,
+  };
+};
+
+export const handler = makeIdempotent(
+  async (event: Request, _context: Context): Promise<Response> => {
+    try {
+      const payment = await createSubscriptionPayment(event);
+
+      return {
+        paymentId: payment.id,
+        message: 'success',
+        statusCode: 200,
+      };
+    } catch (error) {
+      throw new Error('Error creating payment');
+    }
+  },
+  {
+    persistenceStore,
+    config,
+  }
+);

packages/idempotency/src/IdempotencyConfig.ts

@@ -2,7 +2,10 @@ import { PowertoolsFunctions } from '@aws-lambda-powertools/jmespath/functions';
 import type { JMESPathParsingOptions } from '@aws-lambda-powertools/jmespath/types';
 import type { Context } from 'aws-lambda';
 import { EnvironmentVariablesService } from './config/EnvironmentVariablesService.js';
-import type { IdempotencyConfigOptions } from './types/IdempotencyOptions.js';
+import type {
+  IdempotencyConfigOptions,
+  ResponseHook,
+} from './types/IdempotencyOptions.js';
 
 /**
  * Configuration for the idempotency feature.
@@ -52,6 +55,10 @@ class IdempotencyConfig {
    * @default false
    */
   public throwOnNoIdempotencyKey: boolean;
+  /**
+   * A hook that runs when an idempotent request is made.
+   */
+  public responseHook?: ResponseHook;
 
   /**
    * Use the local cache to store idempotency keys.
@@ -70,6 +77,7 @@ class IdempotencyConfig {
     this.maxLocalCacheSize = config.maxLocalCacheSize ?? 1000;
     this.hashFunction = config.hashFunction ?? 'md5';
     this.lambdaContext = config.lambdaContext;
+    this.responseHook = config.responseHook;
     this.#envVarsService = new EnvironmentVariablesService();
     this.#enabled = this.#envVarsService.getIdempotencyEnabled();
   }

packages/idempotency/src/IdempotencyHandler.ts

@@ -87,12 +87,14 @@ export class IdempotencyHandler<Func extends AnyFunction> {
   /**
    * Takes an idempotency key and returns the idempotency record from the persistence layer.
    *
+   * If a response hook is provided in the idempotency configuration, it will be called before returning the response.
+   *
    * If the idempotency record is not COMPLETE, then it will throw an error based on the status of the record.
    *
    * @param idempotencyRecord The idempotency record stored in the persistence layer
    * @returns The result of the function if the idempotency record is in a terminal state
    */
-  public static determineResultFromIdempotencyRecord(
+  public determineResultFromIdempotencyRecord(
     idempotencyRecord: IdempotencyRecord
   ): JSONValue {
     if (idempotencyRecord.getStatus() === IdempotencyRecordStatus.EXPIRED) {
@@ -115,7 +117,14 @@ export class IdempotencyHandler<Func extends AnyFunction> {
       );
     }
 
-    return idempotencyRecord.getResponse();
+    const response = idempotencyRecord.getResponse();
+
+    // If a response hook is provided, call it to allow the user to modify the response
+    if (this.#idempotencyConfig.responseHook) {
+      return this.#idempotencyConfig.responseHook(response, idempotencyRecord);
+    }
+
+    return response;
   }
 
   /**
@@ -381,9 +390,7 @@ export class IdempotencyHandler<Func extends AnyFunction> {
 
         returnValue.isIdempotent = true;
         returnValue.result =
-          IdempotencyHandler.determineResultFromIdempotencyRecord(
-            idempotencyRecord
-          );
+          this.determineResultFromIdempotencyRecord(idempotencyRecord);
 
         return returnValue;
       }

packages/idempotency/src/types/IdempotencyOptions.ts

@@ -2,6 +2,7 @@ import type { JSONValue } from '@aws-lambda-powertools/commons/types';
 import type { Context, Handler } from 'aws-lambda';
 import type { IdempotencyConfig } from '../IdempotencyConfig.js';
 import type { BasePersistenceLayer } from '../persistence/BasePersistenceLayer.js';
+import type { IdempotencyRecord } from '../persistence/IdempotencyRecord.js';
 
 /**
  * Configuration options for the idempotency utility.
@@ -147,6 +148,14 @@ type IdempotencyHandlerOptions = {
   thisArg?: Handler;
 };
 
+/**
+ * A hook that runs when an idempotent request is made.
+ */
+type ResponseHook = (
+  response: JSONValue,
+  record: IdempotencyRecord
+) => JSONValue;
+
 /**
  * Idempotency configuration options
  */
@@ -183,6 +192,10 @@ type IdempotencyConfigOptions = {
    * AWS Lambda Context object containing information about the current invocation, function, and execution environment
    */
   lambdaContext?: Context;
+  /**
+   * A hook that runs when an idempotent request is made
+   */
+  responseHook?: ResponseHook;
 };
 
 export type {
@@ -191,4 +204,5 @@ export type {
   ItempotentFunctionOptions,
   IdempotencyLambdaHandlerOptions,
   IdempotencyHandlerOptions,
+  ResponseHook,
 };

packages/idempotency/src/types/index.ts

@@ -12,6 +12,7 @@ export type {
   IdempotencyHandlerOptions,
   ItempotentFunctionOptions,
   AnyFunction,
+  ResponseHook,
 } from './IdempotencyOptions.js';
 export type {
   DynamoDBPersistenceOptions,