test(idempotency): improve integration tests for utility

.github/workflows/run-e2e-tests.yml

@@ -19,7 +19,7 @@ jobs:
       contents: read
     strategy:
       matrix:
-        package: [logger, metrics, tracer, parameters]
+        package: [logger, metrics, tracer, parameters, idempotency]
         version: [14, 16, 18]
       fail-fast: false
     steps:

packages/idempotency/README.md

@@ -53,12 +53,37 @@ Next, review the IAM permissions attached to your AWS Lambda function and make s
 
 ### Function wrapper
 
-You can make any function idempotent, and safe to retry, by wrapping it using the `makeFunctionIdempotent` higher-order function.
+You can make any function idempotent, and safe to retry, by wrapping it using the `makeIdempotent` higher-order function.
 
 The function wrapper takes a reference to the function to be made idempotent as first argument, and an object with options as second argument.
 
+When you wrap your Lambda handler function, the utility uses the content of the `event` parameter to handle the idempotency logic.
+
 ```ts
-import { makeFunctionIdempotent } from '@aws-lambda-powertools/idempotency';
+import { makeIdempotent } from '@aws-lambda-powertools/idempotency';
+import { DynamoDBPersistenceLayer } from '@aws-lambda-powertools/idempotency/dynamodb';
+import type { Context, APIGatewayProxyEvent } from 'aws-lambda';
+
+const persistenceStore = new DynamoDBPersistenceLayer({
+  tableName: 'idempotencyTableName',
+});
+
+const myHandler = async (
+  event: APIGatewayProxyEvent,
+  _context: Context
+): Promise<void> => {
+  // your code goes here here
+};
+
+export const handler = makeIdempotent(myHandler, {
+  persistenceStore,
+});  
+```
+
+You can also use the `makeIdempotent` function to wrap any other arbitrary function, not just Lambda handlers.
+
+```ts
+import { makeIdempotent } from '@aws-lambda-powertools/idempotency';
 import { DynamoDBPersistenceLayer } from '@aws-lambda-powertools/idempotency/dynamodb';
 import type { Context, SQSEvent, SQSRecord } from 'aws-lambda';
 
@@ -70,20 +95,76 @@ const processingFunction = async (payload: SQSRecord): Promise<void> => {
   // your code goes here here
 };
 
+const processIdempotently = makeIdempotent(processingFunction, {
+  persistenceStore,
+});
+
 export const handler = async (
   event: SQSEvent,
   _context: Context
 ): Promise<void> => {
   for (const record of event.Records) {
-    await makeFunctionIdempotent(processingFunction, {
-      dataKeywordArgument: 'transactionId',
-      persistenceStore,
-    });
+    await processIdempotently(record);
   }
 };
 ```
 
-Note that we are specifying a `dataKeywordArgument` option, this tells the Idempotency utility which field(s) will be used as idempotency key.
+If your function has multiple arguments, you can use the `dataIndexArgument` option to specify which argument should be used as the idempotency key.
+
+```ts
+import { makeIdempotent } from '@aws-lambda-powertools/idempotency';
+import { DynamoDBPersistenceLayer } from '@aws-lambda-powertools/idempotency/dynamodb';
+import type { Context, SQSEvent, SQSRecord } from 'aws-lambda';
+
+const persistenceStore = new DynamoDBPersistenceLayer({
+  tableName: 'idempotencyTableName',
+});
+
+const processingFunction = async (payload: SQSRecord, customerId: string): Promise<void> => {
+  // your code goes here here
+};
+
+const processIdempotently = makeIdempotent(processingFunction, {
+  persistenceStore,
+  // this tells the utility to use the second argument (`customerId`) as the idempotency key
+  dataIndexArgument: 1, 
+});
+
+export const handler = async (
+  event: SQSEvent,
+  _context: Context
+): Promise<void> => {
+  for (const record of event.Records) {
+    await processIdempotently(record, 'customer-123');
+  }
+};
+```
+
+Note that you can also specify a JMESPath expression in the Idempotency config object to select a subset of the event payload as the idempotency key. This is useful when dealing with payloads that contain timestamps or request ids. 
+
+```ts
+import { makeIdempotent, IdempotencyConfig } from '@aws-lambda-powertools/idempotency';
+import { DynamoDBPersistenceLayer } from '@aws-lambda-powertools/idempotency/dynamodb';
+import type { Context, APIGatewayProxyEvent } from 'aws-lambda';
+
+const persistenceStore = new DynamoDBPersistenceLayer({
+  tableName: 'idempotencyTableName',
+});
+
+const myHandler = async (
+  event: APIGatewayProxyEvent,
+  _context: Context
+): Promise<void> => {
+  // your code goes here here
+};
+
+export const handler = makeIdempotent(myHandler, {
+  persistenceStore,
+  config: new IdempotencyConfig({
+    eventKeyJmespath: 'requestContext.identity.user',
+  }),
+});  
+```
 
 Check the [docs](https://docs.powertools.aws.dev/lambda/typescript/latest/utilities/idempotency/) for more examples.
 

packages/idempotency/src/makeIdempotent.ts

@@ -53,7 +53,7 @@ const isOptionsWithDataIndexArgument = (
  * };
  *
  * // we use wrapper to make processing function idempotent with DynamoDBPersistenceLayer
- * const processIdempotently = makeFunctionIdempotent(processRecord, {
+ * const processIdempotently = makeIdempotent(processRecord, {
  *   persistenceStore: new DynamoDBPersistenceLayer()
  *   dataKeywordArgument: 'transactionId', // keyword argument to hash the payload and the result
  * });

packages/idempotency/src/middleware/makeHandlerIdempotent.ts

@@ -1,6 +1,9 @@
 import { IdempotencyHandler } from '../IdempotencyHandler';
 import { IdempotencyConfig } from '../IdempotencyConfig';
-import { cleanupMiddlewares } from '@aws-lambda-powertools/commons/lib/middleware';
+import {
+  cleanupMiddlewares,
+  IDEMPOTENCY_KEY,
+} from '@aws-lambda-powertools/commons/lib/middleware';
 import {
   IdempotencyInconsistentStateError,
   IdempotencyItemAlreadyExistsError,
@@ -9,51 +12,94 @@ import {
 import { IdempotencyRecord } from '../persistence';
 import { MAX_RETRIES } from '../constants';
 import type { IdempotencyLambdaHandlerOptions } from '../types';
+import type { BasePersistenceLayerInterface } from '../persistence';
 import {
   MiddlewareLikeObj,
   MiddyLikeRequest,
   JSONValue,
 } from '@aws-lambda-powertools/commons';
 
+/**
+ * @internal
+ * Utility function to get the persistence store from the request internal storage
+ *
+ * @param request The Middy request object
+ * @returns The persistence store from the request internal
+ */
+const getPersistenceStoreFromRequestInternal = (
+  request: MiddyLikeRequest
+): BasePersistenceLayerInterface => {
+  const persistenceStore = request.internal[
+    `${IDEMPOTENCY_KEY}.idempotencyPersistenceStore`
+  ] as BasePersistenceLayerInterface;
+
+  return persistenceStore;
+};
+
+/**
+ * @internal
+ * Utility function to set the persistence store in the request internal storage
+ *
+ * @param request The Middy request object
+ * @param persistenceStore The persistence store to set in the request internal
+ */
+const setPersistenceStoreInRequestInternal = (
+  request: MiddyLikeRequest,
+  persistenceStore: BasePersistenceLayerInterface
+): void => {
+  request.internal[`${IDEMPOTENCY_KEY}.idempotencyPersistenceStore`] =
+    persistenceStore;
+};
+
+/**
+ * @internal
+ * Utility function to set a flag in the request internal storage to skip the idempotency middleware
+ * This is used to skip the idempotency middleware when the idempotency key is not present in the payload
+ * or when idempotency is disabled
+ *
+ * @param request The Middy request object
+ */
+const setIdempotencySkipFlag = (request: MiddyLikeRequest): void => {
+  request.internal[`${IDEMPOTENCY_KEY}.skip`] = true;
+};
+
+/**
+ * @internal
+ * Utility function to get the idempotency key from the request internal storage
+ * and determine if the request should skip the idempotency middleware
+ *
+ * @param request The Middy request object
+ * @returns Whether the idempotency middleware should be skipped
+ */
+const shouldSkipIdempotency = (request: MiddyLikeRequest): boolean => {
+  return request.internal[`${IDEMPOTENCY_KEY}.skip`] === true;
+};
+
 /**
  * A middy middleware to make your Lambda Handler idempotent.
  *
  * @example
  * ```typescript
- * import {
- *   makeHandlerIdempotent,
- *   DynamoDBPersistenceLayer,
- * } from '@aws-lambda-powertools/idempotency';
+ * import { makeHandlerIdempotent } from '@aws-lambda-powertools/idempotency/middleware';
+ * import { DynamoDBPersistenceLayer } from '@aws-lambda-powertools/idempotency/dynamodb';
  * import middy from '@middy/core';
  *
- * const dynamoDBPersistenceLayer = new DynamoDBPersistenceLayer({
- *   tableName: 'idempotencyTable',
+ * const persistenceStore = new DynamoDBPersistenceLayer({
+ *  tableName: 'idempotencyTable',
  * });
  *
- * const lambdaHandler = async (_event: unknown, _context: unknown) => {
- *   //...
- * };
- *
- * export const handler = middy(lambdaHandler)
- *  .use(makeHandlerIdempotent({ persistenceStore: dynamoDBPersistenceLayer }));
+ * export const handler = middy(
+ *   async (_event: unknown, _context: unknown): Promise<void> => {
+ *     // your code goes here
+ *   }
+ * ).use(makeHandlerIdempotent({ persistenceStore: dynamoDBPersistenceLayer }));
  * ```
  *
  * @param options - Options for the idempotency middleware
  */
 const makeHandlerIdempotent = (
   options: IdempotencyLambdaHandlerOptions
 ): MiddlewareLikeObj => {
-  const idempotencyConfig = options.config
-    ? options.config
-    : new IdempotencyConfig({});
-  const persistenceStore = options.persistenceStore;
-  persistenceStore.configure({
-    config: idempotencyConfig,
-  });
-
-  // keep the flag for after and onError checks
-  let shouldSkipIdempotency = false;
-
   /**
    * Function called before the handler is executed.
    *
@@ -76,18 +122,34 @@ const makeHandlerIdempotent = (
     request: MiddyLikeRequest,
     retryNo = 0
   ): Promise<unknown | void> => {
+    const idempotencyConfig = options.config
+      ? options.config
+      : new IdempotencyConfig({});
+    const persistenceStore = options.persistenceStore;
+    persistenceStore.configure({
+      config: idempotencyConfig,
+    });
+
     if (
+      !idempotencyConfig.isEnabled() ||
       IdempotencyHandler.shouldSkipIdempotency(
         idempotencyConfig.eventKeyJmesPath,
         idempotencyConfig.throwOnNoIdempotencyKey,
         request.event as JSONValue
       )
     ) {
       // set the flag to skip checks in after and onError
-      shouldSkipIdempotency = true;
+      setIdempotencySkipFlag(request);
 
       return;
     }
+
+    /**
+     * Store the persistence store in the request internal so that it can be
+     * used in after and onError
+     */
+    setPersistenceStoreInRequestInternal(request, persistenceStore);
+
     try {
       await persistenceStore.saveInProgress(
         request.event as JSONValue,
@@ -129,6 +191,7 @@ const makeHandlerIdempotent = (
       }
     }
   };
+
   /**
    * Function called after the handler has executed successfully.
    *
@@ -139,9 +202,10 @@ const makeHandlerIdempotent = (
    * @param request - The Middy request object
    */
   const after = async (request: MiddyLikeRequest): Promise<void> => {
-    if (shouldSkipIdempotency) {
+    if (shouldSkipIdempotency(request)) {
       return;
     }
+    const persistenceStore = getPersistenceStoreFromRequestInternal(request);
     try {
       await persistenceStore.saveSuccess(
         request.event as JSONValue,
@@ -164,9 +228,10 @@ const makeHandlerIdempotent = (
    * @param request - The Middy request object
    */
   const onError = async (request: MiddyLikeRequest): Promise<void> => {
-    if (shouldSkipIdempotency) {
+    if (shouldSkipIdempotency(request)) {
       return;
     }
+    const persistenceStore = getPersistenceStoreFromRequestInternal(request);
     try {
       await persistenceStore.deleteRecord(request.event as JSONValue);
     } catch (error) {
@@ -177,19 +242,11 @@ const makeHandlerIdempotent = (
     }
   };
 
-  if (idempotencyConfig.isEnabled()) {
-    return {
-      before,
-      after,
-      onError,
-    };
-  } else {
-    return {
-      before: () => {
-        return undefined;
-      },
-    };
-  }
+  return {
+    before,
+    after,
+    onError,
+  };
 };
 
 export { makeHandlerIdempotent };

packages/idempotency/src/persistence/BasePersistenceLayerInterface.ts

@@ -1,10 +1,11 @@
 import { IdempotencyRecord } from './IdempotencyRecord';
 import type { BasePersistenceLayerOptions } from '../types/BasePersistenceLayer';
 
+// TODO: move this to types folder
 interface BasePersistenceLayerInterface {
   configure(options?: BasePersistenceLayerOptions): void;
   isPayloadValidationEnabled(): boolean;
-  saveInProgress(data: unknown): Promise<void>;
+  saveInProgress(data: unknown, remainingTimeInMillis?: number): Promise<void>;
   saveSuccess(data: unknown, result: unknown): Promise<void>;
   deleteRecord(data: unknown): Promise<void>;
   getRecord(data: unknown): Promise<IdempotencyRecord>;