Cap the number of errors and warnings for bulk KV put (#9776)

Co-authored-by: Dario Piotrowicz <[email protected]>

Author: vicb

.changeset/shaggy-frogs-tap.md

@@ -0,0 +1,5 @@
+---
+"wrangler": patch
+---
+
+Cap the number of errors and warnings for bulk KV put to avoid consuming too much memory

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

@@ -1,5 +1,6 @@
 import { writeFileSync } from "node:fs";
 import { http, HttpResponse } from "msw";
+import { BATCH_MAX_ERRORS_WARNINGS } from "../kv/helpers";
 import { endEventLoop } from "./helpers/end-event-loop";
 import { mockAccountId, mockApiToken } from "./helpers/mock-account-id";
 import { mockConsoleMethods } from "./helpers/mock-console";
@@ -1609,6 +1610,79 @@ describe("wrangler", () => {
 			          "
 		        `);
 			});
+
+			it("should cap the number of errors", async () => {
+				const keyValues = [...Array(BATCH_MAX_ERRORS_WARNINGS + 5).keys()];
+
+				writeFileSync("./keys.json", JSON.stringify(keyValues));
+				await expect(
+					runWrangler(`kv bulk put --namespace-id some-namespace-id keys.json`)
+				).rejects.toThrowErrorMatchingInlineSnapshot(`
+					[Error: Unexpected JSON input from "keys.json".
+					Each item in the array should be an object that matches:
+
+					interface KeyValue {
+					  key: string;
+					  value: string;
+					  expiration?: number;
+					  expiration_ttl?: number;
+					  metadata?: object;
+					  base64?: boolean;
+					}
+
+					The item at index 0 is 0
+					The item at index 1 is 1
+					The item at index 2 is 2
+					The item at index 3 is 3
+					The item at index 4 is 4
+					The item at index 5 is 5
+					The item at index 6 is 6
+					The item at index 7 is 7
+					The item at index 8 is 8
+					The item at index 9 is 9
+					The item at index 10 is 10
+					The item at index 11 is 11
+					...]
+				`);
+
+				expect(std.warn).toMatchInlineSnapshot(`""`);
+			});
+
+			it("should cap the number of warnings", async () => {
+				const keyValues: KeyValue[] = new Array(
+					BATCH_MAX_ERRORS_WARNINGS + 5
+				).fill({
+					key: "k",
+					value: "v",
+					invalid: true,
+				});
+				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(1);
+
+				expect(std.warn).toMatchInlineSnapshot(`
+					"▲ [WARNING] Unexpected key-value properties in \\"keys.json\\".
+
+					  The item at index 0 contains unexpected properties: [\\"invalid\\"].
+					  The item at index 1 contains unexpected properties: [\\"invalid\\"].
+					  The item at index 2 contains unexpected properties: [\\"invalid\\"].
+					  The item at index 3 contains unexpected properties: [\\"invalid\\"].
+					  The item at index 4 contains unexpected properties: [\\"invalid\\"].
+					  The item at index 5 contains unexpected properties: [\\"invalid\\"].
+					  The item at index 6 contains unexpected properties: [\\"invalid\\"].
+					  The item at index 7 contains unexpected properties: [\\"invalid\\"].
+					  The item at index 8 contains unexpected properties: [\\"invalid\\"].
+					  The item at index 9 contains unexpected properties: [\\"invalid\\"].
+					  The item at index 10 contains unexpected properties: [\\"invalid\\"].
+					  The item at index 11 contains unexpected properties: [\\"invalid\\"].
+					  ...
+
+					"
+				`);
+			});
 		});
 
 		describe("delete", () => {

packages/wrangler/src/kv/helpers.ts

@@ -16,6 +16,9 @@ const API_MAX = 10000;
 // The const below are lowered from the API's true capacity to help avoid
 // hammering it with large requests.
 export const BATCH_KEY_MAX = API_MAX / 10;
+// Limit the number of errors or warnings to logs during a bulk put.
+// They might end up filling memory for invalid inputs.
+export const BATCH_MAX_ERRORS_WARNINGS = 12;
 
 type KvArgs = {
 	binding?: string;

packages/wrangler/src/kv/index.ts

@@ -16,6 +16,7 @@ import { parseJSON, readFileSync, readFileSyncToBuffer } from "../parse";
 import { requireAuth } from "../user";
 import { getValidBindingName } from "../utils/getValidBindingName";
 import {
+	BATCH_MAX_ERRORS_WARNINGS,
 	createKVNamespace,
 	deleteKVBulkKeyValue,
 	deleteKVKeyValue,
@@ -687,20 +688,32 @@ export const kvBulkPutCommand = createCommand({
 			);
 		}
 
+		let maxNumberOfErrorsReached = false;
 		const errors: string[] = [];
+		let maxNumberOfWarningsReached = false;
 		const warnings: string[] = [];
 		for (let i = 0; i < content.length; i++) {
 			const keyValue = content[i];
-			if (!isKVKeyValue(keyValue)) {
-				errors.push(`The item at index ${i} is ${JSON.stringify(keyValue)}`);
+			if (!isKVKeyValue(keyValue) && !maxNumberOfErrorsReached) {
+				if (errors.length === BATCH_MAX_ERRORS_WARNINGS) {
+					maxNumberOfErrorsReached = true;
+					errors.push("...");
+				} else {
+					errors.push(`The item at index ${i} is ${JSON.stringify(keyValue)}`);
+				}
 			} else {
 				const props = unexpectedKVKeyValueProps(keyValue);
-				if (props.length > 0) {
-					warnings.push(
-						`The item at index ${i} contains unexpected properties: ${JSON.stringify(
-							props
-						)}.`
-					);
+				if (props.length > 0 && !maxNumberOfWarningsReached) {
+					if (warnings.length === BATCH_MAX_ERRORS_WARNINGS) {
+						maxNumberOfWarningsReached = true;
+						warnings.push("...");
+					} else {
+						warnings.push(
+							`The item at index ${i} contains unexpected properties: ${JSON.stringify(
+								props
+							)}.`
+						);
+					}
 				}
 			}
 		}