feat(idempotency): add support for custom key prefix (#3532)

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

Author: shdq

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

@@ -51,6 +51,10 @@ export class IdempotencyHandler<Func extends AnyFunction> {
    * Idempotency configuration options.
    */
   readonly #idempotencyConfig: IdempotencyConfig;
+  /**
+   * Custom prefix to be used when generating the idempotency key.
+   */
+  readonly #keyPrefix: string | undefined;
   /**
    * Persistence layer used to store the idempotency records.
    */
@@ -69,18 +73,21 @@ export class IdempotencyHandler<Func extends AnyFunction> {
       idempotencyConfig,
       functionArguments,
       persistenceStore,
+      keyPrefix,
       thisArg,
     } = options;
     this.#functionToMakeIdempotent = functionToMakeIdempotent;
     this.#functionPayloadToBeHashed = functionPayloadToBeHashed;
     this.#idempotencyConfig = idempotencyConfig;
+    this.#keyPrefix = keyPrefix;
     this.#functionArguments = functionArguments;
     this.#thisArg = thisArg;
 
     this.#persistenceStore = persistenceStore;
 
     this.#persistenceStore.configure({
       config: this.#idempotencyConfig,
+      keyPrefix: this.#keyPrefix,
     });
   }
 

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

@@ -79,7 +79,7 @@ function makeIdempotent<Func extends AnyFunction>(
   fn: Func,
   options: ItempotentFunctionOptions<Parameters<Func>>
 ): (...args: Parameters<Func>) => ReturnType<Func> {
-  const { persistenceStore, config } = options;
+  const { persistenceStore, config, keyPrefix } = options;
   const idempotencyConfig = config ? config : new IdempotencyConfig({});
 
   if (!idempotencyConfig.isEnabled()) return fn;
@@ -102,6 +102,7 @@ function makeIdempotent<Func extends AnyFunction>(
       functionToMakeIdempotent: fn,
       idempotencyConfig: idempotencyConfig,
       persistenceStore: persistenceStore,
+      keyPrefix: keyPrefix,
       functionArguments: args,
       functionPayloadToBeHashed,
       thisArg: this,

Diff for: packages/idempotency/src/middleware/makeHandlerIdempotent.ts

@@ -117,15 +117,18 @@ const makeHandlerIdempotent = (
       ? options.config
       : new IdempotencyConfig({});
     const persistenceStore = options.persistenceStore;
+    const keyPrefix = options.keyPrefix;
     persistenceStore.configure({
       config: idempotencyConfig,
+      keyPrefix: keyPrefix,
     });
 
     const idempotencyHandler = new IdempotencyHandler({
       functionToMakeIdempotent: /* v8 ignore next */ () => ({}),
       functionArguments: [],
       idempotencyConfig,
       persistenceStore,
+      keyPrefix,
       functionPayloadToBeHashed: undefined,
     });
     setIdempotencyHandlerInRequestInternal(request, idempotencyHandler);

Diff for: packages/idempotency/src/persistence/BasePersistenceLayer.ts

@@ -46,14 +46,15 @@ abstract class BasePersistenceLayer implements BasePersistenceLayerInterface {
   /**
    * Initialize the base persistence layer from the configuration settings
    *
-   * @param {BasePersistenceLayerConfigureOptions} config - configuration object for the persistence layer
+   * @param {BasePersistenceLayerConfigureOptions} options - configuration object for the persistence layer
    */
-  public configure(config: BasePersistenceLayerOptions): void {
-    // Extracting the idempotency config from the config object for easier access
-    const { config: idempotencyConfig } = config;
+  public configure(options: BasePersistenceLayerOptions): void {
+    const { config: idempotencyConfig, keyPrefix, functionName } = options;
 
-    if (config?.functionName && config.functionName.trim() !== '') {
-      this.idempotencyKeyPrefix = `${this.idempotencyKeyPrefix}.${config.functionName}`;
+    if (keyPrefix?.trim()) {
+      this.idempotencyKeyPrefix = keyPrefix.trim();
+    } else if (functionName?.trim()) {
+      this.idempotencyKeyPrefix = `${this.idempotencyKeyPrefix}.${functionName.trim()}`;
     }
 
     // Prevent reconfiguration

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

@@ -4,6 +4,7 @@ import type { IdempotencyRecord } from '../persistence/IdempotencyRecord.js';
 type BasePersistenceLayerOptions = {
   config: IdempotencyConfig;
   functionName?: string;
+  keyPrefix?: string;
 };
 
 interface BasePersistenceLayerInterface {

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

@@ -19,6 +19,7 @@ import type { IdempotencyRecord } from '../persistence/IdempotencyRecord.js';
 type IdempotencyLambdaHandlerOptions = {
   persistenceStore: BasePersistenceLayer;
   config?: IdempotencyConfig;
+  keyPrefix?: string;
 };
 
 /**
@@ -137,6 +138,10 @@ type IdempotencyHandlerOptions = {
    * Idempotency configuration options.
    */
   idempotencyConfig: IdempotencyConfig;
+  /**
+   * The custom idempotency key prefix.
+   */
+  keyPrefix?: string;
   /**
    * Persistence layer used to store the idempotency records.
    */

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

@@ -1,9 +1,10 @@
 import type { LambdaInterface } from '@aws-lambda-powertools/commons/types';
 import context from '@aws-lambda-powertools/testing-utils/context';
 import type { Context } from 'aws-lambda';
-import { describe, expect, it } from 'vitest';
-import { idempotent } from '../../src/index.js';
+import { describe, expect, it, vi } from 'vitest';
+import { idempotent, IdempotencyConfig } from '../../src/index.js';
 import { PersistenceLayerTestClass } from '../helpers/idempotencyUtils.js';
+import { BasePersistenceLayer } from '../../src/persistence/BasePersistenceLayer.js';
 
 describe('Given a class with a function to decorate', () => {
   it('maintains the scope of the decorated function', async () => {
@@ -35,4 +36,41 @@ describe('Given a class with a function to decorate', () => {
     // Assess
     expect(result).toBe('private foo');
   });
+
+  it('passes the custom keyPrefix to the persistenceStore', async () => {
+    // Prepare
+    const configureSpy = vi.spyOn(BasePersistenceLayer.prototype, 'configure');
+    const idempotencyConfig = new IdempotencyConfig({});
+
+    class TestClass implements LambdaInterface {
+      @idempotent({
+        persistenceStore: new PersistenceLayerTestClass(),
+        config: idempotencyConfig,
+        keyPrefix: 'my-custom-prefix',
+      })
+      public async handler(
+        _event: unknown,
+        _context: Context
+      ): Promise<boolean> {
+        return true;
+      }
+    }
+
+    const handlerClass = new TestClass();
+    const handler = handlerClass.handler.bind(handlerClass);
+
+    // Act
+    const result = await handler({}, context);
+
+    // Assess
+    expect(result).toBeTruthy();
+
+    expect(configureSpy).toHaveBeenCalled();
+    const configureCallArgs = configureSpy.mock.calls[0][0]; // Extract first call's arguments
+    expect(configureCallArgs.config).toBe(idempotencyConfig);
+    expect(configureCallArgs.keyPrefix).toBe('my-custom-prefix');
+
+    // Restore the spy
+    configureSpy.mockRestore();
+  });
 });

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

@@ -393,6 +393,45 @@ describe('Function: makeIdempotent', () => {
       expect(saveSuccessSpy).toHaveBeenCalledTimes(0);
     }
   );
+
+  it.each([
+    {
+      type: 'wrapper',
+    },
+    { type: 'middleware' },
+  ])(
+    'passes keyPrefix correctly in idempotency handler ($type)',
+    async ({ type }) => {
+      // Prepare
+      const keyPrefix = 'my-custom-prefix';
+      const options = {
+        ...mockIdempotencyOptions,
+        keyPrefix,
+        config: new IdempotencyConfig({
+          eventKeyJmesPath: 'idempotencyKey',
+        }),
+      };
+      const handler =
+        type === 'wrapper'
+          ? makeIdempotent(fnSuccessfull, options)
+          : middy(fnSuccessfull).use(makeHandlerIdempotent(options));
+
+      const configureSpy = vi.spyOn(
+        mockIdempotencyOptions.persistenceStore,
+        'configure'
+      );
+    
+      // Act
+      const result = await handler(event, context);
+
+      // Assess
+      expect(result).toBe(true);
+      expect(configureSpy).toHaveBeenCalledWith(
+        expect.objectContaining({ keyPrefix })
+      );
+    }
+  );
+
   it('uses the first argument when when wrapping an arbitrary function', async () => {
     // Prepare
     const config = new IdempotencyConfig({});

Diff for: packages/idempotency/tests/unit/persistence/BasePersistenceLayer.test.ts

@@ -80,6 +80,34 @@ describe('Class: BasePersistenceLayer', () => {
       );
     });
 
+    it('trims the function name before appending as key prefix', () => {
+      // Prepare
+      const config = new IdempotencyConfig({});
+      const persistenceLayer = new PersistenceLayerTestClass();
+
+      // Act
+      persistenceLayer.configure({ config, functionName: ' my-function ' });
+
+      // Assess
+      expect(persistenceLayer.idempotencyKeyPrefix).toBe(
+        'my-lambda-function.my-function'
+      );
+    });
+
+    it('appends custom prefix to the idempotence key prefix', () => {
+      // Prepare
+      const config = new IdempotencyConfig({});
+      const persistenceLayer = new PersistenceLayerTestClass();
+
+      // Act
+      persistenceLayer.configure({ config, keyPrefix: 'my-custom-prefix' });
+
+      // Assess
+      expect(persistenceLayer.idempotencyKeyPrefix).toBe(
+        'my-custom-prefix'
+      );
+    });
+
     it('uses default config when no option is provided', () => {
       // Prepare
       const config = new IdempotencyConfig({});