From 7ed675e3a43cfd996496bf1be2b31d34bde36664 Mon Sep 17 00:00:00 2001 From: Glen Maddern Date: Fri, 28 Jun 2024 10:01:00 +1000 Subject: [PATCH] Added `d1 export --local` support (#6073) This uses a new miniflare-only debug pragma `PRAGMA miniflare_d1_export` to trigger the dump as D1 provides no in-Worker API (which Miniflare uses) to generate exports. --- .changeset/heavy-tomatoes-sing.md | 6 + .../src/workers/d1/database.worker.ts | 37 ++++ packages/miniflare/src/workers/d1/dumpSql.ts | 160 ++++++++++++++++++ .../wrangler/src/__tests__/d1/export.test.ts | 90 ++++++++-- packages/wrangler/src/d1/export.ts | 83 +++++++-- 5 files changed, 350 insertions(+), 26 deletions(-) create mode 100644 .changeset/heavy-tomatoes-sing.md create mode 100644 packages/miniflare/src/workers/d1/dumpSql.ts diff --git a/.changeset/heavy-tomatoes-sing.md b/.changeset/heavy-tomatoes-sing.md new file mode 100644 index 000000000000..26fef9affe6b --- /dev/null +++ b/.changeset/heavy-tomatoes-sing.md @@ -0,0 +1,6 @@ +--- +"miniflare": minor +"wrangler": minor +--- + +Added D1 export support for local databases diff --git a/packages/miniflare/src/workers/d1/database.worker.ts b/packages/miniflare/src/workers/d1/database.worker.ts index aedd061806ba..d0880bed911f 100644 --- a/packages/miniflare/src/workers/d1/database.worker.ts +++ b/packages/miniflare/src/workers/d1/database.worker.ts @@ -11,6 +11,7 @@ import { viewToBuffer, } from "miniflare:shared"; import { z } from "miniflare:zod"; +import { dumpSql } from "./dumpSql"; const D1ValueSchema = z.union([ z.number(), @@ -27,6 +28,11 @@ const D1QuerySchema = z.object({ type D1Query = z.infer; const D1QueriesSchema = z.union([D1QuerySchema, z.array(D1QuerySchema)]); +const D1_EXPORT_PRAGMA = `PRAGMA miniflare_d1_export(?,?,?);`; +type D1ExportPragma = [ + { sql: typeof D1_EXPORT_PRAGMA; params: [number, number, ...string[]] }, +]; + const D1ResultsFormatSchema = z .enum(["ARRAY_OF_OBJECTS", "ROWS_AND_COLUMNS", "NONE"]) .catch("ARRAY_OF_OBJECTS"); @@ -201,6 +207,11 @@ export class D1DatabaseObject extends MiniflareDurableObject { let queries = D1QueriesSchema.parse(await req.json()); if (!Array.isArray(queries)) queries = [queries]; + // Special local-mode-only handlers + if (this.#isExportPragma(queries)) { + return this.#doExportData(queries); + } + const { searchParams } = new URL(req.url); const resultsFormat = D1ResultsFormatSchema.parse( searchParams.get("resultsFormat") @@ -208,4 +219,30 @@ export class D1DatabaseObject extends MiniflareDurableObject { return Response.json(this.#txn(queries, resultsFormat)); }; + + #isExportPragma(queries: D1Query[]): queries is D1ExportPragma { + return ( + queries.length === 1 && + queries[0].sql === D1_EXPORT_PRAGMA && + (queries[0].params?.length || 0) >= 2 + ); + } + + #doExportData( + queries: [ + { sql: typeof D1_EXPORT_PRAGMA; params: [number, number, ...string[]] }, + ] + ) { + const [noSchema, noData, ...tables] = queries[0].params; + const options = { + noSchema: Boolean(noSchema), + noData: Boolean(noData), + tables, + }; + return Response.json({ + success: true, + results: [Array.from(dumpSql(this.state.storage.sql, options))], + meta: {}, + }); + } } diff --git a/packages/miniflare/src/workers/d1/dumpSql.ts b/packages/miniflare/src/workers/d1/dumpSql.ts new file mode 100644 index 000000000000..652c755c017b --- /dev/null +++ b/packages/miniflare/src/workers/d1/dumpSql.ts @@ -0,0 +1,160 @@ +// NOTE: this function duplicates the logic inside SQLite's shell.c.in as close +// as possible, with any deviations noted. +import { SqlStorage } from "@cloudflare/workers-types/experimental"; + +export function* dumpSql( + db: SqlStorage, + options?: { + noSchema?: boolean; + noData?: boolean; + tables?: string[]; + } +) { + yield `PRAGMA defer_foreign_keys=TRUE;`; + + // Empty set means include all tables + const filterTables = new Set(options?.tables || []); + const { noData, noSchema } = options || {}; + + // Taken from SQLite shell.c.in https://github.com/sqlite/sqlite/blob/105c20648e1b05839fd0638686b95f2e3998abcb/src/shell.c.in#L8463-L8469 + // @ts-ignore (SqlStorageStatement needs to be callable) + const tables_cursor = db.prepare(` + SELECT name, type, sql + FROM sqlite_schema AS o + WHERE (true) AND type=='table' + AND sql NOT NULL + ORDER BY tbl_name='sqlite_sequence', rowid; + `)(); + const tables: any[] = Array.from(tables_cursor); + + for (const { name: table, sql } of tables) { + if (filterTables.size > 0 && !filterTables.has(table)) continue; + + if (table === "sqlite_sequence") { + if (!noSchema) yield `DELETE FROM sqlite_sequence;`; + } else if (table.match(/^sqlite_stat./)) { + // This feels like it should really appear _after_ the contents of sqlite_stat[1,4] but I'm choosing + // to match SQLite's dump output exactly so writing it immediately like they do. + if (!noSchema) yield `ANALYZE sqlite_schema;`; + } else if (sql.startsWith(`CREATE VIRTUAL TABLE`)) { + throw new Error( + `D1 Export error: cannot export databases with Virtual Tables (fts5)` + ); + } else if (table.startsWith("_cf_") || table.startsWith("sqlite_")) { + continue; + } else { + // SQLite dump has an extremely weird behaviour where, if the table was explicitly + // quoted i.e. "Table", then in the dump it has `IF NOT EXISTS` injected. I don't understand + // why, but on the off chance there's a good reason to I am following suit. + if (sql.match(/CREATE TABLE ['"].*/)) { + if (!noSchema) yield `CREATE TABLE IF NOT EXISTS ${sql.substring(13)};`; + } else { + if (!noSchema) yield `${sql};`; + } + } + + if (noData) continue; + const columns_cursor = db.exec(`PRAGMA table_info="${table}"`); + const columns = Array.from(columns_cursor); + const select = `SELECT ${columns.map((c) => c.name).join(", ")} + FROM "${table}";`; + const rows_cursor = db.exec(select); + for (const dataRow of rows_cursor.raw()) { + const formattedCells = dataRow.map((cell: unknown, i: number) => { + const colType = columns[i].type; + const cellType = typeof cell; + if (cell === null) { + return "NULL"; + } else if (colType === "INTEGER" || cellType === "number") { + return cell; + } else if (colType === "TEXT" || cellType === "string") { + return outputQuotedEscapedString(cell); + } else if (cell instanceof ArrayBuffer) { + return `X'${Array.prototype.map + .call(new Uint8Array(cell), (b) => b.toString(16)) + .join("")}'`; + } else { + console.log({ colType, cellType, cell, column: columns[i] }); + return "ERROR"; + } + }); + + yield `INSERT INTO ${sqliteQuote(table)} VALUES(${formattedCells.join(",")});`; + } + } + + if (!noSchema) { + // Taken from SQLite shell.c.in https://github.com/sqlite/sqlite/blob/105c20648e1b05839fd0638686b95f2e3998abcb/src/shell.c.in#L8473-L8478 + const rest_of_schema = db.exec( + `SELECT name, sql FROM sqlite_schema AS o WHERE (true) AND sql NOT NULL AND type IN ('index', 'trigger', 'view') ORDER BY type COLLATE NOCASE /* DESC */;` + ); + // 'DESC' appears in the code linked above but the observed behaviour of SQLite appears otherwise + for (const { name, sql } of rest_of_schema) { + if (filterTables.size > 0 && !filterTables.has(name as string)) continue; + yield `${sql};`; + } + } +} + +// Ported `output_quoted_escaped_string` from https://github.com/sqlite/sqlite/blob/master/src/shell.c.in#L1799-L1862 +function outputQuotedEscapedString(cell: unknown) { + let lfs = false; + let crs = false; + + const quotesOrNewlinesRegexp = /'|(\n)|(\r)/g; + + // Function to replace ' with '', while also tracking whether the string contains any \r or \n chars + const escapeQuotesDetectingNewlines = (_: string, lf: string, cr: string) => { + if (lf) { + lfs = true; + return `\\n`; + } + if (cr) { + crs = true; + return `\\r`; + } + return `''`; + }; + + const escaped_string = (cell as string).replace( + quotesOrNewlinesRegexp, + escapeQuotesDetectingNewlines + ); + let output_string = `'${escaped_string}'`; + if (crs) output_string = `replace(${output_string},'\\r',char(13))`; + if (lfs) output_string = `replace(${output_string},'\\n',char(10))`; + return output_string; +} + +// Ported from quoteChar: https://github.com/sqlite/sqlite/blob/378bf82e2bc09734b8c5869f9b148efe37d29527/src/shell.c.in#L990 +export function sqliteQuote(token: string) { + // Empty input + return token.length === 0 || + // Doesn't start with alpha or underscore + !token.match(/^[a-zA-Z_]/) || + token.match(/\W/) || + SQLITE_KEYWORDS.has(token.toUpperCase()) + ? `"${token}"` + : token; +} + +// List taken from `aKeywordTable` inhttps://github.com/sqlite/sqlite/blob/378bf82e2bc09734b8c5869f9b148efe37d29527/tool/mkkeywordhash.c#L172 +// prettier-ignore +export const SQLITE_KEYWORDS = new Set([ + "ABORT", "ACTION", "ADD", "AFTER", "ALL", "ALTER", "ALWAYS", "ANALYZE", "AND", "AS", "ASC", + "ATTACH", "AUTOINCREMENT", "BEFORE", "BEGIN", "BETWEEN", "BY", "CASCADE", "CASE", "CAST", + "CHECK", "COLLATE", "COLUMN", "COMMIT", "CONFLICT", "CONSTRAINT", "CREATE", "CROSS", "CURRENT", + "CURRENT_DATE", "CURRENT_TIME", "CURRENT_TIMESTAMP", "DATABASE", "DEFAULT", "DEFERRED", + "DEFERRABLE", "DELETE", "DESC", "DETACH", "DISTINCT", "DO", "DROP", "END", "EACH", "ELSE", + "ESCAPE", "EXCEPT", "EXCLUSIVE", "EXCLUDE", "EXISTS", "EXPLAIN", "FAIL", "FILTER", "FIRST", + "FOLLOWING", "FOR", "FOREIGN", "FROM", "FULL", "GENERATED", "GLOB", "GROUP", "GROUPS", "HAVING", + "IF", "IGNORE", "IMMEDIATE", "IN", "INDEX", "INDEXED", "INITIALLY", "INNER", "INSERT", "INSTEAD", + "INTERSECT", "INTO", "IS", "ISNULL", "JOIN", "KEY", "LAST", "LEFT", "LIKE", "LIMIT", "MATCH", + "MATERIALIZED", "NATURAL", "NO", "NOT", "NOTHING", "NOTNULL", "NULL", "NULLS", "OF", "OFFSET", + "ON", "OR", "ORDER", "OTHERS", "OUTER", "OVER", "PARTITION", "PLAN", "PRAGMA", "PRECEDING", + "PRIMARY", "QUERY", "RAISE", "RANGE", "RECURSIVE", "REFERENCES", "REGEXP", "REINDEX", "RELEASE", + "RENAME", "REPLACE", "RESTRICT", "RETURNING", "RIGHT", "ROLLBACK", "ROW", "ROWS", "SAVEPOINT", + "SELECT", "SET", "TABLE", "TEMP", "TEMPORARY", "THEN", "TIES", "TO", "TRANSACTION", "TRIGGER", + "UNBOUNDED", "UNION", "UNIQUE", "UPDATE", "USING", "VACUUM", "VALUES", "VIEW", "VIRTUAL", "WHEN", + "WHERE", "WINDOW", "WITH", "WITHOUT" +]); diff --git a/packages/wrangler/src/__tests__/d1/export.test.ts b/packages/wrangler/src/__tests__/d1/export.test.ts index 4d7ba237a85a..eb43abf01813 100644 --- a/packages/wrangler/src/__tests__/d1/export.test.ts +++ b/packages/wrangler/src/__tests__/d1/export.test.ts @@ -3,7 +3,7 @@ import { http, HttpResponse } from "msw"; import { mockAccountId, mockApiToken } from "../helpers/mock-account-id"; import { mockConsoleMethods } from "../helpers/mock-console"; import { useMockIsTTY } from "../helpers/mock-istty"; -import { mockGetMemberships, mockOAuthFlow } from "../helpers/mock-oauth-flow"; +import { mockGetMemberships } from "../helpers/mock-oauth-flow"; import { msw } from "../helpers/msw"; import { runInTempDir } from "../helpers/run-in-tmp"; import { runWrangler } from "../helpers/run-wrangler"; @@ -14,19 +14,90 @@ describe("execute", () => { mockApiToken(); mockConsoleMethods(); runInTempDir(); - const { mockOAuthServerCallback } = mockOAuthFlow(); const { setIsTTY } = useMockIsTTY(); + it("should throw if output is missing", async () => { await expect(runWrangler("d1 export db --local")).rejects.toThrowError( `Missing required argument: output` ); }); - it("should reject --local mode (for now)", async () => { - await expect( - runWrangler("d1 export db --local --output /tmp/test.sql") - ).rejects.toThrowError( - `Local imports/exports will be coming in a future version of Wrangler.` + it("should handle local", async () => { + setIsTTY(false); + writeWranglerToml({ + d1_databases: [ + { binding: "DATABASE", database_name: "db", database_id: "xxxx" }, + ], + }); + + // Verify the basic command works with an empty DB + await runWrangler("d1 export db --local --output test-local.sql"); + expect(fs.readFileSync("test-local.sql", "utf8")).toBe( + "PRAGMA defer_foreign_keys=TRUE;" + ); + + // Fill with data + fs.writeFileSync( + "data.sql", + ` + CREATE TABLE foo(id INTEGER PRIMARY KEY, value TEXT); + CREATE TABLE bar(id INTEGER PRIMARY KEY, value TEXT); + INSERT INTO foo (value) VALUES ('xxx'),('yyy'),('zzz'); + INSERT INTO bar (value) VALUES ('aaa'),('bbb'),('ccc'); + ` + ); + await runWrangler("d1 execute db --local --file data.sql"); + + // SQL output expectations + const create_foo = "CREATE TABLE foo(id INTEGER PRIMARY KEY, value TEXT);"; + const create_bar = "CREATE TABLE bar(id INTEGER PRIMARY KEY, value TEXT);"; + const insert_foo = [ + "INSERT INTO foo VALUES(1,'xxx');", + "INSERT INTO foo VALUES(2,'yyy');", + "INSERT INTO foo VALUES(3,'zzz');", + ]; + const insert_bar = [ + "INSERT INTO bar VALUES(1,'aaa');", + "INSERT INTO bar VALUES(2,'bbb');", + "INSERT INTO bar VALUES(3,'ccc');", + ]; + + // Full export + await runWrangler("d1 export db --local --output test-full.sql"); + expect(fs.readFileSync("test-full.sql", "utf8")).toBe( + [ + "PRAGMA defer_foreign_keys=TRUE;", + create_foo, + ...insert_foo, + create_bar, + ...insert_bar, + ].join("\n") + ); + + // Schema only + await runWrangler( + "d1 export db --local --output test-schema.sql --no-data" + ); + expect(fs.readFileSync("test-schema.sql", "utf8")).toBe( + ["PRAGMA defer_foreign_keys=TRUE;", create_foo, create_bar].join("\n") + ); + + // Data only + await runWrangler( + "d1 export db --local --output test-data.sql --no-schema" + ); + expect(fs.readFileSync("test-data.sql", "utf8")).toBe( + ["PRAGMA defer_foreign_keys=TRUE;", ...insert_foo, ...insert_bar].join( + "\n" + ) + ); + + // Foo only + await runWrangler( + "d1 export db --local --output test-data.sql --table foo" + ); + expect(fs.readFileSync("test-data.sql", "utf8")).toBe( + ["PRAGMA defer_foreign_keys=TRUE;", create_foo, ...insert_foo].join("\n") ); }); @@ -37,7 +108,6 @@ describe("execute", () => { { binding: "DATABASE", database_name: "db", database_id: "xxxx" }, ], }); - mockOAuthServerCallback(); mockGetMemberships([ { id: "IG-88", account: { id: "1701", name: "enterprise" } }, ]); @@ -102,7 +172,7 @@ describe("execute", () => { }) ); - await runWrangler("d1 export db --remote --output /tmp/test.sql"); - expect(fs.readFileSync("/tmp/test.sql", "utf8")).toBe(mockSqlContent); + await runWrangler("d1 export db --remote --output test-remote.sql"); + expect(fs.readFileSync("test-remote.sql", "utf8")).toBe(mockSqlContent); }); }); diff --git a/packages/wrangler/src/d1/export.ts b/packages/wrangler/src/d1/export.ts index ac3dc9f4d10d..3eff1c774834 100644 --- a/packages/wrangler/src/d1/export.ts +++ b/packages/wrangler/src/d1/export.ts @@ -1,15 +1,19 @@ import fs from "node:fs/promises"; +import path from "node:path"; import chalk from "chalk"; +import { Miniflare } from "miniflare"; import { fetch } from "undici"; import { printWranglerBanner } from ".."; import { fetchResult } from "../cfetch"; import { readConfig } from "../config"; +import { getLocalPersistencePath } from "../dev/get-local-persistence-path"; import { UserError } from "../errors"; import { logger } from "../logger"; import { APIError } from "../parse"; +import { readableRelative } from "../paths"; import { requireAuth } from "../user"; import { Name } from "./options"; -import { getDatabaseByNameOrBinding } from "./utils"; +import { getDatabaseByNameOrBinding, getDatabaseInfoFromConfig } from "./utils"; import type { Config } from "../config"; import type { CommonYargsArgv, @@ -71,12 +75,7 @@ export const Handler = async (args: HandlerOptions): Promise => { await printWranglerBanner(); const config = readConfig(args.config, args); - if (local) { - throw new UserError( - `Local imports/exports will be coming in a future version of Wrangler.` - ); - } - if (!remote) { + if (!local && !remote) { throw new UserError(`You must specify either --local or --remote`); } @@ -91,17 +90,69 @@ export const Handler = async (args: HandlerOptions): Promise => { : [table] : []; - const result = await exportRemotely( - config, - name, - output, - tables, - !schema, - !data - ); - return result; + if (local) { + return await exportLocal(config, name, output, tables, !schema, !data); + } else { + return await exportRemotely(config, name, output, tables, !schema, !data); + } }; +async function exportLocal( + config: Config, + name: string, + output: string, + tables: string[], + noSchema: boolean, + noData: boolean +) { + const localDB = getDatabaseInfoFromConfig(config, name); + if (!localDB) { + throw new UserError( + `Couldn't find a D1 DB with the name or binding '${name}' in wrangler.toml.` + ); + } + + const id = localDB.previewDatabaseUuid ?? localDB.uuid; + + // TODO: should we allow customising persistence path? + // Should it be --persist-to for consistency (even though this isn't persisting anything)? + const persistencePath = getLocalPersistencePath(undefined, config.configPath); + const d1Persist = path.join(persistencePath, "v3", "d1"); + + logger.log( + `🌀 Exporting local database ${name} (${id}) from ${readableRelative( + d1Persist + )}:` + ); + logger.log( + "🌀 To export your remote database, add a --remote flag to your wrangler command." + ); + + const mf = new Miniflare({ + modules: true, + script: "export default {}", + d1Persist, + d1Databases: { DATABASE: id }, + }); + const db = await mf.getD1Database("DATABASE"); + logger.log(`🌀 Exporting SQL to ${output}...`); + + try { + // Special local-only export pragma. Query must be exactly this string to work. + const dump = await db + .prepare(`PRAGMA miniflare_d1_export(?,?,?);`) + .bind(noSchema, noData, ...tables) + .raw(); + await fs.writeFile(output, dump[0].join("\n")); + } catch (e) { + throw new UserError((e as Error).message); + } finally { + await mf.dispose(); + } + + logger.log(`Done!`); +} + async function exportRemotely( config: Config, name: string,