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

.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.

.changeset/smooth-tables-admire.md

@@ -0,0 +1,8 @@
+---
+"wrangler": patch
+---
+
+fix: limit bulk delete API requests to batches of 5,000
+
+The `kv:bulk delete` command now batches up delete 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(`""`);
       });
@@ -1222,7 +1249,12 @@ describe("wrangler", () => {
             expect(new Headers(headers ?? []).get("Content-Type")).toEqual(
               "application/json"
             );
-            expect(JSON.parse(body as string)).toEqual(expectedKeys);
+            expect(JSON.parse(body as string)).toEqual(
+              expectedKeys.slice(
+                (requests.count - 1) * 5000,
+                requests.count * 5000
+              )
+            );
             return null;
           }
         );
@@ -1246,6 +1278,29 @@ describe("wrangler", () => {
         expect(std.err).toMatchInlineSnapshot(`""`);
       });
 
+      it("should delete the keys in batches of 5000 parsed from a file", async () => {
+        const keys = new Array(12000).fill("some-key");
+        writeFileSync("./keys.json", JSON.stringify(keys));
+        mockConfirm({
+          text: `Are you sure you want to delete all the keys read from "keys.json" from kv-namespace with id "some-namespace-id"?`,
+          result: true,
+        });
+        const requests = mockDeleteRequest("some-namespace-id", keys);
+        await runWrangler(
+          `kv:bulk delete --namespace-id some-namespace-id keys.json`
+        );
+        expect(requests.count).toEqual(3);
+        expect(std.out).toMatchInlineSnapshot(`
+          "Deleted 0 of 12000.
+          Deleted 5000 of 12000.
+          Deleted 10000 of 12000.
+          Deleted 12000 of 12000.
+          Success!"
+        `);
+        expect(std.warn).toMatchInlineSnapshot(`""`);
+        expect(std.err).toMatchInlineSnapshot(`""`);
+      });
+
       it("should not delete the keys if the user confirms no", async () => {
         const keys = ["someKey1", "ns:someKey2"];
         writeFileSync("./keys.json", JSON.stringify(keys));
@@ -1323,9 +1378,9 @@ describe("wrangler", () => {
         ).rejects.toThrowErrorMatchingInlineSnapshot(`
                 "Unexpected JSON input from \\"keys.json\\".
                 Expected an array of strings.
-                The item at index 1 is a number: 12354
-                The item at index 2 is a object: [object Object]
-                The item at index 3 is a object: null"
+                The item at index 1 is type: \\"number\\" - 12354
+                The item at index 2 is type: \\"object\\" - {\\"key\\":\\"someKey\\"}
+                The item at index 3 is type: \\"object\\" - null"
               `);
         expect(std.out).toMatchInlineSnapshot(`""`);
         expect(std.warn).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(
@@ -2098,7 +2107,9 @@ export async function main(argv: string[]): Promise<void> {
               const key = content[i];
               if (typeof key !== "string") {
                 errors.push(
-                  `The item at index ${i} is a ${typeof key}: ${key}`
+                  `The item at index ${i} is type: "${typeof key}" - ${JSON.stringify(
+                    key
+                  )}`
                 );
               }
             }
@@ -2112,7 +2123,14 @@ export async function main(argv: string[]): Promise<void> {
 
             const accountId = await requireAuth(config);
 
-            await deleteBulkKeyValue(accountId, namespaceId, content);
+            await deleteBulkKeyValue(
+              accountId,
+              namespaceId,
+              content,
+              (index, total) => {
+                console.log(`Deleted ${index} of ${total}.`);
+              }
+            );
 
             console.log("Success!");
           }

packages/wrangler/src/kv.tsx

@@ -2,6 +2,12 @@ import { URLSearchParams } from "node:url";
 import { fetchListResult, fetchResult, fetchKVGetValue } from "./cfetch";
 import type { Config, Environment } from "./config";
 
+/** The largest number of kv items we can pass to the API in a single request. */
+const API_MAX = 10000;
+// The const below are halved from the API's true capacity to help avoid
+// hammering it with large requests.
+const BATCH_KEY_MAX = API_MAX / 2;
+
 type KvArgs = {
   binding?: string;
   "namespace-id"?: string;
@@ -163,31 +169,51 @@ 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(
   accountId: string,
   namespaceId: string,
-  keys: string[]
+  keys: string[],
+  progressCallback: (index: number, total: number) => void
 ) {
-  return await fetchResult(
-    `/accounts/${accountId}/storage/kv/namespaces/${namespaceId}/bulk`,
-    {
-      method: "DELETE",
-      body: JSON.stringify(keys),
-      headers: { "Content-Type": "application/json" },
+  for (let index = 0; index < keys.length; index += BATCH_KEY_MAX) {
+    if (progressCallback && keys.length > BATCH_KEY_MAX) {
+      progressCallback(index, keys.length);
     }
-  );
+
+    await fetchResult(
+      `/accounts/${accountId}/storage/kv/namespaces/${namespaceId}/bulk`,
+      {
+        method: "DELETE",
+        body: JSON.stringify(keys.slice(index, index + BATCH_KEY_MAX)),
+        headers: { "Content-Type": "application/json" },
+      }
+    );
+  }
+  if (progressCallback && keys.length > BATCH_KEY_MAX) {
+    progressCallback(keys.length, keys.length);
+  }
 }
 
 export function getNamespaceId(

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