Allow KV and R2 to be referenced to too

cloudflare · Jun 17, 2023 · 2debdc9 · 2debdc9
1 parent 6c047aa
commit 2debdc9
Show file tree

Hide file tree

Showing 4 changed files with 54 additions and 14 deletions.
diff --git a/fixtures/pages-workerjs-directory/package.json b/fixtures/pages-workerjs-directory/package.json
@@ -5,7 +5,7 @@
 	"sideEffects": false,
 	"scripts": {
 		"check:type": "tsc",
-		"dev": "npx wrangler pages dev public --d1=D1 --d1=PUT=elsewhere --port 8794",
+		"dev": "npx wrangler pages dev public --d1=D1 --d1=PUT=elsewhere --kv KV --kv KV_REF=other_kv --r2 R2 --r2=R2_REF=other_r2 --port 8794",
 		"test": "npx vitest run",
 		"test:ci": "npx vitest run",
 		"type:tests": "tsc -p ./tests/tsconfig.json"

diff --git a/fixtures/pages-workerjs-directory/public/_worker.js/index.js b/fixtures/pages-workerjs-directory/public/_worker.js/index.js
@@ -24,6 +24,22 @@ export default {
 			return new Response(JSON.stringify(values));
 		}
 
+		if (pathname === "/kv") {
+			await env.KV.put("key", "value");
+
+			await env.KV_REF.put("key", "value");
+
+			return new Response("saved");
+		}
+
+		if (pathname === "/r2") {
+			await env.R2.put("key", "value");
+
+			await env.R2_REF.put("key", "value");
+
+			return new Response("saved");
+		}
+
 		if (pathname !== "/") {
 			return new Response((await import(`./${pathname.slice(1)}`)).default);
 		}

diff --git a/fixtures/pages-workerjs-directory/tests/index.test.ts b/fixtures/pages-workerjs-directory/tests/index.test.ts
@@ -11,7 +11,15 @@ describe("Pages _worker.js/ directory", () => {
 		const { ip, port, stop } = await runWranglerPagesDev(
 			resolve(__dirname, ".."),
 			"public",
-			["--port=0", "--d1=D1", "--d1=PUT=elsewhere"]
+			[
+				"--port=0",
+				"--d1=D1",
+				"--d1=PUT=elsewhere",
+				"--kv=KV",
+				"--kv=KV_REF=other_kv",
+				"--r2=R2",
+				"--r2=R2_REF=other_r2",
+			]
 		);
 		await expect(
 			fetch(`http://${ip}:${port}/`).then((resp) => resp.text())
@@ -38,6 +46,18 @@ describe("Pages _worker.js/ directory", () => {
 		expect(
 			existsSync(join(__dirname, "../.wrangler/state/v3/d1/elsewhere"))
 		).toBeTruthy();
+		expect(
+			existsSync(join(__dirname, "../.wrangler/state/v3/kv/KV"))
+		).toBeTruthy();
+		expect(
+			existsSync(join(__dirname, "../.wrangler/state/v3/kv/other_kv"))
+		).toBeTruthy();
+		expect(
+			existsSync(join(__dirname, "../.wrangler/state/v3/r2/R2"))
+		).toBeTruthy();
+		expect(
+			existsSync(join(__dirname, "../.wrangler/state/v3/r2/other_r2"))
+		).toBeTruthy();
 	});
 
 	it("should bundle", async ({ expect }) => {

diff --git a/packages/wrangler/src/pages/dev.ts b/packages/wrangler/src/pages/dev.ts
@@ -32,9 +32,7 @@ const DURABLE_OBJECTS_BINDING_REGEXP = new RegExp(
 	/^(?<binding>[^=]+)=(?<className>[^@\s]+)(@(?<scriptName>.*)$)?$/
 );
 
-const D1_BINDING_REGEXP = new RegExp(
-	/^(?<binding>[^=]+)(?:=(?<databaseId>[^\s]+))?$/
-);
+const BINDING_REGEXP = new RegExp(/^(?<binding>[^=]+)(?:=(?<ref>[^\s]+))?$/);
 
 export function Options(yargs: CommonYargsArgv) {
 	return yargs
@@ -524,10 +522,14 @@ export const Handler = async ({
 				.map((binding) => binding.toString().split("="))
 				.map(([key, ...values]) => [key, values.join("=")])
 		),
-		kv: kvs.map((binding) => ({
-			binding: binding.toString(),
-			id: binding.toString(),
-		})),
+		kv: kvs.map((kv) => {
+			const { binding, ref } = BINDING_REGEXP.exec(kv.toString())?.groups || {};
+
+			return {
+				binding,
+				id: ref || kv.toString(),
+			};
+		}),
 		durableObjects: durableObjects
 			.map((durableObject) => {
 				const { binding, className, scriptName } =
@@ -549,8 +551,10 @@ export const Handler = async ({
 				};
 			})
 			.filter(Boolean) as AdditionalDevProps["durableObjects"],
-		r2: r2s.map((binding) => {
-			return { binding: binding.toString(), bucket_name: binding.toString() };
+		r2: r2s.map((r2) => {
+			const { binding, ref } = BINDING_REGEXP.exec(r2.toString())?.groups || {};
+
+			return { binding, bucket_name: ref || binding.toString() };
 		}),
 		rules: usingWorkerDirectory
 			? [
@@ -568,12 +572,12 @@ export const Handler = async ({
 			processEntrypoint: true,
 			additionalModules: modules,
 			d1Databases: d1s.map((d1) => {
-				const { binding, databaseId } =
-					D1_BINDING_REGEXP.exec(d1.toString())?.groups || {};
+				const { binding, ref } =
+					BINDING_REGEXP.exec(d1.toString())?.groups || {};
 
 				return {
 					binding,
-					database_id: databaseId || d1.toString(),
+					database_id: ref || d1.toString(),
 					database_name: `local-${d1}`,
 				};
 			}),