@squiz/db-lib 1.71.3 → 1.73.0

package/src/dynamodb/AbstractDynamoDbRepository.ts

@@ -4,20 +4,45 @@ import {
   UpdateCommandOutput,
   PutCommandInput,
   DeleteCommandInput,
+  BatchWriteCommandInput,
+  BatchWriteCommandOutput,
+  BatchGetCommandInput,
+  BatchGetCommandOutput,
 } from '@aws-sdk/lib-dynamodb';
 
 import { Transaction, DynamoDbManager, MissingKeyValuesError, InvalidDbSchemaError } from '..';
 import { InvalidDataFormatError } from '../error/InvalidDataFormatError';
 
+export type QueryFilterTypeBeginsWith = {
+  type: 'begins_with';
+  keyword: string;
+};
+
+export type QueryOptions = {
+  // whether to match sort key along with the partition key
+  // if set to true this will return 1 item maximum
+  useSortKey?: boolean;
+  // table index to use (one of tables GSIs)
+  index?: keyof TableIndexes;
+  // number of items to limit in the result
+  limit?: number;
+  // result order based on the sort key
+  order?: 'desc' | 'asc';
+  // filter operation on the sort key
+  filter?: QueryFilterTypeBeginsWith;
+};
+
 interface Reader<T> {
-  queryItems(partialItem: Partial<T>, useSortKey?: boolean, index?: keyof TableIndexes): Promise<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>, newValue: Partial<T>): Promise<T | undefined>;
+  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>;
@@ -45,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>
 {
@@ -104,20 +131,127 @@ export abstract class AbstractDynamoDbRepository<SHAPE extends object, DATA_CLAS
     return this.hydrateItem(output.Item);
   }
 
+  /**
+   * Get the single item each matching the key fields value in the given
+   * partial items. Will throw MissingKeyValuesError if key field values
+   * are missing
+   * Uses batchGet() to request 100 items in a batch
+   *
+   * @param item
+   *
+   * @throws MissingKeyValuesError
+   */
+  public async getItems(items: Partial<SHAPE>[]): Promise<DATA_CLASS[]> {
+    // this is the maximum items allowed by BatchGetItem()
+    const batchSize = 100;
+
+    let result: DATA_CLASS[] = [];
+    for (let i = 0; i < items.length; i += batchSize) {
+      const batchResult = await this.getBatchItems(items.slice(i, i + batchSize));
+      result = result.concat(batchResult);
+    }
+    return result;
+  }
+
+  /**
+   * Returns the batch items from the items primary key
+   *
+   * @param items
+   * @returns
+   */
+  private async getBatchItems(items: Partial<SHAPE>[]): Promise<DATA_CLASS[]> {
+    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),
+      });
+    }
+    // keys.push({
+    //   [this.keys.pk.attributeName]: 'foo1',
+    //   [this.keys.sk.attributeName]: 'foo2',
+    // });
+
+    return keys;
+  }
+
   /**
    * Finds all the items matching the partition key or
    * the gsi key (when gsi index name is specified)
    *
    * @param item
-   * @param useSortKey
-   * @param index
+   * @param options
    * @throws MissingKeyValuesError
    */
-  public async queryItems(
-    item: Partial<SHAPE>,
-    useSortKey: boolean = false,
-    index?: keyof TableIndexes,
-  ): Promise<DATA_CLASS[]> {
+  public async queryItems(item: Partial<SHAPE>, options?: QueryOptions): Promise<DATA_CLASS[]> {
+    const useSortKey = options?.useSortKey || options?.filter;
+    const index = options?.index;
+
     let pkName = this.keys.pk.attributeName;
     let skName = this.keys.sk.attributeName;
     let indexName = null;
@@ -136,8 +270,8 @@ export abstract class AbstractDynamoDbRepository<SHAPE extends object, DATA_CLAS
     const expressionAttributeNames: Record<string, string> = { '#pkName': pkName };
     const expressionAttributeValues: Record<string, unknown> = { ':pkValue': pk };
     if (useSortKey) {
-      const sk = this.getKey(item, skName);
-      keyConditionExpression.push('#skName = :skValue');
+      const sk = options?.filter?.keyword ?? this.getKey(item, skName);
+      keyConditionExpression.push(this.getFilterKeyConditionExpression(options?.filter));
       expressionAttributeNames['#skName'] = skName;
       expressionAttributeValues[':skValue'] = sk;
     }
@@ -151,27 +285,43 @@ export abstract class AbstractDynamoDbRepository<SHAPE extends object, DATA_CLAS
     if (indexName) {
       queryCommandInput['IndexName'] = String(indexName);
     }
-    const output = await this.client.query(queryCommandInput);
 
+    if (options?.order !== undefined) {
+      queryCommandInput.ScanIndexForward = options.order === 'asc';
+    }
+    if (options?.limit !== undefined) {
+      queryCommandInput.Limit = options.limit;
+    }
+
+    const output = await this.client.query(queryCommandInput);
     return !output.Items ? [] : output.Items.map((item) => this.hydrateItem(item));
   }
 
+  /**
+   * Evaluate filter condition for sort key
+   * @param filter
+   * @returns string
+   */
+  private getFilterKeyConditionExpression(filter: QueryOptions['filter']): string {
+    if (filter === undefined) {
+      return '#skName = :skValue';
+    } else if (filter.type === 'begins_with') {
+      return 'begins_with(#skName, :skValue)';
+    }
+    throw new Error(`Invalid query filter type: ${(filter as any)?.type}`);
+  }
+
   /**
    * Update the existing item matching the key fields value
-   * in the passed in partialItem
-   * @param partialItem
+   * in the passed newValue
    * @param newValue
    * @param transaction
    *
    * @returns Promise<SHAPE | undefined>
    * @throws MissingKeyValuesError
    */
-  public async updateItem(
-    partialItem: Partial<SHAPE>,
-    newValue: Exclude<Partial<SHAPE>, Record<string, never>>,
-    transaction: Transaction = {},
-  ): Promise<DATA_CLASS | undefined> {
-    const oldValue = await this.getItem(partialItem);
+  public async updateItem(newValue: Partial<SHAPE>, transaction: Transaction = {}): Promise<DATA_CLASS | undefined> {
+    const oldValue = await this.getItem(newValue);
     if (oldValue === undefined) {
       return undefined;
     }
@@ -185,20 +335,29 @@ export abstract class AbstractDynamoDbRepository<SHAPE extends object, DATA_CLAS
     const expressionAttributeNames: Record<string, string> = {};
     const expressionAttributeValues: Record<string, unknown> = {};
     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;
+      }
+
       const propName = `#${modelProperty}`;
-      const propValue = `:${modelProperty}`;
+      const propValuePlaceHolder = `:${modelProperty}`;
 
-      updateExpression.push(`${propName} = ${propValue}`);
+      updateExpression.push(`${propName} = ${propValuePlaceHolder}`);
       expressionAttributeNames[propName] = modelProperty;
-
-      expressionAttributeValues[propValue] = newValue[modelProperty as keyof SHAPE] ?? null;
+      expressionAttributeValues[propValuePlaceHolder] = propValue;
+    }
+    if (!updateExpression.length) {
+      // nothing to update
+      return value;
     }
 
     const updateCommandInput = {
       TableName: this.tableName,
       Key: {
-        [this.keys.pk.attributeName]: this.getPk(partialItem),
-        [this.keys.sk.attributeName]: this.getSk(partialItem),
+        [this.keys.pk.attributeName]: this.getPk(newValue),
+        [this.keys.sk.attributeName]: this.getSk(newValue),
       },
       UpdateExpression: 'SET ' + updateExpression.join(','),
       ExpressionAttributeValues: expressionAttributeValues,
@@ -240,12 +399,17 @@ export abstract class AbstractDynamoDbRepository<SHAPE extends object, DATA_CLAS
    *
    * @param value
    * @param transaction
+   * @param additionalValue Additional item properties that are not part of the DATA_CLASS
    *
    * @returns Promise<SHAPE>
    * @throws DuplicateItemError
    * @throws MissingKeyValuesError
    */
-  public async createItem(value: DATA_CLASS, transaction: Transaction = {}): Promise<DATA_CLASS> {
+  public async createItem(
+    value: DATA_CLASS,
+    transaction: Transaction = {},
+    additionalValue: Partial<SHAPE> = {},
+  ): Promise<DATA_CLASS> {
     this.assertValueMatchesModel(value);
 
     const columns: any = {};
@@ -255,14 +419,14 @@ export abstract class AbstractDynamoDbRepository<SHAPE extends object, DATA_CLAS
     this.convertSelectedValuesToJsonString(columns);
 
     const keyFields: Record<string, unknown> = {
-      [this.keys.pk.attributeName]: this.getPk(value),
-      [this.keys.sk.attributeName]: this.getSk(value),
+      [this.keys.pk.attributeName]: this.getPk({ ...value, ...additionalValue }),
+      [this.keys.sk.attributeName]: this.getSk({ ...value, ...additionalValue }),
     };
 
     Object.keys(this.indexes).forEach((key) => {
       const index = this.indexes[key];
-      keyFields[index.pk.attributeName] = this.getKey(value, index.pk.attributeName);
-      keyFields[index.sk.attributeName] = this.getKey(value, index.sk.attributeName);
+      keyFields[index.pk.attributeName] = this.getKey({ ...value, ...additionalValue }, index.pk.attributeName);
+      keyFields[index.sk.attributeName] = this.getKey({ ...value, ...additionalValue }, index.sk.attributeName);
     });
 
     const putCommandInput: PutCommandInput = {
@@ -447,6 +611,17 @@ 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