CFSQL-1196: Extend D1 local dev experience to work with Sessions API bookmarks (#8445)

CFSQL-1196: Extend D1 local dev experience to work with Sessions API bookmarks

Author: lambrospetrou

.changeset/fuzzy-regions-post.md

@@ -0,0 +1,5 @@
+---
+"miniflare": minor
+---
+
+D1 local developer experience supports sessions API bookmarks

fixtures/d1-read-replication-app/.gitignore

@@ -0,0 +1,2 @@
+dist
+.wrangler

fixtures/d1-read-replication-app/README.md

@@ -0,0 +1,3 @@
+# Notes about testing with this fixture
+
+- This includes tests related to the D1 read replication feature.

fixtures/d1-read-replication-app/package.json

@@ -0,0 +1,22 @@
+{
+	"name": "d1-read-replication-app",
+	"private": true,
+	"scripts": {
+		"check:type": "tsc",
+		"start": "wrangler dev",
+		"test:ci": "vitest run",
+		"test:watch": "vitest",
+		"type:tests": "tsc -p ./tests/tsconfig.json"
+	},
+	"devDependencies": {
+		"@cloudflare/workers-tsconfig": "workspace:*",
+		"@cloudflare/workers-types": "^4.20250224.0",
+		"typescript": "catalog:default",
+		"undici": "catalog:default",
+		"vitest": "catalog:default",
+		"wrangler": "workspace:*"
+	},
+	"volta": {
+		"extends": "../../package.json"
+	}
+}

fixtures/d1-read-replication-app/src/index.ts

@@ -0,0 +1,41 @@
+// These will become default after read replication is open.
+type D1DatabaseWithSessionsAPI = D1Database & {
+	// constraintOrBookmark: "first-primary" | "first-unconstrained" | string
+	withSession(constraintOrBookmark: string): D1DatabaseSession;
+};
+
+type D1DatabaseSession = Pick<D1Database, "prepare" | "batch"> & {
+	getBookmark(): string;
+};
+
+export interface Env {
+	DB01: D1DatabaseWithSessionsAPI;
+}
+
+export default {
+	async fetch(request: Request, env: Env) {
+		const url = new URL(request.url);
+
+		let bookmark = "first-primary";
+		let q = "select 1;";
+
+		if (url.pathname === "/sql") {
+			bookmark = url.searchParams.get("bookmark");
+			q = url.searchParams.get("q");
+		}
+
+		const session = env.DB01.withSession(bookmark);
+		// Dummy select to get the bookmark before the main query.
+		await session.prepare("select 1").all();
+		const bookmarkBefore = session.getBookmark();
+		// Now do the main query requested.
+		const result = await session.prepare(q).all();
+		const bookmarkAfter = session.getBookmark();
+
+		return Response.json({
+			bookmarkBefore,
+			bookmarkAfter,
+			result,
+		});
+	},
+};

fixtures/d1-read-replication-app/tests/index.test.ts

@@ -0,0 +1,82 @@
+import { resolve } from "node:path";
+import { afterAll, beforeAll, describe, it } from "vitest";
+import { runWranglerDev } from "../../shared/src/run-wrangler-long-lived";
+
+describe("d1-sessions-api - getBookmark", () => {
+	describe("with wrangler dev", () => {
+		let ip: string, port: number, stop: (() => Promise<unknown>) | undefined;
+
+		beforeAll(async () => {
+			({ ip, port, stop } = await runWranglerDev(resolve(__dirname, ".."), [
+				"--port=0",
+				"--inspector-port=0",
+			]));
+		});
+
+		afterAll(async () => {
+			await stop?.();
+		});
+
+		it("should respond with bookmarks before and after a session query", async ({
+			expect,
+		}) => {
+			let response = await fetch(`http://${ip}:${port}`);
+			let parsed = await response.json();
+			expect(response.status).toBe(200);
+			expect(parsed).toMatchObject({
+				bookmarkBefore: expect.stringMatching(/\w{8}-\w{8}-\w{8}-\w{32}/),
+				bookmarkAfter: expect.stringMatching(/\w{8}-\w{8}-\w{8}-\w{32}/),
+			});
+		});
+
+		it("should progress the bookmark after a write", async ({ expect }) => {
+			let response = await fetch(
+				`http://${ip}:${port}?q=${encodeURIComponent("create table if not exists users1(id text);")}`
+			);
+			let parsed = (await response.json()) as {
+				bookmarkAfter: string;
+				bookmarkBefore: string;
+			};
+			expect(response.status).toBe(200);
+			expect(parsed).toMatchObject({
+				bookmarkBefore: expect.stringMatching(/\w{8}-\w{8}-\w{8}-\w{32}/),
+				bookmarkAfter: expect.stringMatching(/\w{8}-\w{8}-\w{8}-\w{32}/),
+			});
+			expect(
+				parsed.bookmarkAfter > parsed.bookmarkBefore,
+				`before[${parsed.bookmarkBefore}] !== after[${parsed.bookmarkAfter}]`
+			).toEqual(true);
+		});
+
+		it("should maintain the latest bookmark after many queries", async ({
+			expect,
+		}) => {
+			let responses = [];
+
+			for (let i = 0; i < 10; i++) {
+				const resp = await fetch(
+					`http://${ip}:${port}?q=${encodeURIComponent(`create table if not exists users${i}(id text);`)}`
+				);
+				let parsed = (await resp.json()) as {
+					bookmarkAfter: string;
+					bookmarkBefore: string;
+				};
+				expect(resp.status).toBe(200);
+				responses.push(parsed);
+
+				expect(
+					parsed.bookmarkAfter > parsed.bookmarkBefore,
+					`before[${parsed.bookmarkBefore}] !== after[${parsed.bookmarkAfter}]`
+				).toEqual(true);
+			}
+
+			const lastBookmark = responses.at(-1)?.bookmarkAfter;
+			responses.slice(0, -1).forEach((parsed) => {
+				expect(
+					parsed.bookmarkAfter < lastBookmark!,
+					`previous after[${parsed.bookmarkAfter}] !< lastBookmark[${lastBookmark}]`
+				).toEqual(true);
+			});
+		});
+	});
+});

fixtures/d1-read-replication-app/tests/tsconfig.json

@@ -0,0 +1,7 @@
+{
+	"extends": "@cloudflare/workers-tsconfig/tsconfig.json",
+	"compilerOptions": {
+		"types": ["node"]
+	},
+	"include": ["**/*.ts", "../../../node-types.d.ts"]
+}

fixtures/d1-read-replication-app/tsconfig.json

@@ -0,0 +1,13 @@
+{
+	"compilerOptions": {
+		"target": "ES2020",
+		"module": "CommonJS",
+		"lib": ["ES2020"],
+		"types": ["@cloudflare/workers-types"],
+		"moduleResolution": "node",
+		"noEmit": true,
+		"skipLibCheck": true
+	},
+	"include": ["**/*.ts"],
+	"exclude": ["tests"]
+}

fixtures/d1-read-replication-app/wrangler.toml

@@ -0,0 +1,10 @@
+name = "d1-read-replication-app"
+main = "src/index.ts"
+compatibility_date = "2025-03-11"
+compatibility_flags = ["nodejs_compat", "experimental", "enable_d1_with_sessions_api"]
+
+[[d1_databases]]
+binding = "DB01" # i.e. available in your Worker on env.DB
+database_name = "UPDATE_THIS_FOR_REMOTE_USE"
+preview_database_id = "UPDATE_THIS_FOR_REMOTE_USE"
+database_id = "UPDATE_THIS_FOR_REMOTE_USE"

packages/miniflare/src/workers/d1/database.worker.ts

@@ -66,6 +66,9 @@ interface D1FailureResponse {
 	error: string;
 }
 
+// See https://github.com/cloudflare/workerd/blob/bcaf44290296c71f396175c30b43cff6b570231f/src/cloudflare/internal/d1-api.ts#L68-L72
+const D1_SESSION_COMMIT_TOKEN_HTTP_HEADER = "x-cf-d1-session-commit-token";
+
 export class D1Error extends HttpError {
 	constructor(readonly cause: unknown) {
 		super(500);
@@ -217,7 +220,12 @@ export class D1DatabaseObject extends MiniflareDurableObject {
 			searchParams.get("resultsFormat")
 		);
 
-		return Response.json(this.#txn(queries, resultsFormat));
+		return Response.json(this.#txn(queries, resultsFormat), {
+			headers: {
+				[D1_SESSION_COMMIT_TOKEN_HTTP_HEADER]:
+					await this.state.storage.getCurrentBookmark(),
+			},
+		});
 	};
 
 	#isExportPragma(queries: D1Query[]): queries is D1ExportPragma {