fix(idempotency): pass lambda context remaining time to save inprogress (#1540)

* fix in_progress_expiration timestamp in idempotency handler

* add lambda context mock to tests

* adjust handler signature with context

* revert tests to create config, if not passed, undefined temporary

* decorator can now register context too

Author: Alexander Schueren

packages/idempotency/src/IdempotencyHandler.ts

@@ -141,7 +141,8 @@ export class IdempotencyHandler<U> {
 
     try {
       await this.persistenceStore.saveInProgress(
-        this.functionPayloadToBeHashed
+        this.functionPayloadToBeHashed,
+        this.idempotencyConfig.lambdaContext?.getRemainingTimeInMillis()
       );
     } catch (e) {
       if (e instanceof IdempotencyItemAlreadyExistsError) {

packages/idempotency/src/idempotentDecorator.ts

@@ -5,6 +5,16 @@ import {
 } from './types';
 import { IdempotencyHandler } from './IdempotencyHandler';
 import { IdempotencyConfig } from './IdempotencyConfig';
+import { Context } from 'aws-lambda';
+
+const isContext = (arg: unknown): arg is Context => {
+  return (
+    arg !== undefined &&
+    arg !== null &&
+    typeof arg === 'object' &&
+    'getRemainingTimeInMillis' in arg
+  );
+};
 
 /**
  * use this function to narrow the type of options between IdempotencyHandlerOptions and IdempotencyFunctionOptions
@@ -28,13 +38,20 @@ const idempotent = function (
     descriptor: PropertyDescriptor
   ) {
     const childFunction = descriptor.value;
-    descriptor.value = function (record: GenericTempRecord) {
+    descriptor.value = function (
+      record: GenericTempRecord,
+      ...args: unknown[]
+    ) {
       const functionPayloadtoBeHashed = isFunctionOption(options)
         ? record[(options as IdempotencyFunctionOptions).dataKeywordArgument]
         : record;
       const idempotencyConfig = options.config
         ? options.config
         : new IdempotencyConfig({});
+      const context = args[0];
+      if (isContext(context)) {
+        idempotencyConfig.registerLambdaContext(context);
+      }
       const idempotencyHandler = new IdempotencyHandler<GenericTempRecord>({
         functionToMakeIdempotent: childFunction,
         functionPayloadToBeHashed: functionPayloadtoBeHashed,

packages/idempotency/tests/e2e/makeFunctionIdempotent.test.FunctionCode.ts

@@ -2,6 +2,7 @@ import type { Context } from 'aws-lambda';
 import { DynamoDBPersistenceLayer } from '../../src/persistence/DynamoDBPersistenceLayer';
 import { makeFunctionIdempotent } from '../../src';
 import { Logger } from '@aws-lambda-powertools/logger';
+import { IdempotencyConfig } from '../../src';
 
 const IDEMPOTENCY_TABLE_NAME =
   process.env.IDEMPOTENCY_TABLE_NAME || 'table_name';
@@ -32,15 +33,19 @@ const processRecord = (record: Record<string, unknown>): string => {
   return 'Processing done: ' + record['foo'];
 };
 
+const idempotencyConfig = new IdempotencyConfig({});
+
 const processIdempotently = makeFunctionIdempotent(processRecord, {
   persistenceStore: dynamoDBPersistenceLayer,
   dataKeywordArgument: 'foo',
+  config: idempotencyConfig,
 });
 
 export const handler = async (
   _event: EventRecords,
   _context: Context
 ): Promise<void> => {
+  idempotencyConfig.registerLambdaContext(_context);
   for (const record of _event.records) {
     const result = await processIdempotently(record);
     logger.info(result.toString());
@@ -52,12 +57,14 @@ export const handler = async (
 const processIdempotentlyCustomized = makeFunctionIdempotent(processRecord, {
   persistenceStore: ddbPersistenceLayerCustomized,
   dataKeywordArgument: 'foo',
+  config: idempotencyConfig,
 });
 
 export const handlerCustomized = async (
   _event: EventRecords,
   _context: Context
 ): Promise<void> => {
+  idempotencyConfig.registerLambdaContext(_context);
   for (const record of _event.records) {
     const result = await processIdempotentlyCustomized(record);
     logger.info(result.toString());

packages/idempotency/tests/e2e/makeFunctionIdempotent.test.ts

@@ -109,6 +109,7 @@ describe('Idempotency e2e test function wrapper, default settings', () => {
           { id: 3, foo: 'bar' },
         ],
       };
+      const invokeStart = Date.now();
       await invokeFunction(
         functionNameDefault,
         2,
@@ -136,6 +137,10 @@ describe('Idempotency e2e test function wrapper, default settings', () => {
         })
       );
       expect(resultFirst?.Item?.data).toEqual('Processing done: bar');
+      expect(resultFirst?.Item?.expiration).toBeGreaterThan(Date.now() / 1000);
+      expect(resultFirst?.Item?.in_progress_expiration).toBeGreaterThan(
+        invokeStart
+      );
       expect(resultFirst?.Item?.status).toEqual('COMPLETED');
 
       const resultSecond = await ddb.send(
@@ -145,6 +150,10 @@ describe('Idempotency e2e test function wrapper, default settings', () => {
         })
       );
       expect(resultSecond?.Item?.data).toEqual('Processing done: baz');
+      expect(resultSecond?.Item?.expiration).toBeGreaterThan(Date.now() / 1000);
+      expect(resultSecond?.Item?.in_progress_expiration).toBeGreaterThan(
+        invokeStart
+      );
       expect(resultSecond?.Item?.status).toEqual('COMPLETED');
     },
     TEST_CASE_TIMEOUT

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

@@ -14,6 +14,7 @@ import { BasePersistenceLayer, IdempotencyRecord } from '../../src/persistence';
 import { IdempotencyHandler } from '../../src/IdempotencyHandler';
 import { IdempotencyConfig } from '../../src/';
 import { MAX_RETRIES } from '../../src/constants';
+import { Context } from 'aws-lambda';
 
 class PersistenceLayerTestClass extends BasePersistenceLayer {
   protected _deleteRecord = jest.fn();
@@ -247,6 +248,39 @@ describe('Class IdempotencyHandler', () => {
       expect(mockGetRecord).toHaveBeenCalledTimes(0);
       expect(mockSaveSuccessfulResult).toHaveBeenCalledTimes(0);
     });
+
+    test('when lambdaContext is registered, we pass it to saveInProgress', async () => {
+      const mockSaveInProgress = jest.spyOn(
+        mockIdempotencyOptions.persistenceStore,
+        'saveInProgress'
+      );
+
+      const mockLambaContext: Context = {
+        getRemainingTimeInMillis(): number {
+          return 1000; // we expect this number to be passed to saveInProgress
+        },
+      } as Context;
+      const idempotencyHandlerWithContext = new IdempotencyHandler({
+        functionToMakeIdempotent: mockFunctionToMakeIdempotent,
+        functionPayloadToBeHashed: mockFunctionPayloadToBeHashed,
+        persistenceStore: mockIdempotencyOptions.persistenceStore,
+        fullFunctionPayload: mockFullFunctionPayload,
+        idempotencyConfig: new IdempotencyConfig({
+          lambdaContext: mockLambaContext,
+        }),
+      });
+
+      mockFunctionToMakeIdempotent.mockImplementation(() =>
+        Promise.resolve('result')
+      );
+
+      await expect(idempotencyHandlerWithContext.processIdempotency()).resolves;
+
+      expect(mockSaveInProgress).toBeCalledWith(
+        mockFunctionPayloadToBeHashed,
+        mockLambaContext.getRemainingTimeInMillis()
+      );
+    });
   });
 
   describe('Method: getFunctionResult', () => {

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

@@ -15,6 +15,8 @@ import {
   IdempotencyPersistenceLayerError,
 } from '../../src/Exceptions';
 import { IdempotencyConfig } from '../../src';
+import { Context } from 'aws-lambda';
+import { helloworldContext } from '@aws-lambda-powertools/commons/lib/samples/resources/contexts';
 
 const mockSaveInProgress = jest
   .spyOn(BasePersistenceLayer.prototype, 'saveInProgress')
@@ -26,6 +28,10 @@ const mockGetRecord = jest
   .spyOn(BasePersistenceLayer.prototype, 'getRecord')
   .mockImplementation();
 
+const dummyContext = helloworldContext;
+
+const mockConfig: IdempotencyConfig = new IdempotencyConfig({});
+
 class PersistenceLayerTestClass extends BasePersistenceLayer {
   protected _deleteRecord = jest.fn();
   protected _getRecord = jest.fn();
@@ -39,21 +45,25 @@ class TestinClassWithLambdaHandler {
   @idempotentLambdaHandler({
     persistenceStore: new PersistenceLayerTestClass(),
   })
-  public testing(record: Record<string, unknown>): string {
+  // eslint-disable-next-line @typescript-eslint/no-unused-vars
+  public testing(record: Record<string, unknown>, context: Context): string {
     functionalityToDecorate(record);
 
     return 'Hi';
   }
 }
 
 class TestingClassWithFunctionDecorator {
-  public handler(record: Record<string, unknown>): string {
+  public handler(record: Record<string, unknown>, context: Context): string {
+    mockConfig.registerLambdaContext(context);
+
     return this.proccessRecord(record);
   }
 
   @idempotentFunction({
     persistenceStore: new PersistenceLayerTestClass(),
     dataKeywordArgument: 'testingKey',
+    config: mockConfig,
   })
   public proccessRecord(record: Record<string, unknown>): string {
     functionalityToDecorate(record);
@@ -72,11 +82,14 @@ describe('Given a class with a function to decorate', (classWithLambdaHandler =
 
   describe('When wrapping a function with no previous executions', () => {
     beforeEach(async () => {
-      await classWithFunctionDecorator.handler(inputRecord);
+      await classWithFunctionDecorator.handler(inputRecord, dummyContext);
     });
 
     test('Then it will save the record to INPROGRESS', () => {
-      expect(mockSaveInProgress).toBeCalledWith(keyValueToBeSaved);
+      expect(mockSaveInProgress).toBeCalledWith(
+        keyValueToBeSaved,
+        dummyContext.getRemainingTimeInMillis()
+      );
     });
 
     test('Then it will call the function that was decorated', () => {
@@ -92,11 +105,14 @@ describe('Given a class with a function to decorate', (classWithLambdaHandler =
   });
   describe('When wrapping a function with no previous executions', () => {
     beforeEach(async () => {
-      await classWithLambdaHandler.testing(inputRecord);
+      await classWithLambdaHandler.testing(inputRecord, dummyContext);
     });
 
     test('Then it will save the record to INPROGRESS', () => {
-      expect(mockSaveInProgress).toBeCalledWith(inputRecord);
+      expect(mockSaveInProgress).toBeCalledWith(
+        inputRecord,
+        dummyContext.getRemainingTimeInMillis()
+      );
     });
 
     test('Then it will call the function that was decorated', () => {
@@ -122,14 +138,17 @@ describe('Given a class with a function to decorate', (classWithLambdaHandler =
         new IdempotencyRecord(idempotencyOptions)
       );
       try {
-        await classWithLambdaHandler.testing(inputRecord);
+        await classWithLambdaHandler.testing(inputRecord, dummyContext);
       } catch (e) {
         resultingError = e as Error;
       }
     });
 
     test('Then it will attempt to save the record to INPROGRESS', () => {
-      expect(mockSaveInProgress).toBeCalledWith(inputRecord);
+      expect(mockSaveInProgress).toBeCalledWith(
+        inputRecord,
+        dummyContext.getRemainingTimeInMillis()
+      );
     });
 
     test('Then it will get the previous execution record', () => {
@@ -159,14 +178,17 @@ describe('Given a class with a function to decorate', (classWithLambdaHandler =
         new IdempotencyRecord(idempotencyOptions)
       );
       try {
-        await classWithLambdaHandler.testing(inputRecord);
+        await classWithLambdaHandler.testing(inputRecord, dummyContext);
       } catch (e) {
         resultingError = e as Error;
       }
     });
 
     test('Then it will attempt to save the record to INPROGRESS', () => {
-      expect(mockSaveInProgress).toBeCalledWith(inputRecord);
+      expect(mockSaveInProgress).toBeCalledWith(
+        inputRecord,
+        dummyContext.getRemainingTimeInMillis()
+      );
     });
 
     test('Then it will get the previous execution record', () => {
@@ -195,11 +217,14 @@ describe('Given a class with a function to decorate', (classWithLambdaHandler =
       mockGetRecord.mockResolvedValue(
         new IdempotencyRecord(idempotencyOptions)
       );
-      await classWithLambdaHandler.testing(inputRecord);
+      await classWithLambdaHandler.testing(inputRecord, dummyContext);
     });
 
     test('Then it will attempt to save the record to INPROGRESS', () => {
-      expect(mockSaveInProgress).toBeCalledWith(inputRecord);
+      expect(mockSaveInProgress).toBeCalledWith(
+        inputRecord,
+        dummyContext.getRemainingTimeInMillis()
+      );
     });
 
     test('Then it will get the previous execution record', () => {
@@ -215,7 +240,7 @@ describe('Given a class with a function to decorate', (classWithLambdaHandler =
     class TestinClassWithLambdaHandlerWithConfig {
       @idempotentLambdaHandler({
         persistenceStore: new PersistenceLayerTestClass(),
-        config: new IdempotencyConfig({}),
+        config: new IdempotencyConfig({ lambdaContext: dummyContext }),
       })
       public testing(record: Record<string, unknown>): string {
         functionalityToDecorate(record);
@@ -237,7 +262,10 @@ describe('Given a class with a function to decorate', (classWithLambdaHandler =
     });
 
     test('Then it will attempt to save the record to INPROGRESS', () => {
-      expect(mockSaveInProgress).toBeCalledWith(inputRecord);
+      expect(mockSaveInProgress).toBeCalledWith(
+        inputRecord,
+        dummyContext.getRemainingTimeInMillis()
+      );
     });
 
     test('Then an IdempotencyPersistenceLayerError is thrown', () => {