From f8d5a04eae16a170ecea43ddddf9583fdd0cf790 Mon Sep 17 00:00:00 2001 From: Ciprian Serbu Date: Wed, 26 Nov 2025 15:03:36 +0200 Subject: [PATCH 1/5] feat(dynamodb): add batch get item paginators for low level and document client --- .../codegen/AddDocumentClientPlugin.java | 1 + .../src/pagination/BatchGetPaginator.ts | 48 +++++++++++++++++++ lib/lib-dynamodb/src/pagination/index.ts | 1 + packages/util-dynamodb/src/index.ts | 1 + .../util-dynamodb/src/paginateBatchGetItem.ts | 40 ++++++++++++++++ 5 files changed, 91 insertions(+) create mode 100644 lib/lib-dynamodb/src/pagination/BatchGetPaginator.ts create mode 100644 packages/util-dynamodb/src/paginateBatchGetItem.ts diff --git a/codegen/smithy-aws-typescript-codegen/src/main/java/software/amazon/smithy/aws/typescript/codegen/AddDocumentClientPlugin.java b/codegen/smithy-aws-typescript-codegen/src/main/java/software/amazon/smithy/aws/typescript/codegen/AddDocumentClientPlugin.java index 4b910a0afa9ea..8668f53c87279 100644 --- a/codegen/smithy-aws-typescript-codegen/src/main/java/software/amazon/smithy/aws/typescript/codegen/AddDocumentClientPlugin.java +++ b/codegen/smithy-aws-typescript-codegen/src/main/java/software/amazon/smithy/aws/typescript/codegen/AddDocumentClientPlugin.java @@ -105,6 +105,7 @@ private void writeAdditionalFiles( writerFactory.accept(String.format("%s%s/index.ts", DocumentClientUtils.DOC_CLIENT_PREFIX, DocumentClientPaginationGenerator.PAGINATION_FOLDER), writer -> { writer.write("export * from './Interfaces';"); + writer.write("export * from './BatchGetPaginator';"); for (OperationShape operation : overridenOperationsList) { if (operation.hasTrait(PaginatedTrait.ID)) { String paginationFileName = diff --git a/lib/lib-dynamodb/src/pagination/BatchGetPaginator.ts b/lib/lib-dynamodb/src/pagination/BatchGetPaginator.ts new file mode 100644 index 0000000000000..01fa8a120e4ac --- /dev/null +++ b/lib/lib-dynamodb/src/pagination/BatchGetPaginator.ts @@ -0,0 +1,48 @@ +import { type DynamoDBClient } from "@aws-sdk/client-dynamodb"; + +import { BatchGetCommand, BatchGetCommandInput, BatchGetCommandOutput } from "../commands/BatchGetCommand"; +import { DynamoDBDocumentClient } from "../DynamoDBDocumentClient"; + +/** + * @public + * + * Async generator that issues {@link BatchGetCommand}s repeatedly until all keys are processed or an error response is received. + * + * @see {@link paginateBatchGetItem} for a variant that uses the {@link DynamoDBClient | low level DynamoDB client}. + * + * @example + * + * ``` + * const client = new DynamoDBClient(); + * const docClient = DynamoDBDocumentClient.from(client); + * const input: BatchGetCommandInput = { + * RequestItems: { + * table1: Keys: [...], + * table2: Keys: [...], + * } + * }; + * + * let pageNumber = 1; + * for await (const page of paginateBatchGet({ client: docClient }, input)) { + * console.log("page:", pageNumber++); + * console.log("items:", page.Responses); + * console.log("unprocessed:", page.UnprocessedKeys); // will be returned in the next page(s) + * } + * ``` + */ +export async function* paginateBatchGet( + config: { + client: DynamoDBDocumentClient; + }, + input: BatchGetCommandInput +): AsyncGenerator { + let RequestItems = input.RequestItems; + + while (RequestItems && Object.keys(RequestItems).length > 0) { + const cmd = new BatchGetCommand({ ...input, RequestItems }); + const response = await config.client.send(cmd); + RequestItems = { ...response.UnprocessedKeys }; + + yield response; + } +} diff --git a/lib/lib-dynamodb/src/pagination/index.ts b/lib/lib-dynamodb/src/pagination/index.ts index 844e9cbb4adf4..803dd24fa8217 100644 --- a/lib/lib-dynamodb/src/pagination/index.ts +++ b/lib/lib-dynamodb/src/pagination/index.ts @@ -1,3 +1,4 @@ +export * from "./BatchGetPaginator"; // smithy-typescript generated code export * from "./Interfaces"; export * from "./QueryPaginator"; diff --git a/packages/util-dynamodb/src/index.ts b/packages/util-dynamodb/src/index.ts index 60877561d3349..4ac3c4445f036 100644 --- a/packages/util-dynamodb/src/index.ts +++ b/packages/util-dynamodb/src/index.ts @@ -3,4 +3,5 @@ export * from "./convertToAttr"; export * from "./convertToNative"; export * from "./marshall"; export * from "./models"; +export * from "./paginateBatchGetItem"; export * from "./unmarshall"; diff --git a/packages/util-dynamodb/src/paginateBatchGetItem.ts b/packages/util-dynamodb/src/paginateBatchGetItem.ts new file mode 100644 index 0000000000000..58e0a58572c2e --- /dev/null +++ b/packages/util-dynamodb/src/paginateBatchGetItem.ts @@ -0,0 +1,40 @@ +import { BatchGetItemCommand, BatchGetItemCommandInput, DynamoDBClient } from "@aws-sdk/client-dynamodb"; +import { type DynamoDBDocumentClient } from "@aws-sdk/lib-dynamodb"; + +/** + * @public + * + * Async generator that issues {@link BatchGetItemCommand}s repeatedly until all keys are processed or an error response is received. + * + * @see {@link paginateBatchGetItem} for a variant that uses the {@link DynamoDBDocumentClient | DynamoDB document client}. + * + * @example + * + * ``` + * const client = new DynamoDBClient(); + * const input: BatchGetCommandInput = { + * RequestItems: { + * table1: Keys: [...], + * table2: Keys: [...], + * } + * }; + * + * let pageNumber = 1; + * for await (const page of paginateBatchGetItem({ client }, input)) { + * console.log("page:", pageNumber++); + * console.log("items:", page.Responses); + * console.log("unprocessed:", page.UnprocessedKeys); // will be returned in the next page(s) + * } + * ``` + */ +export async function* paginateBatchGetItem(client: DynamoDBClient, input: BatchGetItemCommandInput) { + let RequestItems = input.RequestItems; + + while (RequestItems && Object.keys(RequestItems).length > 0) { + const cmd = new BatchGetItemCommand({ ...input, RequestItems }); + const response = await client.send(cmd); + RequestItems = { ...response.UnprocessedKeys }; + + yield response; + } +} From c9b28303e215ec8898261e012ff753538006c176 Mon Sep 17 00:00:00 2001 From: Ciprian Serbu Date: Wed, 26 Nov 2025 15:03:36 +0200 Subject: [PATCH 2/5] improve tsdoc --- lib/lib-dynamodb/src/pagination/BatchGetPaginator.ts | 6 +++--- packages/util-dynamodb/src/paginateBatchGetItem.ts | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/lib/lib-dynamodb/src/pagination/BatchGetPaginator.ts b/lib/lib-dynamodb/src/pagination/BatchGetPaginator.ts index 01fa8a120e4ac..7818e836f13b9 100644 --- a/lib/lib-dynamodb/src/pagination/BatchGetPaginator.ts +++ b/lib/lib-dynamodb/src/pagination/BatchGetPaginator.ts @@ -4,15 +4,15 @@ import { BatchGetCommand, BatchGetCommandInput, BatchGetCommandOutput } from ".. import { DynamoDBDocumentClient } from "../DynamoDBDocumentClient"; /** - * @public - * * Async generator that issues {@link BatchGetCommand}s repeatedly until all keys are processed or an error response is received. * + * @public + * * @see {@link paginateBatchGetItem} for a variant that uses the {@link DynamoDBClient | low level DynamoDB client}. * * @example * - * ``` + * ```typescript * const client = new DynamoDBClient(); * const docClient = DynamoDBDocumentClient.from(client); * const input: BatchGetCommandInput = { diff --git a/packages/util-dynamodb/src/paginateBatchGetItem.ts b/packages/util-dynamodb/src/paginateBatchGetItem.ts index 58e0a58572c2e..bc74859a96612 100644 --- a/packages/util-dynamodb/src/paginateBatchGetItem.ts +++ b/packages/util-dynamodb/src/paginateBatchGetItem.ts @@ -2,15 +2,15 @@ import { BatchGetItemCommand, BatchGetItemCommandInput, DynamoDBClient } from "@ import { type DynamoDBDocumentClient } from "@aws-sdk/lib-dynamodb"; /** - * @public - * * Async generator that issues {@link BatchGetItemCommand}s repeatedly until all keys are processed or an error response is received. * + * @public + * * @see {@link paginateBatchGetItem} for a variant that uses the {@link DynamoDBDocumentClient | DynamoDB document client}. * * @example * - * ``` + * ```typescript * const client = new DynamoDBClient(); * const input: BatchGetCommandInput = { * RequestItems: { From 05edaf11413bbea9affad7cbf9f952fa3bee87f1 Mon Sep 17 00:00:00 2001 From: Ciprian Serbu Date: Wed, 26 Nov 2025 15:03:36 +0200 Subject: [PATCH 3/5] add tests (positive and negative) --- packages/core/integ/pagination.integ.spec.ts | 133 +++++++++++++++++++ 1 file changed, 133 insertions(+) diff --git a/packages/core/integ/pagination.integ.spec.ts b/packages/core/integ/pagination.integ.spec.ts index 7b23733d4d460..603ae8fd23e46 100644 --- a/packages/core/integ/pagination.integ.spec.ts +++ b/packages/core/integ/pagination.integ.spec.ts @@ -1,5 +1,6 @@ import { requireRequestsFrom } from "@aws-sdk/aws-util-test/src"; import { DynamoDB, paginateScan, ScanCommandInput } from "@aws-sdk/client-dynamodb"; +import { BatchGetCommandInput, paginateBatchGet } from "@aws-sdk/lib-dynamodb"; import { HttpResponse } from "@smithy/protocol-http"; import { describe, expect, test as it } from "vitest"; @@ -94,4 +95,136 @@ describe("pagination", () => { }); expect.assertions(7); }); + + it("processes batch items until all items are processed or an error is received", async () => { + const ddb = new DynamoDB({ + credentials: { + accessKeyId: "INTEG_TEST", + secretAccessKey: "INTEG_TEST", + }, + region: "us-west-2", + }); + + requireRequestsFrom(ddb) + .toMatch( + // first page request + { + hostname: /dynamodb/, + body(b) { + expect(b).toContain( + '"RequestItems":{"test":{"Keys":[{"id":{"S":"1"}},{"id":{"S":"2"}},{"id":{"S":"3"}},{"id":{"S":"4"}},{"id":{"S":"5"}}]}}' + ); + }, + }, + // second page request + { + hostname: /dynamodb/, + body(b) { + expect(b).toContain('"RequestItems":{"test":{"Keys":[{"id":{"S":"4"}},{"id":{"S":"3"}}]}}'); + }, + }, + // invalid request (duplicate key) + { + hostname: /dynamodb/, + body(b) { + expect(b).toContain('"RequestItems":{"test":{"Keys":[{"id":{"S":"1"}},{"id":{"S":"1"}}]}}'); + }, + } + ) + .respondWith( + // first page response + new HttpResponse({ + statusCode: 200, + headers: {}, + body: Buffer.from( + JSON.stringify({ + Responses: { + test: [ + { id: { S: "2" }, name: { S: "Item 2" } }, + { id: { S: "1" }, name: { S: "Item 1" } }, + ], + }, + UnprocessedKeys: { + test: { + Keys: [{ id: { S: "4" } }, { id: { S: "3" } }], + }, + }, + }) + ), + }), + // second page response + new HttpResponse({ + statusCode: 200, + headers: {}, + body: Buffer.from( + JSON.stringify({ + Responses: { + test: [ + { id: { S: "3" }, name: { S: "Item 3" } }, + { id: { S: "4" }, name: { S: "Item 4" } }, + ], + }, + UnprocessedKeys: {}, + }) + ), + }), + // error response + new HttpResponse({ + statusCode: 400, + headers: {}, + body: Buffer.from( + JSON.stringify({ + message: "Provided list of item keys contains duplicates", + }) + ), + }) + ); + + const requestParams: BatchGetCommandInput = { + RequestItems: { + test: { Keys: [{ id: "1" }, { id: "2" }, { id: "3" }, { id: "4" }, { id: "5" }] }, + }, + }; + + let pages = 0; + for await (const page of paginateBatchGet({ client: ddb }, requestParams)) { + pages += 1; + if (pages === 1) { + expect(page.Responses?.test).toEqual([ + { id: "2", name: "Item 2" }, + { id: "1", name: "Item 1" }, + ]); + } else { + expect(page.Responses?.test).toEqual([ + { id: "3", name: "Item 3" }, + { id: "4", name: "Item 4" }, + ]); + } + } + + expect(pages).toEqual(2); + + let thrownError; + + try { + for await (const page of paginateBatchGet( + { client: ddb }, + { + RequestItems: { + test: { Keys: [{ id: "1" }, { id: "1" }] }, + }, + } + )) { + void page; + throw new Error("Received unexpected page"); + } + } catch (error) { + thrownError = error; + } + + expect(thrownError).toBeInstanceOf(Error); + expect((thrownError as Error).message).toBe("Provided list of item keys contains duplicates"); + + expect.assertions(11); + }); }); From 1af69d052e8e5be8db58620e9afdecc3a40ac548 Mon Sep 17 00:00:00 2001 From: Ciprian Serbu Date: Wed, 26 Nov 2025 15:03:36 +0200 Subject: [PATCH 4/5] explain the tests in a comment --- packages/core/integ/pagination.integ.spec.ts | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/packages/core/integ/pagination.integ.spec.ts b/packages/core/integ/pagination.integ.spec.ts index 603ae8fd23e46..c30ba19b206eb 100644 --- a/packages/core/integ/pagination.integ.spec.ts +++ b/packages/core/integ/pagination.integ.spec.ts @@ -96,6 +96,19 @@ describe("pagination", () => { expect.assertions(7); }); + /** + * This test makes a DynamoDB paginated batch get request for 5 items, with keys 1-2-3-4-5, in this exact order. + * + * The first returned page contains items 2 and 1 (order switched to simulate the unpredictability of the order of the + * items returned by the DDB API BatchGetItem command), plus unprocessed keys 3 and 4. The second page contains the + * items 3 and 4, and no further unprocessed keys. + * + * Item 5 is asked for, but we consider that the table does not contain it, so it's not returned at all. That's a + * valid use case and does not generate an error. + * + * In the second part of the test, another paginated request is done for 2 items, with keys 1 and 1. So the same + * item is requested twice. As the API will return an error, we want to catch the generated SDK exception. + */ it("processes batch items until all items are processed or an error is received", async () => { const ddb = new DynamoDB({ credentials: { From aa4f6aa1152e0d7483c729cc7e50d3e81d7abea7 Mon Sep 17 00:00:00 2001 From: Ciprian Serbu Date: Thu, 27 Nov 2025 13:31:57 +0200 Subject: [PATCH 5/5] rename batch paginator functions to avoid future collisions with their generated versions --- lib/lib-dynamodb/src/pagination/BatchGetPaginator.ts | 4 ++-- packages/core/integ/pagination.integ.spec.ts | 6 +++--- packages/util-dynamodb/src/index.ts | 2 +- .../{paginateBatchGetItem.ts => paginatedBatchGetItem.ts} | 6 +++--- 4 files changed, 9 insertions(+), 9 deletions(-) rename packages/util-dynamodb/src/{paginateBatchGetItem.ts => paginatedBatchGetItem.ts} (78%) diff --git a/lib/lib-dynamodb/src/pagination/BatchGetPaginator.ts b/lib/lib-dynamodb/src/pagination/BatchGetPaginator.ts index 7818e836f13b9..ee1a5e177911b 100644 --- a/lib/lib-dynamodb/src/pagination/BatchGetPaginator.ts +++ b/lib/lib-dynamodb/src/pagination/BatchGetPaginator.ts @@ -8,7 +8,7 @@ import { DynamoDBDocumentClient } from "../DynamoDBDocumentClient"; * * @public * - * @see {@link paginateBatchGetItem} for a variant that uses the {@link DynamoDBClient | low level DynamoDB client}. + * @see {@link paginatedBatchGet} for a variant that uses the {@link DynamoDBClient | low level DynamoDB client}. * * @example * @@ -30,7 +30,7 @@ import { DynamoDBDocumentClient } from "../DynamoDBDocumentClient"; * } * ``` */ -export async function* paginateBatchGet( +export async function* paginatedBatchGet( config: { client: DynamoDBDocumentClient; }, diff --git a/packages/core/integ/pagination.integ.spec.ts b/packages/core/integ/pagination.integ.spec.ts index c30ba19b206eb..b50afaf92e321 100644 --- a/packages/core/integ/pagination.integ.spec.ts +++ b/packages/core/integ/pagination.integ.spec.ts @@ -1,6 +1,6 @@ import { requireRequestsFrom } from "@aws-sdk/aws-util-test/src"; import { DynamoDB, paginateScan, ScanCommandInput } from "@aws-sdk/client-dynamodb"; -import { BatchGetCommandInput, paginateBatchGet } from "@aws-sdk/lib-dynamodb"; +import { BatchGetCommandInput, paginatedBatchGet } from "@aws-sdk/lib-dynamodb"; import { HttpResponse } from "@smithy/protocol-http"; import { describe, expect, test as it } from "vitest"; @@ -200,7 +200,7 @@ describe("pagination", () => { }; let pages = 0; - for await (const page of paginateBatchGet({ client: ddb }, requestParams)) { + for await (const page of paginatedBatchGet({ client: ddb }, requestParams)) { pages += 1; if (pages === 1) { expect(page.Responses?.test).toEqual([ @@ -220,7 +220,7 @@ describe("pagination", () => { let thrownError; try { - for await (const page of paginateBatchGet( + for await (const page of paginatedBatchGet( { client: ddb }, { RequestItems: { diff --git a/packages/util-dynamodb/src/index.ts b/packages/util-dynamodb/src/index.ts index 4ac3c4445f036..6e3fd8beb4f53 100644 --- a/packages/util-dynamodb/src/index.ts +++ b/packages/util-dynamodb/src/index.ts @@ -3,5 +3,5 @@ export * from "./convertToAttr"; export * from "./convertToNative"; export * from "./marshall"; export * from "./models"; -export * from "./paginateBatchGetItem"; +export * from "./paginatedBatchGetItem"; export * from "./unmarshall"; diff --git a/packages/util-dynamodb/src/paginateBatchGetItem.ts b/packages/util-dynamodb/src/paginatedBatchGetItem.ts similarity index 78% rename from packages/util-dynamodb/src/paginateBatchGetItem.ts rename to packages/util-dynamodb/src/paginatedBatchGetItem.ts index bc74859a96612..53d5b8ee602dd 100644 --- a/packages/util-dynamodb/src/paginateBatchGetItem.ts +++ b/packages/util-dynamodb/src/paginatedBatchGetItem.ts @@ -6,7 +6,7 @@ import { type DynamoDBDocumentClient } from "@aws-sdk/lib-dynamodb"; * * @public * - * @see {@link paginateBatchGetItem} for a variant that uses the {@link DynamoDBDocumentClient | DynamoDB document client}. + * @see {@link paginatedBatchGetItem} for a variant that uses the {@link DynamoDBDocumentClient | DynamoDB document client}. * * @example * @@ -27,12 +27,12 @@ import { type DynamoDBDocumentClient } from "@aws-sdk/lib-dynamodb"; * } * ``` */ -export async function* paginateBatchGetItem(client: DynamoDBClient, input: BatchGetItemCommandInput) { +export async function* paginatedBatchGetItem(config: { client: DynamoDBClient }, input: BatchGetItemCommandInput) { let RequestItems = input.RequestItems; while (RequestItems && Object.keys(RequestItems).length > 0) { const cmd = new BatchGetItemCommand({ ...input, RequestItems }); - const response = await client.send(cmd); + const response = await config.client.send(cmd); RequestItems = { ...response.UnprocessedKeys }; yield response;