DynamoDB Toolbox Get All Query Pages Typescript Serverless Lambda

getAllQueryPages.ts

/** Usage **/
// const results = await getAllQueryPages(
//   SomeModel.query("MY_PK", {
//     index: GSI.GSI1.IndexName,
//   })
// );


/** getAllQueryPages.ts **/
import { DocumentClient } from "aws-sdk/lib/dynamodb/document_client";
import { Replace } from "dynamodb-toolbox";

// DynamoDB queries have a 1Mo limit, after that there is a pagination system
// This function fetch all available pages
export const getAllQueryPages = async <T>(
  dynamoDbQuery: Promise<Replace<DocumentClient.QueryOutput, "Items", T[]>>
): Promise<T[]> => {
  let currentQueryResult = await dynamoDbQuery;
  let results = currentQueryResult.Items;
  let iteration = 0;

  // Limit to prevent infinite loop
  const MAX_QUERY_LIMIT = 10;
  while (currentQueryResult.LastEvaluatedKey !== undefined) {
    if (iteration > MAX_QUERY_LIMIT) {
      throw Error("Paginating over to many queries!");
    }
    // @ts-expect-error the next() function is badly typed in dynamodb toolbox but it exists
    // eslint-disable-next-line @typescript-eslint/no-unsafe-call
    currentQueryResult = await currentQueryResult.next();
    results = results.concat(currentQueryResult.Items);
    iteration += 1;
  }

  return results;
};

  

/** getAllQueryPages.test.ts **/
import { DocumentClient } from "aws-sdk/lib/dynamodb/document_client";
import { Replace } from "dynamodb-toolbox";
import { getAllQueryPages } from "./getAllQueryPages";

type TestQueryType = Promise<
  Replace<DocumentClient.QueryOutput, "Items", number[]>
>;

describe("getAllQueryPages", () => {
  it("should not paginate if at the end of request", async () => {
    const queryMock: TestQueryType = Promise.resolve({ Items: [1] });

    const results = await getAllQueryPages(queryMock);
    expect(results).toEqual([1]);
  });

  it("should paginate if more than one page", async () => {
    const queryMock: TestQueryType = Promise.resolve({
      Items: [1],
      LastEvaluatedKey: { test1: "value1" },
      next: () =>
        Promise.resolve({
          Items: [2],
          LastEvaluatedKey: { test2: "value2" },
          next: () => Promise.resolve({ Items: [3] }),
        }),
    });

    const results = await getAllQueryPages(queryMock);
    expect(results).toEqual([1, 2, 3]);
  });

  it("should not do infinite pagination", async () => {
    const getNextQuery = (deep: number) => () =>
      Promise.resolve({
        Items: [2],
        LastEvaluatedKey: { test2: "value2" },
        next: deep < 100 ? getNextQuery(deep + 1) : () => Promise.resolve(),
      });

    const queryMock: TestQueryType = Promise.resolve({
      Items: [1],
      LastEvaluatedKey: { test1: "value1" },
      next: getNextQuery(0),
    });

    await expect(() => getAllQueryPages(queryMock)).rejects.toBeInstanceOf(
      Error
    );
  });
});