@squiz/db-lib 1.72.0 → 1.74.0

package/src/dynamodb/AbstractDynamoDbRepository.spec.ts

@@ -16,6 +16,8 @@ import {
   TransactWriteCommandInput,
   BatchGetCommand,
   BatchGetCommandInput,
+  BatchWriteCommand,
+  BatchWriteCommandInput,
 } from '@aws-sdk/lib-dynamodb';
 import { ConditionalCheckFailedException } from '@aws-sdk/client-dynamodb';
 import { DynamoDbManager, Transaction } from './DynamoDbManager';
@@ -230,6 +232,8 @@ describe('AbstractRepository', () => {
         Attributes: {
           name: 'foo',
           age: 99,
+          // country attribute is part of gsi key
+          // hence updating this will also update gsi key value
           country: 'au-updated',
           data: {},
         },
@@ -237,12 +241,14 @@ describe('AbstractRepository', () => {
       const input: UpdateCommandInput = {
         TableName: TABLE_NAME,
         Key: { pk: 'test_item#foo', sk: '#meta' },
-        UpdateExpression: 'SET #country = :country',
+        UpdateExpression: 'SET #country = :country, #gsi1_pk = :gsi1_pk',
         ExpressionAttributeNames: {
           '#country': 'country',
+          '#gsi1_pk': 'gsi1_pk',
         },
         ExpressionAttributeValues: {
           ':country': 'au-updated',
+          ':gsi1_pk': 'country#au-updated',
         },
         ConditionExpression: `attribute_exists(pk)`,
       };
@@ -263,6 +269,60 @@ describe('AbstractRepository', () => {
       );
     });
 
+    it('should only update the changed attributes', async () => {
+      ddbClientMock.on(GetCommand).resolves({
+        $metadata: {
+          httpStatusCode: 200,
+        },
+        Item: {
+          name: 'foo',
+          age: 99,
+          country: 'au',
+          data: {},
+        },
+      });
+      ddbClientMock.on(UpdateCommand).resolves({
+        $metadata: {
+          httpStatusCode: 200,
+        },
+        Attributes: {
+          name: 'foo',
+          age: 99,
+          country: 'au',
+          data: { active: true },
+        },
+      });
+      const input: UpdateCommandInput = {
+        TableName: TABLE_NAME,
+        Key: { pk: 'test_item#foo', sk: '#meta' },
+        UpdateExpression: 'SET #data = :data',
+        ExpressionAttributeNames: {
+          '#data': 'data',
+        },
+        ExpressionAttributeValues: {
+          ':data': { active: true },
+        },
+        ConditionExpression: `attribute_exists(pk)`,
+      };
+
+      const updateItem = {
+        name: 'foo',
+        age: 99,
+        // this is the only change attribute value
+        data: { active: true },
+      };
+      const result = await repository.updateItem(updateItem);
+      expect(ddbClientMock).toHaveReceivedNthCommandWith(2, UpdateCommand, input);
+      expect(result).toEqual(
+        new TestItem({
+          name: 'foo',
+          age: 99,
+          country: 'au',
+          data: { active: true },
+        }),
+      );
+    });
+
     it('should not trigger update request if the input attributes are same as in the existing item', async () => {
       ddbClientMock.on(GetCommand).resolves({
         $metadata: {
@@ -557,6 +617,156 @@ describe('AbstractRepository', () => {
       ]);
     });
 
+    it('should retry if unprocessed keys returned', async () => {
+      const input1: BatchGetCommandInput = {
+        RequestItems: {
+          [TABLE_NAME]: {
+            Keys: [
+              { pk: 'test_item#foo', sk: '#meta' },
+              { pk: 'test_item#foo2', sk: '#meta' },
+            ],
+          },
+        },
+      };
+      const input2: BatchGetCommandInput = {
+        RequestItems: {
+          [TABLE_NAME]: {
+            Keys: [{ pk: 'test_item#foo2', sk: '#meta' }],
+          },
+        },
+      };
+
+      ddbClientMock.on(BatchGetCommand, input1).resolves({
+        $metadata: {
+          httpStatusCode: 200,
+        },
+        Responses: {
+          [TABLE_NAME]: [
+            {
+              name: 'foo',
+              age: 99,
+              country: 'au',
+              data: {},
+              data2: '{"foo":"bar","num":123}',
+            },
+          ],
+        },
+        UnprocessedKeys: {
+          [TABLE_NAME]: {
+            Keys: [{ pk: 'test_item#foo2', sk: '#meta' }],
+          },
+        },
+      });
+
+      ddbClientMock.on(BatchGetCommand, input2).resolves({
+        $metadata: {
+          httpStatusCode: 200,
+        },
+        Responses: {
+          [TABLE_NAME]: [
+            {
+              name: 'foo2',
+              age: 999,
+              country: 'au',
+              data: {},
+              data2: '{"foo":"bar","num":123}',
+            },
+          ],
+        },
+        UnprocessedKeys: {},
+      });
+
+      const requestItems = [{ name: 'foo' }, { name: 'foo2' }];
+      const result = await repository.getItems(requestItems);
+      expect(ddbClientMock).toHaveReceivedNthCommandWith(1, BatchGetCommand, input1);
+      expect(ddbClientMock).toHaveReceivedNthCommandWith(2, BatchGetCommand, input2);
+      expect(result).toEqual([
+        new TestItem({
+          name: 'foo',
+          age: 99,
+          country: 'au',
+          data: {},
+          data2: {
+            foo: 'bar',
+            num: 123,
+          },
+        }),
+        new TestItem({
+          name: 'foo2',
+          age: 999,
+          country: 'au',
+          data: {},
+          data2: {
+            foo: 'bar',
+            num: 123,
+          },
+        }),
+      ]);
+    });
+
+    it('should fail after max retries for unprocessed keys', async () => {
+      const input1: BatchGetCommandInput = {
+        RequestItems: {
+          [TABLE_NAME]: {
+            Keys: [
+              { pk: 'test_item#foo', sk: '#meta' },
+              { pk: 'test_item#foo2', sk: '#meta' },
+            ],
+          },
+        },
+      };
+      const input2: BatchGetCommandInput = {
+        RequestItems: {
+          [TABLE_NAME]: {
+            Keys: [{ pk: 'test_item#foo2', sk: '#meta' }],
+          },
+        },
+      };
+
+      ddbClientMock.on(BatchGetCommand, input1).resolves({
+        $metadata: {
+          httpStatusCode: 200,
+        },
+        Responses: {
+          [TABLE_NAME]: [
+            {
+              name: 'foo',
+              age: 99,
+              country: 'au',
+              data: {},
+              data2: '{"foo":"bar","num":123}',
+            },
+          ],
+        },
+        UnprocessedKeys: {
+          [TABLE_NAME]: {
+            Keys: [{ pk: 'test_item#foo2', sk: '#meta' }],
+          },
+        },
+      });
+
+      ddbClientMock.on(BatchGetCommand, input2).resolves({
+        $metadata: {
+          httpStatusCode: 200,
+        },
+        Responses: {},
+        UnprocessedKeys: {
+          [TABLE_NAME]: {
+            Keys: [{ pk: 'test_item#foo2', sk: '#meta' }],
+          },
+        },
+      });
+
+      const requestItems = [{ name: 'foo' }, { name: 'foo2' }];
+      await expect(repository.getItems(requestItems)).rejects.toEqual(
+        new Error('Maximum allowed retries exceeded for unprocessed items'),
+      );
+      expect(ddbClientMock).toHaveReceivedNthCommandWith(1, BatchGetCommand, input1);
+      expect(ddbClientMock).toHaveReceivedNthCommandWith(2, BatchGetCommand, input2);
+      expect(ddbClientMock).toHaveReceivedNthCommandWith(3, BatchGetCommand, input2);
+      expect(ddbClientMock).toHaveReceivedNthCommandWith(4, BatchGetCommand, input2);
+    });
+
     it('should request BatchGetItem in batch of 100 items to get result', async () => {
       ddbClientMock.on(BatchGetCommand).resolves({
         $metadata: {
@@ -608,6 +818,207 @@ describe('AbstractRepository', () => {
     });
   });
 
+  describe('deleteItems()', () => {
+    it('should use batchWrite() to get result', async () => {
+      ddbClientMock.on(BatchWriteCommand).resolves({
+        $metadata: {
+          httpStatusCode: 200,
+        },
+        ItemCollectionMetrics: {
+          [TABLE_NAME]: [{}],
+        },
+      });
+      const input: BatchWriteCommandInput = {
+        RequestItems: {
+          [TABLE_NAME]: [
+            {
+              DeleteRequest: {
+                Key: { pk: 'test_item#foo', sk: '#meta' },
+              },
+            },
+            {
+              DeleteRequest: {
+                Key: { pk: 'test_item#foo2', sk: '#meta' },
+              },
+            },
+          ],
+        },
+      };
+
+      const requestItems = [{ name: 'foo' }, { name: 'foo2' }];
+      await repository.deleteItems(requestItems);
+      expect(ddbClientMock).toHaveReceivedCommandWith(BatchWriteCommand, input);
+    });
+
+    it('should use re-try if unprocessed items returned', async () => {
+      const input1: BatchWriteCommandInput = {
+        RequestItems: {
+          [TABLE_NAME]: [
+            {
+              DeleteRequest: {
+                Key: { pk: 'test_item#foo', sk: '#meta' },
+              },
+            },
+            {
+              DeleteRequest: {
+                Key: { pk: 'test_item#foo2', sk: '#meta' },
+              },
+            },
+          ],
+        },
+      };
+      const input2: BatchWriteCommandInput = {
+        RequestItems: {
+          [TABLE_NAME]: [
+            {
+              DeleteRequest: {
+                Key: { pk: 'test_item#foo2', sk: '#meta' },
+              },
+            },
+          ],
+        },
+      };
+
+      ddbClientMock.on(BatchWriteCommand, input1).resolves({
+        $metadata: {
+          httpStatusCode: 200,
+        },
+        UnprocessedItems: {
+          [TABLE_NAME]: [
+            {
+              DeleteRequest: {
+                Key: { pk: 'test_item#foo2', sk: '#meta' },
+              },
+            },
+          ],
+        },
+      });
+      ddbClientMock.on(BatchWriteCommand, input2).resolves({
+        $metadata: {
+          httpStatusCode: 200,
+        },
+        UnprocessedItems: {},
+      });
+
+      const requestItems = [{ name: 'foo' }, { name: 'foo2' }];
+      await repository.deleteItems(requestItems);
+      expect(ddbClientMock).toHaveReceivedNthCommandWith(1, BatchWriteCommand, input1);
+      expect(ddbClientMock).toHaveReceivedNthCommandWith(2, BatchWriteCommand, input2);
+    });
+
+    it('should fail after max number of retries', async () => {
+      const input1: BatchWriteCommandInput = {
+        RequestItems: {
+          [TABLE_NAME]: [
+            {
+              DeleteRequest: {
+                Key: { pk: 'test_item#foo', sk: '#meta' },
+              },
+            },
+            {
+              DeleteRequest: {
+                Key: { pk: 'test_item#foo2', sk: '#meta' },
+              },
+            },
+          ],
+        },
+      };
+      const input2: BatchWriteCommandInput = {
+        RequestItems: {
+          [TABLE_NAME]: [
+            {
+              DeleteRequest: {
+                Key: { pk: 'test_item#foo2', sk: '#meta' },
+              },
+            },
+          ],
+        },
+      };
+
+      ddbClientMock.on(BatchWriteCommand).resolves({
+        $metadata: {
+          httpStatusCode: 200,
+        },
+        UnprocessedItems: {
+          [TABLE_NAME]: [
+            {
+              DeleteRequest: {
+                Key: { pk: 'test_item#foo2', sk: '#meta' },
+              },
+            },
+          ],
+        },
+      });
+
+      const requestItems = [{ name: 'foo' }, { name: 'foo2' }];
+      await expect(repository.deleteItems(requestItems)).rejects.toEqual(
+        new Error('Maximum allowed retries exceeded for unprocessed items'),
+      );
+      expect(ddbClientMock).toHaveReceivedNthCommandWith(1, BatchWriteCommand, input1);
+      expect(ddbClientMock).toHaveReceivedNthCommandWith(2, BatchWriteCommand, input2);
+      expect(ddbClientMock).toHaveReceivedNthCommandWith(3, BatchWriteCommand, input2);
+      expect(ddbClientMock).toHaveReceivedNthCommandWith(4, BatchWriteCommand, input2);
+    });
+
+    it('should request batchWrite in batch of 25 items to get result', async () => {
+      ddbClientMock.on(BatchWriteCommand).resolves({
+        $metadata: {
+          httpStatusCode: 200,
+        },
+      });
+
+      const requestItems = [];
+      for (let i = 0; i < 30; i++) {
+        requestItems.push({ name: `foo${i}` });
+      }
+      // keys for first batch request
+      const keys1 = [];
+      for (let i = 0; i < 25; i++) {
+        keys1.push({ pk: `test_item#foo${i}`, sk: '#meta' });
+      }
+      // keys for second batch request
+      const keys2 = [];
+      for (let i = 25; i < 30; i++) {
+        keys2.push({ pk: `test_item#foo${i}`, sk: '#meta' });
+      }
+
+      const input1: BatchWriteCommandInput = {
+        RequestItems: {
+          [TABLE_NAME]: keys1.map((key) => {
+            return {
+              DeleteRequest: {
+                Key: key,
+              },
+            };
+          }),
+        },
+      };
+      const input2: BatchWriteCommandInput = {
+        RequestItems: {
+          [TABLE_NAME]: keys2.map((key) => {
+            return {
+              DeleteRequest: {
+                Key: key,
+              },
+            };
+          }),
+        },
+      };
+
+      await repository.deleteItems(requestItems);
+      expect(ddbClientMock).toHaveReceivedCommandTimes(BatchWriteCommand, 2);
+      expect(ddbClientMock).toHaveReceivedNthCommandWith(1, BatchWriteCommand, input1);
+      expect(ddbClientMock).toHaveReceivedNthCommandWith(2, BatchWriteCommand, input2);
+    });
+
+    it('should throw error if any input item does not includes key field(s)', async () => {
+      const requestItems = [{ name: 'foo' }, { age: 22 }];
+      await expect(repository.deleteItems(requestItems)).rejects.toEqual(
+        new MissingKeyValuesError('Key field "name" must be specified in the input item in entity test-item-entity'),
+      );
+    });
+  });
+
   describe('queryItems()', () => {
     it('should return the items if found', async () => {
       ddbClientMock.on(QueryCommand).resolves({
@@ -1038,16 +1449,18 @@ describe('AbstractRepository', () => {
               ConditionExpression: 'attribute_exists(pk)',
               ExpressionAttributeNames: {
                 '#age': 'age',
+                '#gsi1_sk': 'gsi1_sk',
               },
               ExpressionAttributeValues: {
                 ':age': 55,
+                ':gsi1_sk': 'age#55',
               },
               Key: {
                 pk: 'test_item#foo2',
                 sk: '#meta',
               },
               TableName: 'test-table',
-              UpdateExpression: 'SET #age = :age',
+              UpdateExpression: 'SET #age = :age, #gsi1_sk = :gsi1_sk',
             },
           },
           {

package/src/dynamodb/AbstractDynamoDbRepository.ts

@@ -4,6 +4,10 @@ import {
   UpdateCommandOutput,
   PutCommandInput,
   DeleteCommandInput,
+  BatchWriteCommandInput,
+  BatchWriteCommandOutput,
+  BatchGetCommandInput,
+  BatchGetCommandOutput,
 } from '@aws-sdk/lib-dynamodb';
 
 import { Transaction, DynamoDbManager, MissingKeyValuesError, InvalidDbSchemaError } from '..';
@@ -30,14 +34,15 @@ export type QueryOptions = {
 
 interface Reader<T> {
   queryItems(partialItem: Partial<T>, options?: QueryOptions): Promise<T[]>;
-
   getItem(id: string | Partial<T>): Promise<T | undefined>;
+  getItems(partialItem: Partial<T>[]): Promise<T[]>;
 }
 
 interface Writer<T> {
   createItem(item: Partial<T>): Promise<T>;
   updateItem(partialItem: Partial<T>): Promise<T | undefined>;
   deleteItem(partialItem: Partial<T>): Promise<number>;
+  deleteItems(partialItem: Partial<T>[]): Promise<void>;
 }
 
 type Repository<T> = Reader<T> & Writer<T>;
@@ -65,6 +70,8 @@ export type EntityDefinition = {
   fieldsAsJsonString: string[];
 };
 
+const MAX_REATTEMPTS = 3;
+
 export abstract class AbstractDynamoDbRepository<SHAPE extends object, DATA_CLASS extends SHAPE>
   implements Reader<SHAPE>, Writer<SHAPE>
 {
@@ -153,25 +160,84 @@ export abstract class AbstractDynamoDbRepository<SHAPE extends object, DATA_CLAS
    * @returns
    */
   private async getBatchItems(items: Partial<SHAPE>[]): Promise<DATA_CLASS[]> {
-    const keys: Record<string, string>[] = [];
+    let resultItems: DATA_CLASS[] = [];
+
+    let requestKeys: BatchGetCommandInput['RequestItems'] = {
+      [this.tableName]: {
+        Keys: this.getBatchKeys(items),
+      },
+    };
+    let reattemptsCount = 0;
+    while (requestKeys && Object.keys(requestKeys).length) {
+      if (reattemptsCount++ > MAX_REATTEMPTS) {
+        throw Error('Maximum allowed retries exceeded for unprocessed items');
+      }
+      const output: BatchGetCommandOutput = await this.client.batchGet({
+        RequestItems: requestKeys,
+      });
+      requestKeys = output.UnprocessedKeys;
+      if (output.Responses && output.Responses[this.tableName] && output.Responses[this.tableName].length) {
+        resultItems = resultItems.concat(output.Responses[this.tableName].map((i) => this.hydrateItem(i)));
+      }
+    }
+
+    return resultItems;
+  }
+
+  /**
+   * Delete items in a batch
+   * Uses batchWrite() with 25 items
+   *
+   * @param item
+   *
+   * @throws MissingKeyValuesError
+   */
+  public async deleteItems(items: Partial<SHAPE>[]): Promise<void> {
+    // this is the maximum items allowed by BatchWriteItem()
+    const batchSize = 25;
+
+    for (let i = 0; i < items.length; i += batchSize) {
+      const keys = this.getBatchKeys(items.slice(i, i + batchSize));
+      await this.deleteBatchItems(keys);
+    }
+  }
+
+  private async deleteBatchItems(keys: { [key: string]: string }[]): Promise<void> {
+    let requestItems: BatchWriteCommandInput['RequestItems'] = {
+      [this.tableName]: Object.values(keys).map((key) => {
+        return {
+          DeleteRequest: {
+            Key: key,
+          },
+        };
+      }),
+    };
+    let reattemptsCount = 0;
+    while (requestItems && Object.keys(requestItems).length) {
+      if (reattemptsCount++ > MAX_REATTEMPTS) {
+        throw Error('Maximum allowed retries exceeded for unprocessed items');
+      }
+      const response: BatchWriteCommandOutput = await this.client.batchWrite({
+        RequestItems: requestItems,
+      });
+      requestItems = response.UnprocessedItems;
+    }
+  }
+
+  private getBatchKeys(items: Partial<SHAPE>[]) {
+    const keys: { [key: string]: string }[] = [];
     for (const item of items) {
       keys.push({
         [this.keys.pk.attributeName]: this.getPk(item),
         [this.keys.sk.attributeName]: this.getSk(item),
       });
     }
-    const output = await this.client.batchGet({
-      RequestItems: {
-        [this.tableName]: {
-          Keys: keys,
-        },
-      },
-    });
+    // keys.push({
+    //   [this.keys.pk.attributeName]: 'foo1',
+    //   [this.keys.sk.attributeName]: 'foo2',
+    // });
 
-    if (output.Responses && output.Responses[this.tableName].length) {
-      return output.Responses[this.tableName].map((i) => this.hydrateItem(i));
-    }
-    return [];
+    return keys;
   }
 
   /**
@@ -264,12 +330,15 @@ export abstract class AbstractDynamoDbRepository<SHAPE extends object, DATA_CLAS
     this.assertValueMatchesModel(value);
 
     this.convertSelectedValuesToJsonString(newValue);
+    this.convertSelectedValuesToJsonString(oldValue as Record<string, unknown>);
 
     const updateExpression = [];
     const expressionAttributeNames: Record<string, string> = {};
     const expressionAttributeValues: Record<string, unknown> = {};
+    const updatedAttributes: string[] = [];
     for (const modelProperty of Object.keys(newValue)) {
       const propValue = newValue[modelProperty as keyof SHAPE] ?? null;
+
       if (propValue === oldValue[modelProperty as keyof SHAPE]) {
         // don't need to update the properties that are unchanged
         continue;
@@ -281,19 +350,35 @@ export abstract class AbstractDynamoDbRepository<SHAPE extends object, DATA_CLAS
       updateExpression.push(`${propName} = ${propValuePlaceHolder}`);
       expressionAttributeNames[propName] = modelProperty;
       expressionAttributeValues[propValuePlaceHolder] = propValue;
+      updatedAttributes.push(modelProperty);
     }
-    if (!updateExpression.length) {
+    if (!updatedAttributes.length) {
       // nothing to update
       return value;
     }
 
+    // also update the gsi attributes if needed
+    Object.keys(this.indexes).forEach((key) => {
+      const index = this.indexes[key];
+      [index.pk.attributeName, index.sk.attributeName].forEach((keyAttributeName) => {
+        const keyFormat = this.keysFormat[keyAttributeName];
+        if (updatedAttributes.find((attr) => keyFormat.search(`{${attr}}`) !== -1)) {
+          const propName = `#${keyAttributeName}`;
+          const propValuePlaceHolder = `:${keyAttributeName}`;
+          updateExpression.push(`${propName} = ${propValuePlaceHolder}`);
+          expressionAttributeNames[propName] = keyAttributeName;
+          expressionAttributeValues[propValuePlaceHolder] = this.getKey(value, keyAttributeName);
+        }
+      });
+    });
+
     const updateCommandInput = {
       TableName: this.tableName,
       Key: {
         [this.keys.pk.attributeName]: this.getPk(newValue),
         [this.keys.sk.attributeName]: this.getSk(newValue),
       },
-      UpdateExpression: 'SET ' + updateExpression.join(','),
+      UpdateExpression: 'SET ' + updateExpression.join(', '),
       ExpressionAttributeValues: expressionAttributeValues,
       ExpressionAttributeNames: expressionAttributeNames,
       ConditionExpression: `attribute_exists(${this.keys.pk.attributeName})`,
@@ -545,17 +630,6 @@ export abstract class AbstractDynamoDbRepository<SHAPE extends object, DATA_CLAS
     return keyFormat;
   }
 
-  /**
-   * Whether the given property name is part of the entity's pk/sk string
-   * @param propertyName
-   * @returns boolean
-   */
-  private isPropertyPartOfKeys(propertyName: string) {
-    if (this.keysFormat[this.keys.pk.attributeName].search(`{${propertyName}}`) !== -1) return true;
-    if (this.keysFormat[this.keys.sk.attributeName].search(`{${propertyName}}`) !== -1) return true;
-    return false;
-  }
-
   /**
    * Validate the data matches with "DATA_MODEL"
    * @param value