fix(idempotency): include sk in error msgs when using composite key

packages/idempotency/src/IdempotencyHandler.ts

@@ -120,7 +120,11 @@ export class IdempotencyHandler<Func extends AnyFunction> {
         );
       }
       throw new IdempotencyAlreadyInProgressError(
-        `There is already an execution in progress with idempotency key: ${idempotencyRecord.idempotencyKey}`
+        `There is already an execution in progress with idempotency key: ${idempotencyRecord.idempotencyKey}${
+          idempotencyRecord.sortKey
+            ? ` and sort key: ${idempotencyRecord.sortKey}`
+            : ''
+        }`
       );
     }
 

packages/idempotency/src/persistence/DynamoDBPersistenceLayer.ts

@@ -122,6 +122,7 @@ class DynamoDBPersistenceLayer extends BasePersistenceLayer {
 
     return new IdempotencyRecord({
       idempotencyKey: item[this.keyAttr],
+      sortKey: this.sortKeyAttr && item[this.sortKeyAttr],
       status: item[this.statusAttr],
       expiryTimestamp: item[this.expiryAttr],
       inProgressExpiryTimestamp: item[this.inProgressExpiryAttr],
@@ -207,14 +208,17 @@ class DynamoDBPersistenceLayer extends BasePersistenceLayer {
           item &&
           new IdempotencyRecord({
             idempotencyKey: item[this.keyAttr],
+            sortKey: this.sortKeyAttr && item[this.sortKeyAttr],
             status: item[this.statusAttr],
             expiryTimestamp: item[this.expiryAttr],
             inProgressExpiryTimestamp: item[this.inProgressExpiryAttr],
             responseData: item[this.dataAttr],
             payloadHash: item[this.validationKeyAttr],
           });
         throw new IdempotencyItemAlreadyExistsError(
-          `Failed to put record for already existing idempotency key: ${record.idempotencyKey}`,
+          `Failed to put record for already existing idempotency key: ${record.idempotencyKey}${
+            this.sortKeyAttr ? ` and sort key: ${record.sortKey}` : ''
+          }`,
           idempotencyRecord
         );
       }

packages/idempotency/src/persistence/IdempotencyRecord.ts

@@ -5,6 +5,7 @@ import type {
   IdempotencyRecordOptions,
   IdempotencyRecordStatusValue,
 } from '../types/IdempotencyRecord.js';
+import type { DynamoDBPersistenceLayer } from './DynamoDBPersistenceLayer.js';
 
 /**
  * Class representing an idempotency record.
@@ -19,6 +20,10 @@ class IdempotencyRecord {
    * The idempotency key of the record that is used to identify the record.
    */
   public idempotencyKey: string;
+  /**
+   * An optional sort key that can be used with the {@link DynamoDBPersistenceLayer | `DynamoDBPersistenceLayer`}.
+   */
+  public sortKey?: string;
   /**
    * The expiry timestamp of the in progress record in milliseconds UTC.
    */
@@ -46,6 +51,7 @@ class IdempotencyRecord {
     this.responseData = config.responseData;
     this.payloadHash = config.payloadHash;
     this.status = config.status;
+    this.sortKey = config.sortKey;
   }
 
   /**

packages/idempotency/src/types/IdempotencyRecord.ts

@@ -11,11 +11,35 @@ type IdempotencyRecordStatusValue =
  * Options for creating a new IdempotencyRecord
  */
 type IdempotencyRecordOptions = {
+  /**
+   * The idempotency key of the record that is used to identify the record.
+   */
   idempotencyKey: string;
+  /**
+   * An optional sort key that can be used with the {@link DynamoDBPersistenceLayer | `DynamoDBPersistenceLayer`}.
+   */
+  sortKey?: string;
+  /**
+   * The idempotency record status can be COMPLETED, IN_PROGRESS or EXPIRED.
+   * We check the status during idempotency processing to make sure we don't process an expired record and handle concurrent requests.
+   * {@link constants.IdempotencyRecordStatusValue | IdempotencyRecordStatusValue}
+   */
   status: IdempotencyRecordStatusValue;
+  /**
+   * The expiry timestamp of the record in milliseconds UTC.
+   */
   expiryTimestamp?: number;
+  /**
+   * The expiry timestamp of the in progress record in milliseconds UTC.
+   */
   inProgressExpiryTimestamp?: number;
+  /**
+   * The response data of the request, this will be returned if the payload hash matches.
+   */
   responseData?: JSONValue;
+  /**
+   * The hash of the payload of the request, used for comparing requests.
+   */
   payloadHash?: string;
 };
 

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

@@ -49,25 +49,46 @@ describe('Class IdempotencyHandler', () => {
   });
 
   describe('Method: determineResultFromIdempotencyRecord', () => {
-    it('throws when the record is in progress and within expiry window', async () => {
-      // Prepare
-      const stubRecord = new IdempotencyRecord({
-        idempotencyKey: 'idempotencyKey',
-        expiryTimestamp: Date.now() + 1000, // should be in the future
-        inProgressExpiryTimestamp: 0, // less than current time in milliseconds
-        responseData: { responseData: 'responseData' },
-        payloadHash: 'payloadHash',
-        status: IdempotencyRecordStatus.INPROGRESS,
-      });
-
-      // Act & Assess
-      expect(stubRecord.isExpired()).toBe(false);
-      expect(stubRecord.getStatus()).toBe(IdempotencyRecordStatus.INPROGRESS);
-      expect(() =>
-        idempotentHandler.determineResultFromIdempotencyRecord(stubRecord)
-      ).toThrow(IdempotencyAlreadyInProgressError);
-      expect(mockResponseHook).not.toHaveBeenCalled();
-    });
+    it.each([
+      {
+        keys: {
+          idempotencyKey: 'idempotencyKey',
+          sortKey: 'sk',
+        },
+        expectedErrorMsg:
+          'There is already an execution in progress with idempotency key: idempotencyKey and sort key: sk',
+        case: 'pk & sk',
+      },
+      {
+        keys: {
+          idempotencyKey: 'idempotencyKey',
+        },
+        expectedErrorMsg:
+          'There is already an execution in progress with idempotency key: idempotencyKey',
+        case: 'pk only',
+      },
+    ])(
+      'throws when the record is in progress and within expiry window ($case)',
+      async ({ keys, expectedErrorMsg }) => {
+        // Prepare
+        const stubRecord = new IdempotencyRecord({
+          ...keys,
+          expiryTimestamp: Date.now() + 1000, // should be in the future
+          inProgressExpiryTimestamp: 0, // less than current time in milliseconds
+          responseData: { responseData: 'responseData' },
+          payloadHash: 'payloadHash',
+          status: IdempotencyRecordStatus.INPROGRESS,
+        });
+
+        // Act & Assess
+        expect(stubRecord.isExpired()).toBe(false);
+        expect(stubRecord.getStatus()).toBe(IdempotencyRecordStatus.INPROGRESS);
+        expect(() =>
+          idempotentHandler.determineResultFromIdempotencyRecord(stubRecord)
+        ).toThrow(new IdempotencyAlreadyInProgressError(expectedErrorMsg));
+        expect(mockResponseHook).not.toHaveBeenCalled();
+      }
+    );
 
     it('throws when the record is in progress and outside expiry window', async () => {
       // Prepare

packages/idempotency/tests/unit/persistence/DynamoDbPersistenceLayer.test.ts

@@ -346,41 +346,75 @@ describe('Class: DynamoDBPersistenceLayer', () => {
       persistenceLayerSpy.mockRestore();
     });
 
-    it('throws when called with a record that fails any condition', async () => {
-      // Prepare
-      const record = new IdempotencyRecord({
-        idempotencyKey: dummyKey,
-        status: IdempotencyRecordStatus.EXPIRED,
-        expiryTimestamp: 0,
-      });
-      const expiration = Date.now();
-      client.on(PutItemCommand).rejects(
-        new ConditionalCheckFailedException({
-          $metadata: {
-            httpStatusCode: 400,
-            requestId: 'someRequestId',
-          },
-          message: 'Conditional check failed',
-          Item: {
-            id: { S: 'test-key' },
-            status: { S: 'INPROGRESS' },
-            expiration: { N: expiration.toString() },
-          },
-        })
-      );
-
-      // Act & Assess
-      await expect(persistenceLayer._putRecord(record)).rejects.toThrowError(
-        new IdempotencyItemAlreadyExistsError(
-          `Failed to put record for already existing idempotency key: ${record.idempotencyKey}`,
-          new IdempotencyRecord({
-            idempotencyKey: 'test-key',
-            status: IdempotencyRecordStatus.INPROGRESS,
-            expiryTimestamp: expiration,
+    it.each([
+      {
+        keys: {
+          id: 'idempotency#my-lambda-function',
+          sortKey: dummyKey,
+        },
+        case: 'composite key',
+      },
+      {
+        keys: {
+          id: dummyKey,
+        },
+        case: 'single key',
+      },
+    ])(
+      'throws when called with a record that fails any condition ($case)',
+      async ({ keys }) => {
+        // Prepare
+        const { id, sortKey } = keys;
+
+        const record = new IdempotencyRecord({
+          idempotencyKey: id,
+          sortKey,
+          status: IdempotencyRecordStatus.EXPIRED,
+          expiryTimestamp: 0,
+        });
+        const expiration = Date.now();
+        client.on(PutItemCommand).rejects(
+          new ConditionalCheckFailedException({
+            $metadata: {
+              httpStatusCode: 400,
+              requestId: 'someRequestId',
+            },
+            message: 'Conditional check failed',
+            Item: {
+              id: { S: 'test-key' },
+              ...(sortKey ? { sortKey: { S: sortKey } } : {}),
+              status: { S: 'INPROGRESS' },
+              expiration: { N: expiration.toString() },
+            },
           })
-        )
-      );
-    });
+        );
+        const testPersistenceLayer = sortKey
+          ? new DynamoDBPersistenceLayerTestClass({
+              tableName: dummyTableName,
+              sortKeyAttr: 'sortKey',
+            })
+          : persistenceLayer;
+
+        // Act & Assess
+        await expect(
+          testPersistenceLayer._putRecord(record)
+        ).rejects.toThrowError(
+          new IdempotencyItemAlreadyExistsError(
+            `Failed to put record for already existing idempotency key: ${
+              sortKey
+                ? `${record.idempotencyKey} and sort key: ${sortKey}`
+                : record.idempotencyKey
+            }`,
+            new IdempotencyRecord({
+              idempotencyKey: 'test-key',
+              sortKey,
+              status: IdempotencyRecordStatus.INPROGRESS,
+              expiryTimestamp: expiration,
+            })
+          )
+        );
+      }
+    );
 
     it('throws when encountering an unknown error', async () => {
       // Prepare
@@ -445,7 +479,7 @@ describe('Class: DynamoDBPersistenceLayer', () => {
       );
     });
 
-    it('it builds the request correctly when using composite keys', async () => {
+    it('builds the request correctly when using composite keys', async () => {
       // Prepare
       const persistenceLayer = new DynamoDBPersistenceLayerTestClass({
         tableName: dummyTableName,