feat(idempotency): ability to specify JMESPath custom functions (#3150)

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

Author: arnabrahman

Diff for: docs/utilities/idempotency.md

@@ -573,6 +573,7 @@ Idempotent decorator can be further configured with **`IdempotencyConfig`** as s
 | ----------------------------- | ----------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ |
 | **eventKeyJmespath**          | `''`        | JMESPath expression to extract the idempotency key from the event record using [built-in functions](./jmespath.md#built-in-jmespath-functions){target="_blank"}                                                                            |
 | **payloadValidationJmespath** | `''`        | JMESPath expression to validate that the specified fields haven't changed across requests for the same idempotency key _e.g., payload tampering._                                                                                          |
+| **jmesPathOptions**           | `undefined` | Custom JMESPath functions to use when parsing the JMESPath expressions. See [Custom JMESPath Functions](idempotency.md#custom-jmespath-functions)                                                                                           |
 | **throwOnNoIdempotencyKey**   | `false`     | Throw an error if no idempotency key was found in the request                                                                                                                                                                              |
 | **expiresAfterSeconds**       | 3600        | The number of seconds to wait before a record is expired, allowing a new transaction with the same idempotency key                                                                                                                         |
 | **useLocalCache**             | `false`     | Whether to cache idempotency results in-memory to save on persistence storage latency and costs                                                                                                                                            |
@@ -657,6 +658,16 @@ Without payload validation, we would have returned the same result as we did for
 
 By using **`payloadValidationJmesPath="amount"`**, we prevent this potentially confusing behavior and instead throw an error.
 
+### Custom JMESPath Functions
+
+You can provide custom JMESPath functions for evaluating JMESPath expressions by passing them through the **`jmesPathOptions`** parameter. In this example, we use a custom function, `my_fancy_function`, to parse the payload as a JSON object instead of a string.
+
+=== "Custom JMESPath functions"
+
+    ```typescript hl_lines="16 20 28-29"
+    --8<-- "examples/snippets/idempotency/workingWithCustomJmesPathFunctions.ts"
+    ```
+
 ### Making idempotency key required
 
 If you want to enforce that an idempotency key is required, you can set **`throwOnNoIdempotencyKey`** to `true`.

Diff for: examples/snippets/idempotency/workingWithCustomJmesPathFunctions.ts

@@ -0,0 +1,31 @@
+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 {
+  Functions,
+  PowertoolsFunctions,
+} from '@aws-lambda-powertools/jmespath/functions';
+
+const persistenceStore = new DynamoDBPersistenceLayer({
+  tableName: 'idempotencyTableName',
+});
+
+class MyFancyFunctions extends PowertoolsFunctions {
+  @Functions.signature({
+    argumentsSpecs: [['string']],
+  })
+  public funcMyFancyFunction(value: string): JSONValue {
+    return JSON.parse(value);
+  }
+}
+
+export const handler = makeIdempotent(async () => true, {
+  persistenceStore,
+  config: new IdempotencyConfig({
+    eventKeyJmesPath: 'my_fancy_function(body).["user", "productId"]',
+    jmesPathOptions: new MyFancyFunctions(),
+  }),
+});

Diff for: packages/idempotency/src/IdempotencyConfig.ts

@@ -70,7 +70,9 @@ class IdempotencyConfig {
   public constructor(config: IdempotencyConfigOptions) {
     this.eventKeyJmesPath = config.eventKeyJmesPath ?? '';
     this.payloadValidationJmesPath = config.payloadValidationJmesPath;
-    this.jmesPathOptions = { customFunctions: new PowertoolsFunctions() };
+    this.jmesPathOptions = {
+      customFunctions: config.jmesPathOptions ?? new PowertoolsFunctions(),
+    };
     this.throwOnNoIdempotencyKey = config.throwOnNoIdempotencyKey ?? false;
     this.expiresAfterSeconds = config.expiresAfterSeconds ?? 3600; // 1 hour default
     this.useLocalCache = config.useLocalCache ?? false;

Diff for: packages/idempotency/src/types/IdempotencyOptions.ts

@@ -1,4 +1,5 @@
 import type { JSONValue } from '@aws-lambda-powertools/commons/types';
+import type { Functions } from '@aws-lambda-powertools/jmespath/functions';
 import type { Context, Handler } from 'aws-lambda';
 import type { IdempotencyConfig } from '../IdempotencyConfig.js';
 import type { BasePersistenceLayer } from '../persistence/BasePersistenceLayer.js';
@@ -168,6 +169,10 @@ type IdempotencyConfigOptions = {
    * An optional JMESPath expression to extract the payload to be validated from the event record
    */
   payloadValidationJmesPath?: string;
+  /**
+   * Custom JMESPath functions to use when parsing the JMESPath expressions
+   */
+  jmesPathOptions?: Functions;
   /**
    * Throw an error if no idempotency key was found in the request, defaults to `false`
    */

Diff for: packages/idempotency/tests/unit/IdempotencyConfig.test.ts

@@ -1,3 +1,8 @@
+import type { JSONValue } from '@aws-lambda-powertools/commons/types';
+import {
+  Functions,
+  PowertoolsFunctions,
+} from '@aws-lambda-powertools/jmespath/functions';
 import context from '@aws-lambda-powertools/testing-utils/context';
 import { afterAll, beforeEach, describe, expect, it } from 'vitest';
 import { IdempotencyConfig } from '../../src/index.js';
@@ -32,12 +37,23 @@ describe('Class: IdempotencyConfig', () => {
           useLocalCache: false,
           hashFunction: 'md5',
           lambdaContext: undefined,
+          jmesPathOptions: expect.objectContaining({
+            customFunctions: expect.any(PowertoolsFunctions),
+          }),
         })
       );
     });
 
     it('initializes the config with the provided configs', () => {
       // Prepare
+      class MyFancyFunctions extends Functions {
+        @Functions.signature({
+          argumentsSpecs: [['string']],
+        })
+        public funcMyFancyFunction(value: string): JSONValue {
+          return JSON.parse(value);
+        }
+      }
       const configOptions: IdempotencyConfigOptions = {
         eventKeyJmesPath: 'eventKeyJmesPath',
         payloadValidationJmesPath: 'payloadValidationJmesPath',
@@ -46,6 +62,7 @@ describe('Class: IdempotencyConfig', () => {
         useLocalCache: true,
         hashFunction: 'hashFunction',
         lambdaContext: context,
+        jmesPathOptions: new MyFancyFunctions(),
       };
 
       // Act
@@ -61,6 +78,9 @@ describe('Class: IdempotencyConfig', () => {
           useLocalCache: true,
           hashFunction: 'hashFunction',
           lambdaContext: context,
+          jmesPathOptions: expect.objectContaining({
+            customFunctions: expect.any(MyFancyFunctions),
+          }),
         })
       );
     });

Diff for: packages/jmespath/src/PowertoolsFunctions.ts

@@ -58,4 +58,4 @@ class PowertoolsFunctions extends Functions {
   }
 }
 
-export { PowertoolsFunctions };
+export { PowertoolsFunctions, Functions };