D1: using new polling endpoint for larger exports

packages/wrangler/src/d1/constants.ts

@@ -1,5 +1,13 @@
 export const DEFAULT_MIGRATION_PATH = "./migrations";
 export const DEFAULT_MIGRATION_TABLE = "d1_migrations";
-export const LOCATION_CHOICES = ["weur", "eeur", "apac", "wnam", "enam"];
+export const LOCATION_CHOICES = [
+	"weur",
+	"eeur",
+	"apac",
+	"wnam",
+	"enam",
+	"oc",
+	"eagle",
+];
 // Max number of statements to send in a single /execute call
 export const DEFAULT_BATCH_SIZE = 10_000;

packages/wrangler/src/d1/export.ts

@@ -1,7 +1,9 @@
 import fs from "node:fs/promises";
+import chalk from "chalk";
 import { fetch } from "undici";
 import { printWranglerBanner } from "..";
 import { fetchResult } from "../cfetch";
+import { performApiFetch } from "../cfetch/internal";
 import { readConfig } from "../config";
 import { UserError } from "../errors";
 import { logger } from "../logger";
@@ -80,8 +82,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 +117,79 @@ 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: number = 0
+): Promise<PollingResponse> {
+	const response = await fetchResult<
+		PollingResponse | { success: false; error: string }
+	>(`/accounts/${accountId}/d1/database/${db.uuid}/export`, {
+		method: "POST",
+		headers: {
+			"x-d1-internal-env": "staging",
+		},
+		body: JSON.stringify({
+			outputFormat: "polling",
+			dumpOptions: {
+				...dumpOptions,
+				limit: 790000,
+			},
+			currentBookmark,
+		}),
+	});
+
+	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
+		);
+	}
+}