feat(dynamodb): add batch get item paginators for low level and document client

Author: cipak

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 =

lib/lib-dynamodb/src/pagination/BatchGetPaginator.ts

@@ -0,0 +1,47 @@
+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<BatchGetCommandOutput> {
+  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;
+  }
+}

lib/lib-dynamodb/src/pagination/index.ts

@@ -1,3 +1,4 @@
+export * from "./BatchGetPaginator";
 // smithy-typescript generated code
 export * from "./Interfaces";
 export * from "./QueryPaginator";

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";

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;
+  }
+}