使用 REST API 批量导入 D1

DROP TABLE IF EXISTS TargetD1Table;
CREATE TABLE IF NOT EXISTS TargetD1Table (id INTEGER PRIMARY KEY, text TEXT, date_added TEXT);

# 创建 D1 数据库
npx wrangler d1 create d1-import-tutorial

# 创建 D1 表
npx wrangler d1 execute d1-import-tutorial --command="DROP TABLE IF EXISTS TargetD1Table; CREATE TABLE IF NOT EXISTS TargetD1Table (id INTEGER PRIMARY KEY, text TEXT, date_added TEXT);" --remote

mkdir d1-import-tutorial
cd d1-import-tutorial
npm init -y

import crypto from "crypto";
import { faker } from "@faker-js/faker";

const TARGET_TABLE = "";
const ACCOUNT_ID = "";
const DATABASE_ID = "";
const D1_API_KEY = "";
const D1_URL = `https://api.cloudflare.com/client/v4/accounts/${ACCOUNT_ID}/d1/database/${DATABASE_ID}/import`;
const uploadSize = 500;
const headers = {
	"Content-Type": "application/json",
	Authorization: `Bearer ${D1_API_KEY}`,
};

// 生成假数据
const data = Array.from({ length: uploadSize }, () => ({
	id: Math.floor(Math.random() * 1000000),
	text: faker.lorem.paragraph(),
	date_added: new Date().toISOString().slice(0, 19).replace("T", " "),
}));

// 创建 SQL insert 语句
function makeSqlInsert(data, tableName, skipCols = []) {
	const columns = Object.keys(data[0]).join(",");
	const values = data
		.map((row) => {
			return (
				"(" +
				Object.values(row)
					.map((val) => {
						if (skipCols.includes(val) || val === null || val === "") {
							return "NULL";
						}
						return `'${String(val).replace(/'/g, "").replace(/"/g, "'")}'`;
					})
					.join(",") +
				")"
			);
		})
		.join(",");

	return `INSERT INTO ${tableName} (${columns}) VALUES ${values};`;
}

const sqlInsert = makeSqlInsert(data, TARGET_TABLE);

async function pollImport(bookmark) {
	const payload = {
		action: "poll",
		current_bookmark: bookmark,
	};

	while (true) {
		const pollResponse = await fetch(D1_URL, {
			method: "POST",
			headers,
			body: JSON.stringify(payload),
		});

		const result = await pollResponse.json();
		console.log("Poll Response:", result.result);

		const { success, error } = result.result;

		if (
			success ||
			(!success && error === "Not currently importing anything.")
		) {
			break;
		}

		await new Promise((resolve) => setTimeout(resolve, 1000));
	}
}

// 上传到 D1
async function uploadToD1() {
	// 1. 初始化上传
	const hashStr = crypto.createHash("md5").update(sqlInsert).digest("hex");

	try {
		const initResponse = await fetch(D1_URL, {
			method: "POST",
			headers,
			body: JSON.stringify({
				action: "init",
				etag: hashStr,
			}),
		});

		const uploadData = await initResponse.json();
		const uploadUrl = uploadData.result.upload_url;
		const filename = uploadData.result.filename;

		// 2. 上传到 R2
		const r2Response = await fetch(uploadUrl, {
			method: "PUT",
			body: sqlInsert,
		});

		const r2Etag = r2Response.headers.get("ETag").replace(/"/g, "");

		// 验证 etag
		if (r2Etag !== hashStr) {
			throw new Error("ETag mismatch");
		}

		// 3. 开始摄取
		const ingestResponse = await fetch(D1_URL, {
			method: "POST",
			headers,
			body: JSON.stringify({
				action: "ingest",
				etag: hashStr,
				filename,
			}),
		});

		const ingestData = await ingestResponse.json();
		console.log("Ingestion Response:", ingestData);

		// 4. 轮询
		await pollImport(ingestData.result.at_bookmark);

		return "Import completed successfully";
	} catch (e) {
		console.error("Error:", e);
		return "Import failed";
	}
}

async function runImport() {
	const result = await uploadToD1();
	console.log(result);
}

runImport();

node index.js