@squiz/db-lib 1.77.0 → 1.77.2

package/src/dynamodb/AbstractDynamoDbRepository.spec.ts

@@ -1465,7 +1465,11 @@ describe('AbstractRepository', () => {
           httpStatusCode: 200,
         },
       });
+
+      // Limit does not appear in the input query.
+      // Splitting/pagination is now handled in the queryFullPartition sub function.
       const input: QueryCommandInput = {
+        ExclusiveStartKey: undefined,
         TableName: TABLE_NAME,
         KeyConditionExpression: '#pkName = :pkValue',
         ExpressionAttributeNames: {
@@ -1474,7 +1478,6 @@ describe('AbstractRepository', () => {
         ExpressionAttributeValues: {
           ':pkValue': 'test_item#foo',
         },
-        Limit: 33,
       };
 
       const partialItem = { name: 'foo' };
@@ -1563,6 +1566,181 @@ describe('AbstractRepository', () => {
         new MissingKeyValuesError(`Key field "age" must be specified in the input item in entity test-item-entity`),
       );
     });
+
+    describe('dynamo pagination', () => {
+      const mockItems = [
+        {
+          name: 'foo',
+          age: 10,
+          country: 'au',
+          data: {},
+        },
+        {
+          name: 'fox',
+          age: 11,
+          country: 'au',
+          data: {},
+        },
+        {
+          name: 'bar',
+          age: 12,
+          country: 'au',
+          data: {},
+        },
+        {
+          name: 'baz',
+          age: 13,
+          country: 'au',
+          data: {},
+        },
+        {
+          name: 'qux',
+          age: 14,
+          country: 'au',
+          data: {},
+        },
+      ];
+
+      it('should return all items when there is only 1 page of results', async () => {
+        ddbClientMock.on(QueryCommand).resolves({
+          Items: [mockItems[0], mockItems[1], mockItems[2]],
+        });
+
+        const input: QueryCommandInput = {
+          TableName: TABLE_NAME,
+          KeyConditionExpression: '#pkName = :pkValue',
+          ExpressionAttributeNames: {
+            '#pkName': 'pk',
+          },
+          ExpressionAttributeValues: {
+            ':pkValue': 'test_item#foo',
+          },
+        };
+
+        const partialItem = { name: 'foo' };
+        const result = await repository.queryItems(partialItem);
+        expect(ddbClientMock).toHaveReceivedCommandTimes(QueryCommand, 1);
+        expect(ddbClientMock).toHaveReceivedCommandWith(QueryCommand, input);
+        expect(result).toEqual([new TestItem(mockItems[0]), new TestItem(mockItems[1]), new TestItem(mockItems[2])]);
+      });
+
+      it('should return all items when there are more than 1 page of results', async () => {
+        const mockLastEvaluatedKey = { pk: { S: 'test_item#foo' }, sk: { S: 'name#bar' } };
+
+        ddbClientMock
+          .on(QueryCommand)
+          .resolvesOnce({
+            Items: [mockItems[0], mockItems[1], mockItems[2]],
+            LastEvaluatedKey: mockLastEvaluatedKey,
+          })
+          .resolvesOnce({
+            Items: [mockItems[3], mockItems[4]],
+          });
+
+        const input: QueryCommandInput = {
+          TableName: TABLE_NAME,
+          KeyConditionExpression: '#pkName = :pkValue',
+          ExpressionAttributeNames: {
+            '#pkName': 'pk',
+          },
+          ExpressionAttributeValues: {
+            ':pkValue': 'test_item#foo',
+          },
+        };
+
+        const partialItem = { name: 'foo' };
+        const result = await repository.queryItems(partialItem);
+        expect(ddbClientMock).toHaveReceivedCommandTimes(QueryCommand, 2);
+        expect(ddbClientMock).toHaveReceivedNthCommandWith(1, QueryCommand, input);
+        expect(ddbClientMock).toHaveReceivedNthCommandWith(2, QueryCommand, {
+          ...input,
+          ExclusiveStartKey: mockLastEvaluatedKey,
+        });
+        expect(result).toEqual([
+          new TestItem(mockItems[0]),
+          new TestItem(mockItems[1]),
+          new TestItem(mockItems[2]),
+          new TestItem(mockItems[3]),
+          new TestItem(mockItems[4]),
+        ]);
+      });
+
+      it('should return items but only up to the provided limit when limit is within first page', async () => {
+        const mockLimit = 2;
+        const mockLastEvaluatedKey = { pk: { S: 'test_item#foo' }, sk: { S: 'name#bar' } };
+
+        ddbClientMock
+          .on(QueryCommand)
+          .resolvesOnce({
+            Items: [mockItems[0], mockItems[1], mockItems[2]],
+            LastEvaluatedKey: mockLastEvaluatedKey,
+          })
+          .resolvesOnce({
+            Items: [mockItems[3], mockItems[4]],
+          });
+
+        const input: QueryCommandInput = {
+          TableName: TABLE_NAME,
+          KeyConditionExpression: '#pkName = :pkValue',
+          ExpressionAttributeNames: {
+            '#pkName': 'pk',
+          },
+          ExpressionAttributeValues: {
+            ':pkValue': 'test_item#foo',
+          },
+        };
+
+        // Only expect first two items from the first page of results.
+        const partialItem = { name: 'foo' };
+        const result = await repository.queryItems(partialItem, { limit: mockLimit });
+        expect(ddbClientMock).toHaveReceivedCommandTimes(QueryCommand, 1);
+        expect(ddbClientMock).toHaveReceivedNthCommandWith(1, QueryCommand, input);
+        expect(result).toEqual([new TestItem(mockItems[0]), new TestItem(mockItems[1])]);
+      });
+
+      it('should return items up to the limit when limit spans multiple pages', async () => {
+        const mockLimit = 4;
+        const mockLastEvaluatedKey = { pk: { S: 'test_item#foo' }, sk: { S: 'name#bar' } };
+
+        ddbClientMock
+          .on(QueryCommand)
+          .resolvesOnce({
+            Items: [mockItems[0], mockItems[1], mockItems[2]],
+            LastEvaluatedKey: mockLastEvaluatedKey,
+          })
+          .resolvesOnce({
+            Items: [mockItems[3], mockItems[4]],
+          });
+
+        const input: QueryCommandInput = {
+          TableName: TABLE_NAME,
+          KeyConditionExpression: '#pkName = :pkValue',
+          ExpressionAttributeNames: {
+            '#pkName': 'pk',
+          },
+          ExpressionAttributeValues: {
+            ':pkValue': 'test_item#foo',
+          },
+        };
+
+        // Expect only 4 results in total across both pages.
+        // Expect all 3 results from first page, but only 1 result from second page.
+        const partialItem = { name: 'foo' };
+        const result = await repository.queryItems(partialItem, { limit: mockLimit });
+        expect(ddbClientMock).toHaveReceivedCommandTimes(QueryCommand, 2);
+        expect(ddbClientMock).toHaveReceivedNthCommandWith(1, QueryCommand, input);
+        expect(ddbClientMock).toHaveReceivedNthCommandWith(2, QueryCommand, {
+          ...input,
+          ExclusiveStartKey: mockLastEvaluatedKey,
+        });
+        expect(result).toEqual([
+          new TestItem(mockItems[0]),
+          new TestItem(mockItems[1]),
+          new TestItem(mockItems[2]),
+          new TestItem(mockItems[3]),
+        ]);
+      });
+    });
   });
 
   describe('deleteItem()', () => {

package/src/dynamodb/AbstractDynamoDbRepository.ts

@@ -304,12 +304,8 @@ export abstract class AbstractDynamoDbRepository<SHAPE extends object, DATA_CLAS
     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));
+    return await this.queryFullPartition(queryCommandInput, options?.limit);
   }
 
   /**
@@ -326,6 +322,44 @@ export abstract class AbstractDynamoDbRepository<SHAPE extends object, DATA_CLAS
     throw new Error(`Invalid query filter type: ${(filter as any)?.type}`);
   }
 
+  /**
+   * Get the list of objects by given params, paginating until all results are returned for a full partition.
+   * @param params - The DynamoDB query params.
+   * @param limit - The maximum number of items to return.
+   * @returns Promise of array of items.
+   */
+  private async queryFullPartition(params: QueryCommandInput, limit?: number): Promise<DATA_CLASS[]> {
+    const items: Array<DATA_CLASS> = [];
+    let exclusiveStartKey: QueryCommandInput['ExclusiveStartKey'] = undefined;
+
+    do {
+      const queryInput: QueryCommandInput = {
+        ...params,
+        ExclusiveStartKey: exclusiveStartKey,
+      };
+
+      const result = await this.client.query(queryInput);
+
+      exclusiveStartKey = result.LastEvaluatedKey;
+
+      if (result.Items && result.Items.length) {
+        const resultItems = result.Items.map((item) => this.hydrateItem(item));
+
+        const itemsToAdd = limit === undefined ? resultItems : resultItems.slice(0, limit - items.length);
+
+        items.push(...itemsToAdd);
+
+        const hasReachedLimit = limit !== undefined && items.length >= limit;
+
+        if (hasReachedLimit) {
+          exclusiveStartKey = undefined;
+        }
+      }
+    } while (exclusiveStartKey);
+
+    return items;
+  }
+
   /**
    * Update the existing item matching the key fields value
    * in the passed newValue