feat(idempotency): makeHandlerIdempotent middy middleware (#1474)

* feat: makeHandlerIdempotent middy middleware

Author: dreamorosi

packages/idempotency/src/IdempotencyConfig.ts

@@ -1,3 +1,4 @@
+import { EnvironmentVariablesService } from './config';
 import type { Context } from 'aws-lambda';
 import type { IdempotencyConfigOptions } from './types';
 
@@ -10,6 +11,8 @@ class IdempotencyConfig {
   public payloadValidationJmesPath?: string;
   public throwOnNoIdempotencyKey: boolean;
   public useLocalCache: boolean;
+  readonly #envVarsService: EnvironmentVariablesService;
+  readonly #enabled: boolean = true;
 
   public constructor(config: IdempotencyConfigOptions) {
     this.eventKeyJmesPath = config.eventKeyJmesPath ?? '';
@@ -20,6 +23,17 @@ class IdempotencyConfig {
     this.maxLocalCacheSize = config.maxLocalCacheSize ?? 1000;
     this.hashFunction = config.hashFunction ?? 'md5';
     this.lambdaContext = config.lambdaContext;
+    this.#envVarsService = new EnvironmentVariablesService();
+    this.#enabled = this.#envVarsService.getIdempotencyEnabled();
+  }
+
+  /**
+   * Determines if the idempotency feature is enabled.
+   *
+   * @returns {boolean} Returns true if the idempotency feature is enabled.
+   */
+  public isEnabled(): boolean {
+    return this.#enabled;
   }
 
   public registerLambdaContext(context: Context): void {

packages/idempotency/src/IdempotencyHandler.ts

@@ -8,6 +8,7 @@ import {
 } from './Exceptions';
 import { BasePersistenceLayer, IdempotencyRecord } from './persistence';
 import { IdempotencyConfig } from './IdempotencyConfig';
+import { MAX_RETRIES } from './constants';
 
 export class IdempotencyHandler<U> {
   private readonly fullFunctionPayload: Record<string, unknown>;
@@ -36,9 +37,9 @@ export class IdempotencyHandler<U> {
     });
   }
 
-  public determineResultFromIdempotencyRecord(
+  public static determineResultFromIdempotencyRecord(
     idempotencyRecord: IdempotencyRecord
-  ): Promise<U> | U {
+  ): Promise<unknown> | unknown {
     if (idempotencyRecord.getStatus() === IdempotencyRecordStatus.EXPIRED) {
       throw new IdempotencyInconsistentStateError(
         'Item has expired during processing and may not longer be valid.'
@@ -61,7 +62,7 @@ export class IdempotencyHandler<U> {
       }
     }
 
-    return idempotencyRecord.getResponse() as U;
+    return idempotencyRecord.getResponse();
   }
 
   public async getFunctionResult(): Promise<U> {
@@ -96,26 +97,30 @@ export class IdempotencyHandler<U> {
 
   /**
    * Main entry point for the handler
-   * IdempotencyInconsistentStateError can happen under rare but expected cases
-   * when persistent state changes in the small time between put & get requests.
-   * In most cases we can retry successfully on this exception.
+   *
+   * In some rare cases, when the persistent state changes in small time
+   * window, we might get an `IdempotencyInconsistentStateError`. In such
+   * cases we can safely retry the handling a few times.
    */
   public async handle(): Promise<U> {
-    const MAX_RETRIES = 2;
-    for (let i = 1; i <= MAX_RETRIES; i++) {
+    let e;
+    for (let retryNo = 0; retryNo <= MAX_RETRIES; retryNo++) {
       try {
         return await this.processIdempotency();
-      } catch (e) {
+      } catch (error) {
         if (
-          !(e instanceof IdempotencyAlreadyInProgressError) ||
-          i === MAX_RETRIES
+          error instanceof IdempotencyInconsistentStateError &&
+          retryNo < MAX_RETRIES
         ) {
-          throw e;
+          // Retry
+          continue;
         }
+        // Retries exhausted or other error
+        e = error;
+        break;
       }
     }
-    /* istanbul ignore next */
-    throw new Error('This should never happen');
+    throw e;
   }
 
   public async processIdempotency(): Promise<U> {
@@ -128,7 +133,9 @@ export class IdempotencyHandler<U> {
         const idempotencyRecord: IdempotencyRecord =
           await this.persistenceStore.getRecord(this.functionPayloadToBeHashed);
 
-        return this.determineResultFromIdempotencyRecord(idempotencyRecord);
+        return IdempotencyHandler.determineResultFromIdempotencyRecord(
+          idempotencyRecord
+        ) as U;
       } else {
         throw new IdempotencyPersistenceLayerError();
       }

packages/idempotency/src/config/ConfigServiceInterface.ts

@@ -4,6 +4,8 @@ interface ConfigServiceInterface {
   getServiceName(): string;
 
   getFunctionName(): string;
+
+  getIdempotencyEnabled(): boolean;
 }
 
 export { ConfigServiceInterface };

packages/idempotency/src/config/EnvironmentVariablesService.ts

@@ -21,6 +21,7 @@ class EnvironmentVariablesService
 {
   // Reserved environment variables
   private functionNameVariable = 'AWS_LAMBDA_FUNCTION_NAME';
+  private idempotencyDisabledVariable = 'POWERTOOLS_IDEMPOTENCY_DISABLED';
 
   /**
    * It returns the value of the AWS_LAMBDA_FUNCTION_NAME environment variable.
@@ -30,6 +31,17 @@ class EnvironmentVariablesService
   public getFunctionName(): string {
     return this.get(this.functionNameVariable);
   }
+
+  /**
+   * It returns whether the idempotency feature is enabled or not.
+   *
+   * Reads the value of the POWERTOOLS_IDEMPOTENCY_DISABLED environment variable.
+   *
+   * @returns {boolean}
+   */
+  public getIdempotencyEnabled(): boolean {
+    return !this.isValueTrue(this.get(this.idempotencyDisabledVariable));
+  }
 }
 
 export { EnvironmentVariablesService };

packages/idempotency/src/constants.ts

@@ -0,0 +1,10 @@
+/**
+ * Number of times to retry a request in case of `IdempotencyInconsistentStateError`
+ *
+ * Used in `IdempotencyHandler` and `makeHandlerIdempotent`
+ *
+ * @internal
+ */
+const MAX_RETRIES = 2;
+
+export { MAX_RETRIES };

packages/idempotency/src/makeFunctionIdempotent.ts

@@ -1,41 +1,58 @@
+import type { Context } from 'aws-lambda';
 import type {
   AnyFunctionWithRecord,
   AnyIdempotentFunction,
-  GenericTempRecord,
   IdempotencyFunctionOptions,
 } from './types';
 import { IdempotencyHandler } from './IdempotencyHandler';
 import { IdempotencyConfig } from './IdempotencyConfig';
 
+const isContext = (arg: unknown): arg is Context => {
+  return (
+    arg !== undefined &&
+    arg !== null &&
+    typeof arg === 'object' &&
+    'getRemainingTimeInMillis' in arg
+  );
+};
+
 const makeFunctionIdempotent = function <U>(
   fn: AnyFunctionWithRecord<U>,
   options: IdempotencyFunctionOptions
-): AnyIdempotentFunction<U> {
+): AnyIdempotentFunction<U> | AnyFunctionWithRecord<U> {
+  const idempotencyConfig = options.config
+    ? options.config
+    : new IdempotencyConfig({});
+
   const wrappedFn: AnyIdempotentFunction<U> = function (
-    record: GenericTempRecord
+    ...args: Parameters<AnyFunctionWithRecord<U>>
   ): Promise<U> {
+    const payload = args[0];
+    const context = args[1];
+
     if (options.dataKeywordArgument === undefined) {
       throw new Error(
         `Missing data keyword argument ${options.dataKeywordArgument}`
       );
     }
-    const idempotencyConfig = options.config
-      ? options.config
-      : new IdempotencyConfig({});
+    if (isContext(context)) {
+      idempotencyConfig.registerLambdaContext(context);
+    }
     const idempotencyHandler: IdempotencyHandler<U> = new IdempotencyHandler<U>(
       {
         functionToMakeIdempotent: fn,
-        functionPayloadToBeHashed: record[options.dataKeywordArgument],
+        functionPayloadToBeHashed: payload[options.dataKeywordArgument],
         idempotencyConfig: idempotencyConfig,
         persistenceStore: options.persistenceStore,
-        fullFunctionPayload: record,
+        fullFunctionPayload: payload,
       }
     );
 
     return idempotencyHandler.handle();
   };
 
-  return wrappedFn;
+  if (idempotencyConfig.isEnabled()) return wrappedFn;
+  else return fn;
 };
 
 export { makeFunctionIdempotent };

packages/idempotency/src/middleware/index.ts

@@ -0,0 +1 @@
+export * from './makeHandlerIdempotent';