feat(idempotency): add custom JMESPath functions (#2364)

* feat(idempotency): add custom jmespath functions * docs(idempotency): add custom functions to docs * chore(layers): add package to layer * chore: move jmespath pkg up into build process * feat: enable custom functions * chore: move jmespath pkg up into build process * chore: move jmespath pkg up into build process * chore: update layer setup * refactor: moved Powertools function initialization * chore: added pkg to pre-push hook --------- Co-authored-by: Alexander Schueren <[email protected]>
aws-powertools · Apr 16, 2024 · 9721e7c · 9721e7c
1 parent fc2d709
commit 9721e7c
Show file tree

Hide file tree

Showing 13 changed files with 74 additions and 36 deletions.
diff --git a/.github/actions/cached-node-modules/action.yml b/.github/actions/cached-node-modules/action.yml
@@ -39,12 +39,12 @@ runs:
       # sequence, but still in the correct order.
       run: |
         npm run build -w packages/commons
+        npm run build -w packages/jmespath
         npm run build -w packages/logger & \
         npm run build -w packages/tracer & \
         npm run build -w packages/metrics & \
         npm run build -w packages/parameters & \
         npm run build -w packages/idempotency & \
         npm run build -w packages/batch & \
-        npm run build -w packages/testing & \
-        npm run build -w packages/jmespath
+        npm run build -w packages/testing 
       shell: bash
diff --git a/.husky/pre-push b/.husky/pre-push
@@ -1,5 +1,6 @@
 npm t \
   -w packages/commons \
+  -w packages/jmespath \
   -w packages/logger \
   -w packages/metrics \
   -w packages/tracer \

diff --git a/docs/snippets/idempotency/makeIdempotentJmes.ts b/docs/snippets/idempotency/makeIdempotentJmes.ts
@@ -22,9 +22,9 @@ const createSubscriptionPayment = async (
   };
 };
 
-// Extract the idempotency key from the request headers
+// Deserialize JSON string under the "body" key, then extract the "user" and "productId" keys
 const config = new IdempotencyConfig({
-  eventKeyJmesPath: 'body',
+  eventKeyJmesPath: 'powertools_json(body).["user", "productId"]',
 });
 
 export const handler = makeIdempotent(

diff --git a/docs/utilities/idempotency.md b/docs/utilities/idempotency.md
@@ -220,9 +220,11 @@ Imagine the function executes successfully, but the client never receives the re
 ???+ warning "Deserializing JSON strings in payloads for increased accuracy."
     The payload extracted by the `eventKeyJmesPath` is treated as a string by default. This means there could be differences in whitespace even when the JSON payload itself is identical.
 
+    To alter this behaviour, we can use the [JMESPath built-in function `powertools_json()`](jmespath.md#powertools_json-function) to treat the payload as a JSON object rather than a string.
+
 === "index.ts"
 
-    ```typescript hl_lines="4 26-28 49"
+    ```typescript hl_lines="4 27 49"
     --8<-- "docs/snippets/idempotency/makeIdempotentJmes.ts"
     ```
 

diff --git a/layers/src/layer-publisher-stack.ts b/layers/src/layer-publisher-stack.ts
@@ -62,6 +62,7 @@ export class LayerPublisherStack extends Stack {
               // the name is the same as the npm workspace name
               const utilities = [
                 'commons',
+                'jmespath',
                 'logger',
                 'metrics',
                 'tracer',
@@ -87,8 +88,6 @@ export class LayerPublisherStack extends Stack {
                 'node_modules/async-hook-jl/test',
                 'node_modules/stack-chain/test',
                 'node_modules/shimmer/test',
-                'node_modules/jmespath/artifacts',
-                'node_modules/jmespath/bower.json',
                 'node_modules/obliterator/*.d.ts',
                 'node_modules/strnum/.vscode',
                 'node_modules/strnum/*.test.js',

diff --git a/package-lock.json b/package-lock.json
diff --git a/packages/idempotency/README.md b/packages/idempotency/README.md
@@ -18,7 +18,6 @@ You can use the package in both TypeScript and JavaScript code bases.
   - [Becoming a reference customer](#becoming-a-reference-customer)
   - [Sharing your work](#sharing-your-work)
   - [Using Lambda Layer](#using-lambda-layer)
-- [Credits](#credits)
 - [License](#license)
 
 ## Intro
@@ -158,7 +157,33 @@ export const handler = makeIdempotent(myHandler, {
   config: new IdempotencyConfig({
     eventKeyJmespath: 'requestContext.identity.user',
   }),
-});  
+});
+```
+
+Additionally, you can also use one of the [JMESPath built-in functions](https://docs.powertools.aws.dev/lambda/typescript/latest/utilities/jmespath/#built-in-jmespath-functions) like `powertools_json()` to decode keys and use parts of the payload as the idempotency key.
+
+```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: 'powertools_json(body).["user", "productId"]',
+  }),
+});
 ```
 
 Check the [docs](https://docs.powertools.aws.dev/lambda/typescript/latest/utilities/idempotency/) for more examples.
@@ -311,12 +336,8 @@ Share what you did with Powertools for AWS Lambda (TypeScript) 💞💞. Blog po
 
 ### Using Lambda Layer
 
-This helps us understand who uses Powertools for AWS Lambda (TypeScript) in a non-intrusive way, and helps us gain future investments for other Powertools for AWS Lambda languages. When [using Layers](#lambda-layers), you can add Powertools as a dev dependency (or as part of your virtual env) to not impact the development process.
-
-## Credits
-
-Credits for the Lambda Powertools for AWS Lambda (TypeScript) idea go to [DAZN](https://github.com/getndazn) and their [DAZN Lambda Powertools](https://github.com/getndazn/dazn-lambda-powertools/).
+This helps us understand who uses Powertools for AWS Lambda (TypeScript) in a non-intrusive way, and helps us gain future investments for other Powertools for AWS Lambda languages. When [using Layers](https://docs.powertools.aws.dev/lambda/typescript/latest/#lambda-layer), you can add Powertools as a dev dependency to not impact the development process.
 
 ## License
 
-This library is licensed under the MIT-0 License. See the LICENSE file.
+This library is licensed under the MIT-0 License. See the LICENSE file.
diff --git a/packages/idempotency/package.json b/packages/idempotency/package.json
@@ -101,7 +101,7 @@
   },
   "dependencies": {
     "@aws-lambda-powertools/commons": "^2.0.4",
-    "jmespath": "^0.16.0"
+    "@aws-lambda-powertools/jmespath": "^2.0.4"
   },
   "peerDependencies": {
     "@aws-sdk/client-dynamodb": ">=3.x",
@@ -131,7 +131,6 @@
     "@aws-lambda-powertools/testing-utils": "file:../testing",
     "@aws-sdk/client-dynamodb": "^3.554.0",
     "@aws-sdk/lib-dynamodb": "^3.554.0",
-    "@types/jmespath": "^0.15.0",
     "aws-sdk-client-mock": "^4.0.0",
     "aws-sdk-client-mock-jest": "^4.0.0"
   }

diff --git a/packages/idempotency/src/IdempotencyConfig.ts b/packages/idempotency/src/IdempotencyConfig.ts
@@ -1,6 +1,8 @@
 import { EnvironmentVariablesService } from './config/EnvironmentVariablesService.js';
 import type { Context } from 'aws-lambda';
 import type { IdempotencyConfigOptions } from './types/IdempotencyOptions.js';
+import type { ParsingOptions } from '@aws-lambda-powertools/jmespath/types';
+import { PowertoolsFunctions } from '@aws-lambda-powertools/jmespath/functions';
 
 /**
  * Configuration for the idempotency feature.
@@ -22,6 +24,10 @@ class IdempotencyConfig {
    * @default 'md5'
    */
   public hashFunction: string;
+  /**
+   *
+   */
+  public jmesPathOptions: ParsingOptions;
   /**
    * The lambda context object.
    */
@@ -53,6 +59,7 @@ class IdempotencyConfig {
   public constructor(config: IdempotencyConfigOptions) {
     this.eventKeyJmesPath = config.eventKeyJmesPath ?? '';
     this.payloadValidationJmesPath = config.payloadValidationJmesPath;
+    this.jmesPathOptions = { customFunctions: new PowertoolsFunctions() };
     this.throwOnNoIdempotencyKey = config.throwOnNoIdempotencyKey ?? false;
     this.expiresAfterSeconds = config.expiresAfterSeconds ?? 3600; // 1 hour default
     this.useLocalCache = config.useLocalCache ?? false;

diff --git a/packages/idempotency/src/IdempotencyHandler.ts b/packages/idempotency/src/IdempotencyHandler.ts
@@ -16,7 +16,7 @@ import { BasePersistenceLayer } from './persistence/BasePersistenceLayer.js';
 import { IdempotencyRecord } from './persistence/IdempotencyRecord.js';
 import { IdempotencyConfig } from './IdempotencyConfig.js';
 import { MAX_RETRIES, IdempotencyRecordStatus } from './constants.js';
-import { search } from 'jmespath';
+import { search } from '@aws-lambda-powertools/jmespath';
 
 /**
  * @internal
@@ -275,8 +275,9 @@ export class IdempotencyHandler<Func extends AnyFunction> {
       !this.#idempotencyConfig.throwOnNoIdempotencyKey
     ) {
       const selection = search(
+        this.#idempotencyConfig.eventKeyJmesPath,
         this.#functionPayloadToBeHashed,
-        this.#idempotencyConfig.eventKeyJmesPath
+        this.#idempotencyConfig.jmesPathOptions
       );
 
       return selection === undefined || selection === null;

diff --git a/packages/idempotency/src/persistence/BasePersistenceLayer.ts b/packages/idempotency/src/persistence/BasePersistenceLayer.ts
@@ -1,5 +1,6 @@
 import { createHash, Hash } from 'node:crypto';
-import { search } from 'jmespath';
+import { search } from '@aws-lambda-powertools/jmespath';
+import type { ParsingOptions } from '@aws-lambda-powertools/jmespath/types';
 import type {
   BasePersistenceLayerOptions,
   BasePersistenceLayerInterface,
@@ -36,6 +37,7 @@ abstract class BasePersistenceLayer implements BasePersistenceLayerInterface {
   private throwOnNoIdempotencyKey = false;
   private useLocalCache = false;
   private validationKeyJmesPath?: string;
+  #jmesPathOptions?: ParsingOptions;
 
   public constructor() {
     this.envVarsService = new EnvironmentVariablesService();
@@ -63,6 +65,7 @@ abstract class BasePersistenceLayer implements BasePersistenceLayerInterface {
 
     this.eventKeyJmesPath = idempotencyConfig?.eventKeyJmesPath;
     this.validationKeyJmesPath = idempotencyConfig?.payloadValidationJmesPath;
+    this.#jmesPathOptions = idempotencyConfig.jmesPathOptions;
     this.payloadValidationEnabled =
       this.validationKeyJmesPath !== undefined || false;
     this.throwOnNoIdempotencyKey =
@@ -279,7 +282,11 @@ abstract class BasePersistenceLayer implements BasePersistenceLayerInterface {
    */
   private getHashedIdempotencyKey(data: JSONValue): string {
     if (this.eventKeyJmesPath) {
-      data = search(data, this.eventKeyJmesPath);
+      data = search(
+        this.eventKeyJmesPath,
+        data,
+        this.#jmesPathOptions
+      ) as JSONValue;
     }
 
     if (BasePersistenceLayer.isMissingIdempotencyKey(data)) {
@@ -305,7 +312,11 @@ abstract class BasePersistenceLayer implements BasePersistenceLayerInterface {
    */
   private getHashedPayload(data: JSONValue): string {
     if (this.isPayloadValidationEnabled() && this.validationKeyJmesPath) {
-      data = search(data, this.validationKeyJmesPath);
+      data = search(
+        this.validationKeyJmesPath,
+        data,
+        this.#jmesPathOptions
+      ) as JSONValue;
 
       return this.generateHash(JSON.stringify(data));
     } else {

diff --git a/packages/idempotency/tests/e2e/makeIdempotent.test.FunctionCode.ts b/packages/idempotency/tests/e2e/makeIdempotent.test.FunctionCode.ts
@@ -90,16 +90,17 @@ export const handlerCustomized = async (
  * Test idempotent Lambda handler with JMESPath expression to extract event key.
  */
 export const handlerLambda = makeIdempotent(
-  async (event: { foo: string }, context: Context) => {
+  async (event: { body: string }, context: Context) => {
     logger.addContext(context);
-    logger.info(`foo`, { details: event.foo });
+    const body = JSON.parse(event.body);
+    logger.info('foo', { details: body.foo });
 
-    return event.foo;
+    return body.foo;
   },
   {
     persistenceStore: dynamoDBPersistenceLayer,
     config: new IdempotencyConfig({
-      eventKeyJmesPath: 'foo',
+      eventKeyJmesPath: 'powertools_json(body).foo',
       useLocalCache: true,
     }),
   }

diff --git a/packages/idempotency/tests/e2e/makeIdempotent.test.ts b/packages/idempotency/tests/e2e/makeIdempotent.test.ts
@@ -268,10 +268,12 @@ describe(`Idempotency E2E tests, wrapper function usage`, () => {
     async () => {
       // Prepare
       const payload = {
-        foo: 'bar',
+        body: JSON.stringify({
+          foo: 'bar',
+        }),
       };
       const payloadHash = createHash('md5')
-        .update(JSON.stringify(payload.foo))
+        .update(JSON.stringify('bar'))
         .digest('base64');
 
       // Act