fix: limit bulk put API requests to batches of 5,000

The `kv:bulk put` command now batches up put requests in groups of 5,000,
displaying progress for each request.

Author: petebacondarwin

.changeset/angry-countries-matter.md

@@ -0,0 +1,8 @@
+---
+"wrangler": patch
+---
+
+fix: limit bulk put API requests to batches of 5,000
+
+The `kv:bulk put` command now batches up put requests in groups of 5,000,
+displaying progress for each request.

packages/wrangler/src/__tests__/kv.test.ts

@@ -1130,7 +1130,12 @@ describe("wrangler", () => {
             requests.count++;
             expect(accountId).toEqual("some-account-id");
             expect(namespaceId).toEqual(expectedNamespaceId);
-            expect(JSON.parse(body as string)).toEqual(expectedKeyValues);
+            expect(JSON.parse(body as string)).toEqual(
+              expectedKeyValues.slice(
+                (requests.count - 1) * 5000,
+                requests.count * 5000
+              )
+            );
             return null;
           }
         );
@@ -1150,7 +1155,29 @@ describe("wrangler", () => {
           `kv:bulk put --namespace-id some-namespace-id keys.json`
         );
         expect(requests.count).toEqual(1);
-        expect(std.out).toMatchInlineSnapshot(`""`);
+        expect(std.out).toMatchInlineSnapshot(`"Success!"`);
+        expect(std.warn).toMatchInlineSnapshot(`""`);
+        expect(std.err).toMatchInlineSnapshot(`""`);
+      });
+
+      it("should put the key-values in batches of 5000 parsed from a file", async () => {
+        const keyValues: KeyValue[] = new Array(12000).fill({
+          key: "someKey1",
+          value: "someValue1",
+        });
+        writeFileSync("./keys.json", JSON.stringify(keyValues));
+        const requests = mockPutRequest("some-namespace-id", keyValues);
+        await runWrangler(
+          `kv:bulk put --namespace-id some-namespace-id keys.json`
+        );
+        expect(requests.count).toEqual(3);
+        expect(std.out).toMatchInlineSnapshot(`
+          "Uploaded 0 of 12000.
+          Uploaded 5000 of 12000.
+          Uploaded 10000 of 12000.
+          Uploaded 12000 of 12000.
+          Success!"
+        `);
         expect(std.warn).toMatchInlineSnapshot(`""`);
         expect(std.err).toMatchInlineSnapshot(`""`);
       });

packages/wrangler/src/index.tsx

@@ -2031,7 +2031,16 @@ export async function main(argv: string[]): Promise<void> {
             }
 
             const accountId = await requireAuth(config);
-            await putBulkKeyValue(accountId, namespaceId, content);
+            await putBulkKeyValue(
+              accountId,
+              namespaceId,
+              content,
+              (index, total) => {
+                console.log(`Uploaded ${index} of ${total}.`);
+              }
+            );
+
+            console.log("Success!");
           }
         )
         .command(

packages/wrangler/src/kv.tsx

@@ -169,16 +169,26 @@ export async function getKeyValue(
 export async function putBulkKeyValue(
   accountId: string,
   namespaceId: string,
-  keyValues: KeyValue[]
+  keyValues: KeyValue[],
+  progressCallback: (index: number, total: number) => void
 ) {
-  return await fetchResult(
-    `/accounts/${accountId}/storage/kv/namespaces/${namespaceId}/bulk`,
-    {
-      method: "PUT",
-      body: JSON.stringify(keyValues),
-      headers: { "Content-Type": "application/json" },
+  for (let index = 0; index < keyValues.length; index += BATCH_KEY_MAX) {
+    if (progressCallback && keyValues.length > BATCH_KEY_MAX) {
+      progressCallback(index, keyValues.length);
     }
-  );
+
+    await fetchResult(
+      `/accounts/${accountId}/storage/kv/namespaces/${namespaceId}/bulk`,
+      {
+        method: "PUT",
+        body: JSON.stringify(keyValues.slice(index, index + BATCH_KEY_MAX)),
+        headers: { "Content-Type": "application/json" },
+      }
+    );
+  }
+  if (progressCallback && keyValues.length > BATCH_KEY_MAX) {
+    progressCallback(keyValues.length, keyValues.length);
+  }
 }
 
 export async function deleteBulkKeyValue(

packages/wrangler/src/sites.tsx

@@ -164,7 +164,7 @@ export async function syncAssets(
     }
     manifest[path.relative(siteAssets.baseDirectory, file)] = assetKey;
   }
-  await putBulkKeyValue(accountId, namespace, upload);
+  await putBulkKeyValue(accountId, namespace, upload, () => {});
   return { manifest, namespace };
 }