Allow setting d1 database id in pages dev (#3485)

* Allow setting d1 database id in pages dev

* Allow KV and R2 to be referenced to too

* Fix state directory in fixture test

* Add comments describing the binding regexps

* Remove duplicate import

* More defensively parse bindings

Author: GregBrimble

.changeset/weak-dolphins-carry.md

@@ -0,0 +1,5 @@
+---
+"wrangler": patch
+---
+
+feat: Allow setting a D1 database ID when using `wrangler pages dev` by providing an optional `=<ID>` suffix to the argument like `--d1 BINDING_NAME=database-id`

fixtures/pages-workerjs-directory/package.json

@@ -5,7 +5,7 @@
 	"sideEffects": false,
 	"scripts": {
 		"check:type": "tsc",
-		"dev": "npx wrangler pages dev public --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"

fixtures/pages-workerjs-directory/public/_worker.js/index.js

@@ -15,9 +15,33 @@ export default {
 		}
 
 		if (pathname === "/d1") {
-			const stmt = env.D1.prepare("SELECT 1");
+			const stmt1 = env.D1.prepare("SELECT 1");
+			const values1 = await stmt1.first();
+
+			const stmt = env.PUT.prepare("SELECT 1");
 			const values = await stmt.first();
-			return new Response(JSON.stringify(values));
+
+			if (JSON.stringify(values1) === JSON.stringify(values)) {
+				return new Response(JSON.stringify(values));
+			}
+
+			return new Response("couldn't select 1");
+		}
+
+		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 !== "/") {

fixtures/pages-workerjs-directory/tests/index.test.ts

@@ -1,5 +1,5 @@
 import { execSync } from "node:child_process";
-import { readFileSync } from "node:fs";
+import { existsSync, readFileSync } from "node:fs";
 import { tmpdir } from "node:os";
 import path, { join, resolve } from "node:path";
 import { fetch } from "undici";
@@ -8,10 +8,21 @@ import { runWranglerPagesDev } from "../../shared/src/run-wrangler-long-lived";
 
 describe("Pages _worker.js/ directory", () => {
 	it("should support non-bundling with 'dev'", async ({ expect }) => {
+		const tmpDir = join(tmpdir(), Math.random().toString(36).slice(2));
+
 		const { ip, port, stop } = await runWranglerPagesDev(
 			resolve(__dirname, ".."),
 			"public",
-			["--port=0", "--d1=D1"]
+			[
+				"--port=0",
+				`--persist-to=${tmpDir}`,
+				"--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())
@@ -28,7 +39,20 @@ describe("Pages _worker.js/ directory", () => {
 		await expect(
 			fetch(`http://${ip}:${port}/d1`).then((resp) => resp.text())
 		).resolves.toContain('{"1":1}');
+		await expect(
+			fetch(`http://${ip}:${port}/kv`).then((resp) => resp.text())
+		).resolves.toContain("saved");
+		await expect(
+			fetch(`http://${ip}:${port}/r2`).then((resp) => resp.text())
+		).resolves.toContain("saved");
 		await stop();
+
+		expect(existsSync(join(tmpDir, "./v3/d1/D1"))).toBeTruthy();
+		expect(existsSync(join(tmpDir, "./v3/d1/elsewhere"))).toBeTruthy();
+		expect(existsSync(join(tmpDir, "./v3/kv/KV"))).toBeTruthy();
+		expect(existsSync(join(tmpDir, "./v3/kv/other_kv"))).toBeTruthy();
+		expect(existsSync(join(tmpDir, "./v3/r2/R2"))).toBeTruthy();
+		expect(existsSync(join(tmpDir, "./v3/r2/other_r2"))).toBeTruthy();
 	});
 
 	it("should bundle", async ({ expect }) => {

packages/wrangler/src/pages/dev.ts

@@ -28,10 +28,28 @@ import type {
 } from "../yargs-types";
 import type { RoutesJSONSpec } from "./functions/routes-transformation";
 
+/*
+ * DURABLE_OBJECTS_BINDING_REGEXP matches strings like:
+ * - "binding=className"
+ * - "BINDING=MyClass"
+ * - "BINDING=MyClass@service-name"
+ * Every DO needs a binding (the JS reference) and the exported class name it refers to.
+ * Optionally, users can also provide a service name if they want to reference a DO from another dev session over the dev registry.
+ */
 const DURABLE_OBJECTS_BINDING_REGEXP = new RegExp(
 	/^(?<binding>[^=]+)=(?<className>[^@\s]+)(@(?<scriptName>.*)$)?$/
 );
 
+/* BINDING_REGEXP matches strings like:
+ * - "binding"
+ * - "BINDING"
+ * - "BINDING=ref"
+ * This is used to capture both the binding name (how the binding is used in JS) as well as the reference if provided.
+ * In the case of a D1 database, that's the database ID.
+ * This is useful to people who want to reference the same database in multiple bindings, or a Worker and Pages project dev session want to reference the same database.
+ */
+const BINDING_REGEXP = new RegExp(/^(?<binding>[^=]+)(?:=(?<ref>[^\s]+))?$/);
+
 export function Options(yargs: CommonYargsArgv) {
 	return yargs
 		.positional("directory", {
@@ -520,10 +538,22 @@ 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 || {};
+
+				if (!binding) {
+					logger.warn("Could not parse KV binding:", kv.toString());
+					return;
+				}
+
+				return {
+					binding,
+					id: ref || kv.toString(),
+				};
+			})
+			.filter(Boolean) as AdditionalDevProps["kv"],
 		durableObjects: durableObjects
 			.map((durableObject) => {
 				const { binding, className, scriptName } =
@@ -545,9 +575,19 @@ 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 || {};
+
+				if (!binding) {
+					logger.warn("Could not parse R2 binding:", r2.toString());
+					return;
+				}
+
+				return { binding, bucket_name: ref || binding.toString() };
+			})
+			.filter(Boolean) as AdditionalDevProps["r2"],
 		rules: usingWorkerDirectory
 			? [
 					{
@@ -563,11 +603,23 @@ export const Handler = async ({
 		experimental: {
 			processEntrypoint: true,
 			additionalModules: modules,
-			d1Databases: d1s.map((binding) => ({
-				binding: binding.toString(),
-				database_id: binding.toString(),
-				database_name: `local-${binding}`,
-			})),
+			d1Databases: d1s
+				.map((d1) => {
+					const { binding, ref } =
+						BINDING_REGEXP.exec(d1.toString())?.groups || {};
+
+					if (!binding) {
+						logger.warn("Could not parse D1 binding:", d1.toString());
+						return;
+					}
+
+					return {
+						binding,
+						database_id: ref || d1.toString(),
+						database_name: `local-${d1}`,
+					};
+				})
+				.filter(Boolean) as AdditionalDevProps["d1Databases"],
 			disableExperimentalWarning: true,
 			enablePagesAssetsServiceBinding: {
 				proxyPort,