fix(idempotency): preserve scope of decorated class

Author: dreamorosi

packages/idempotency/src/IdempotencyHandler.ts

@@ -1,3 +1,4 @@
+import type { Handler } from 'aws-lambda';
 import type {
   JSONValue,
   MiddyLikeRequest,
@@ -53,6 +54,12 @@ export class IdempotencyHandler<Func extends AnyFunction> {
    * Persistence layer used to store the idempotency records.
    */
   readonly #persistenceStore: BasePersistenceLayer;
+  /**
+   * The `this` context to be used when calling the function.
+   *
+   * When decorating a class method, this will be the instance of the class.
+   */
+  readonly #thisArg?: Handler;
 
   public constructor(options: IdempotencyHandlerOptions) {
     const {
@@ -61,11 +68,13 @@ export class IdempotencyHandler<Func extends AnyFunction> {
       idempotencyConfig,
       functionArguments,
       persistenceStore,
+      thisArg,
     } = options;
     this.#functionToMakeIdempotent = functionToMakeIdempotent;
     this.#functionPayloadToBeHashed = functionPayloadToBeHashed;
     this.#idempotencyConfig = idempotencyConfig;
     this.#functionArguments = functionArguments;
+    this.#thisArg = thisArg;
 
     this.#persistenceStore = persistenceStore;
 
@@ -121,7 +130,10 @@ export class IdempotencyHandler<Func extends AnyFunction> {
   public async getFunctionResult(): Promise<ReturnType<Func>> {
     let result;
     try {
-      result = await this.#functionToMakeIdempotent(...this.#functionArguments);
+      result = await this.#functionToMakeIdempotent.apply(
+        this.#thisArg,
+        this.#functionArguments
+      );
     } catch (error) {
       await this.#deleteInProgressRecord();
       throw error;
@@ -149,7 +161,10 @@ export class IdempotencyHandler<Func extends AnyFunction> {
   public async handle(): Promise<ReturnType<Func>> {
     // early return if we should skip idempotency completely
     if (this.shouldSkipIdempotency()) {
-      return await this.#functionToMakeIdempotent(...this.#functionArguments);
+      return await this.#functionToMakeIdempotent.apply(
+        this.#thisArg,
+        this.#functionArguments
+      );
     }
 
     let e;

packages/idempotency/src/idempotencyDecorator.ts

@@ -1,3 +1,4 @@
+import type { Handler } from 'aws-lambda';
 import {
   AnyFunction,
   ItempotentFunctionOptions,
@@ -65,7 +66,10 @@ const idempotent = function (
     descriptor: PropertyDescriptor
   ) {
     const childFunction = descriptor.value;
-    descriptor.value = makeIdempotent(childFunction, options);
+
+    descriptor.value = async function (this: Handler, ...args: unknown[]) {
+      return makeIdempotent(childFunction, options).bind(this)(...args);
+    };
 
     return descriptor;
   };

packages/idempotency/src/makeIdempotent.ts

@@ -72,9 +72,10 @@ const isOptionsWithDataIndexArgument = (
  *
  * ```
  */
-const makeIdempotent = <Func extends AnyFunction>(
+/* const makeIdempotent = <Func extends AnyFunction>(
   fn: Func,
-  options: ItempotentFunctionOptions<Parameters<Func>>
+  options: ItempotentFunctionOptions<Parameters<Func>>,
+  thisArg: Handler
 ): ((...args: Parameters<Func>) => ReturnType<Func>) => {
   const { persistenceStore, config } = options;
   const idempotencyConfig = config ? config : new IdempotencyConfig({});
@@ -101,8 +102,44 @@ const makeIdempotent = <Func extends AnyFunction>(
       persistenceStore: persistenceStore,
       functionArguments: args,
       functionPayloadToBeHashed,
+      thisArg,
     }).handle() as ReturnType<Func>;
   };
-};
+}; */
+// eslint-disable-next-line func-style
+function makeIdempotent<Func extends AnyFunction>(
+  fn: Func,
+  options: ItempotentFunctionOptions<Parameters<Func>>
+): (...args: Parameters<Func>) => ReturnType<Func> {
+  const { persistenceStore, config } = options;
+  const idempotencyConfig = config ? config : new IdempotencyConfig({});
+
+  if (!idempotencyConfig.isEnabled()) return fn;
+
+  return function (...args: Parameters<Func>): ReturnType<Func> {
+    let functionPayloadToBeHashed;
+
+    if (isFnHandler(fn, args)) {
+      idempotencyConfig.registerLambdaContext(args[1]);
+      functionPayloadToBeHashed = args[0];
+    } else {
+      if (isOptionsWithDataIndexArgument(options)) {
+        functionPayloadToBeHashed = args[options.dataIndexArgument];
+      } else {
+        functionPayloadToBeHashed = args[0];
+      }
+    }
+
+    return new IdempotencyHandler({
+      functionToMakeIdempotent: fn,
+      idempotencyConfig: idempotencyConfig,
+      persistenceStore: persistenceStore,
+      functionArguments: args,
+      functionPayloadToBeHashed,
+      // @ts-expect-error abc
+      thisArg: this,
+    }).handle() as ReturnType<Func>;
+  };
+}
 
 export { makeIdempotent };

packages/idempotency/src/types/IdempotencyOptions.ts

@@ -1,4 +1,4 @@
-import type { Context } from 'aws-lambda';
+import type { Context, Handler } from 'aws-lambda';
 import type { BasePersistenceLayer } from '../persistence/BasePersistenceLayer.js';
 import type { IdempotencyConfig } from '../IdempotencyConfig.js';
 import type { JSONValue } from '@aws-lambda-powertools/commons/types';
@@ -139,6 +139,12 @@ type IdempotencyHandlerOptions = {
    * Persistence layer used to store the idempotency records.
    */
   persistenceStore: BasePersistenceLayer;
+  /**
+   * The `this` context to be used when calling the function.
+   *
+   * When decorating a class method, this will be the instance of the class.
+   */
+  thisArg?: Handler;
 };
 
 /**

packages/idempotency/tests/unit/idempotencyDecorator.test.ts

@@ -19,6 +19,7 @@ import type { IdempotencyRecordOptions } from '../../src/types/index.js';
 import { Context } from 'aws-lambda';
 import context from '@aws-lambda-powertools/testing-utils/context';
 import { IdempotencyRecordStatus } from '../../src/constants.js';
+import type { LambdaInterface } from '@aws-lambda-powertools/commons/types';
 
 const mockSaveInProgress = jest
   .spyOn(BasePersistenceLayer.prototype, 'saveInProgress')
@@ -86,7 +87,10 @@ describe('Given a class with a function to decorate', (classWithLambdaHandler =
     testingKey: keyValueToBeSaved,
     otherKey: 'thisWillNot',
   };
-  beforeEach(() => jest.clearAllMocks());
+  beforeEach(() => {
+    jest.clearAllMocks();
+    jest.resetAllMocks();
+  });
 
   describe('When wrapping a function with no previous executions', () => {
     beforeEach(async () => {
@@ -313,4 +317,34 @@ describe('Given a class with a function to decorate', (classWithLambdaHandler =
       delete process.env.POWERTOOLS_IDEMPOTENCY_DISABLED;
     });
   });
+
+  it('maintains the scope of the decorated function', async () => {
+    // Prepare
+    class TestClass implements LambdaInterface {
+      private readonly foo = 'foo';
+
+      @idempotent({
+        persistenceStore: new PersistenceLayerTestClass(),
+      })
+      public async handler(
+        _event: unknown,
+        _context: Context
+      ): Promise<string> {
+        return this.privateMethod();
+      }
+
+      public privateMethod(): string {
+        return `private ${this.foo}`;
+      }
+    }
+
+    const handlerClass = new TestClass();
+    const handler = handlerClass.handler.bind(handlerClass);
+
+    // Act
+    const result = await handler({}, context);
+
+    // Assess
+    expect(result).toBe('private foo');
+  });
 });