D1: using new polling endpoint for larger exports

packages/wrangler/src/__tests__/d1/export.test.ts

@@ -45,23 +45,63 @@ describe("execute", () => {
 		msw.use(
 			rest.post(
 				"*/accounts/:accountId/d1/database/:databaseId/export",
-				async (_req, res, ctx) => {
-					return res(
-						ctx.status(200),
-						ctx.json({
-							result: { signedUrl: "https://example.com" },
-							success: true,
-							errors: [],
-							messages: [],
-						})
-					);
+				async (req, res, ctx) => {
+					const body = await req.json();
+					console.log({ body });
+					if (!body.currentBookmark) {
+						return res(
+							ctx.status(202),
+							ctx.json({
+								success: true,
+								result: {
+									type: "export",
+									at_bookmark:
+										"00000018-00000000-00004d71-88bdea5a30e91621c2885390778efd4a",
+									status: "active",
+									messages: [
+										"Generating 80b3c788-ab15-46d9-9cbf-da1dd332bddf-00000018-00000000-00004d71-88bdea5a30e91621c2885390778efd4a.sql",
+										// Out of order uploads ok
+										"Uploaded part 2",
+										"Uploaded part 1",
+									],
+								},
+							})
+						);
+					} else {
+						return res(
+							ctx.status(200),
+							ctx.json({
+								success: true,
+								result: {
+									type: "export",
+									at_bookmark:
+										"00000019-00000000-00004d71-7fe714301c40a6b23b56d03a6d039c7c",
+									status: "complete",
+									result: {
+										filename:
+											"80b3c788-ab15-46d9-9cbf-da1dd332bddf-00000019-00000000-00004d71-7fe714301c40a6b23b56d03a6d039c7c.sql",
+										signedUrl:
+											"https://example.com/80b3c788-ab15-46d9-9cbf-da1dd332bddf-00000019-00000000-00004d71-7fe714301c40a6b23b56d03a6d039c7c.sql",
+									},
+									messages: [
+										"Uploaded part 3",
+										"Uploaded part 4",
+										"Finished uploading 80b3c788-ab15-46d9-9cbf-da1dd332bddf-00000019-00000000-00004d71-7fe714301c40a6b23b56d03a6d039c7c.sql in 4 parts.",
+									],
+								},
+							})
+						);
+					}
 				}
 			)
 		);
 		msw.use(
-			rest.get("https://example.com", async (req, res, ctx) => {
-				return res(ctx.status(200), ctx.text(mockSqlContent));
-			})
+			rest.get(
+				"https://example.com/80b3c788-ab15-46d9-9cbf-da1dd332bddf-00000019-00000000-00004d71-7fe714301c40a6b23b56d03a6d039c7c.sql",
+				async (req, res, ctx) => {
+					return res(ctx.status(200), ctx.text(mockSqlContent));
+				}
+			)
 		);
 
 		await runWrangler("d1 export db --remote --output /tmp/test.sql");

packages/wrangler/src/d1/export.ts

@@ -1,4 +1,5 @@
 import fs from "node:fs/promises";
+import chalk from "chalk";
 import { fetch } from "undici";
 import { printWranglerBanner } from "..";
 import { fetchResult } from "../cfetch";
@@ -80,8 +81,22 @@ export const Handler = async (args: HandlerOptions): Promise<void> => {
 	return result;
 };
 
-type ExportMetadata = {
-	signedUrl: string;
+type PollingResponse = {
+	success: true;
+	result: {
+		type: "export";
+		at_bookmark: string;
+		messages: string[];
+		errors: string[];
+	} & (
+		| {
+				status: "active" | "error";
+		  }
+		| {
+				status: "complete";
+				result: { filename: string; signedUrl: string };
+		  }
+	);
 };
 
 async function exportRemotely(
@@ -101,26 +116,74 @@ async function exportRemotely(
 
 	logger.log(`🌀 Executing on remote database ${name} (${db.uuid}):`);
 	logger.log(`🌀 Creating export...`);
-	const metadata = await fetchResult<ExportMetadata>(
-		`/accounts/${accountId}/d1/database/${db.uuid}/export`,
-		{
-			method: "POST",
-			headers: {
-				"Content-Type": "application/json",
-			},
-			body: JSON.stringify({
-				outputFormat: "file",
-				dumpOptions: {
-					noSchema,
-					noData,
-					tables,
-				},
-			}),
-		}
-	);
+	const dumpOptions = {
+		noSchema,
+		noData,
+		tables,
+	};
+
+	const { result } = await pollExport(accountId, db, dumpOptions, undefined);
 
-	logger.log(`🌀 Downloading SQL to ${output}`);
-	const contents = await fetch(metadata.signedUrl);
+	if (result.status !== "complete")
+		throw new Error(`Error: D1 reset before export completed!`);
+
+	logger.log(`🌀 Downloading SQL to ${output}...`);
+	logger.log(
+		chalk.gray(
+			`If this download fails, you can retry downloading the following URL manually. It is valid for 1 hour: ${result.result.signedUrl}`
+		)
+	);
+	const contents = await fetch(result.result.signedUrl);
 	await fs.writeFile(output, contents.body || "");
 	logger.log(`Done!`);
 }
+
+async function pollExport(
+	accountId: string,
+	db: Database,
+	dumpOptions: {
+		tables: string[];
+		noSchema?: boolean;
+		noData?: boolean;
+	},
+	currentBookmark: string | undefined,
+	num_parts_uploaded = 0
+): Promise<PollingResponse> {
+	const response = await fetchResult<
+		PollingResponse | { success: false; error: string }
+	>(`/accounts/${accountId}/d1/database/${db.uuid}/export`, {
+		method: "POST",
+		body: JSON.stringify({
+			outputFormat: "polling",
+			dumpOptions,
+			currentBookmark,
+		}),
+	});
+
+	console.log({ response });
+	if (!response.success) throw new Error(response.error);
+
+	response.result.messages.forEach((line) => {
+		if (line.startsWith(`Uploaded part`)) {
+			// Part numbers can be reported as complete out-of-order which looks confusing to a user. But their ID has no
+			// special meaning, so just make them sequential.
+			console.log(`🌀 Uploaded part ${++num_parts_uploaded}`);
+		} else {
+			console.log(`🌀 ${line}`);
+		}
+	});
+
+	if (response.result.status === "complete") {
+		return response;
+	} else if (response.result.status === "error") {
+		throw new Error(response.result.errors.join("\n"));
+	} else {
+		return await pollExport(
+			accountId,
+			db,
+			dumpOptions,
+			response.result.at_bookmark,
+			num_parts_uploaded
+		);
+	}
+}