[wrangler] teach wrangler's D1 commands about config (#3055)

* [wrangler] teach wrangler's D1 commands about config * fix lint issue * write tests
cloudflare · Apr 17, 2023 · 5f48c40 · 5f48c40
1 parent 7930471
commit 5f48c40
Show file tree

Hide file tree

Showing 9 changed files with 177 additions and 67 deletions.
diff --git a/.changeset/orange-plants-flash.md b/.changeset/orange-plants-flash.md
@@ -0,0 +1,9 @@
+---
+"wrangler": patch
+---
+
+fix: Teach D1 commands to read auth configuration from wrangler.toml
+
+This PR fixes a bug in how D1 handles a user's accounts. We've updated the D1 commands to read from config (typically via wrangler.toml) before trying to run commands. This means if an `account_id` is defined in config, we'll use that instead of erroring out when there are multiple accounts to pick from.
+
+Fixes #3046
diff --git a/packages/wrangler/src/__tests__/d1/migrate.test.ts b/packages/wrangler/src/__tests__/d1/migrate.test.ts
@@ -1,15 +1,22 @@
 import { cwd } from "process";
+import { rest } from "msw";
 import { reinitialiseAuthTokens } from "../../user";
 import { mockAccountId, mockApiToken } from "../helpers/mock-account-id";
 import { mockConsoleMethods } from "../helpers/mock-console";
+import { mockConfirm } from "../helpers/mock-dialogs";
 import { useMockIsTTY } from "../helpers/mock-istty";
+import { mockGetMemberships, mockOAuthFlow } from "../helpers/mock-oauth-flow";
+import { mockSetTimeout } from "../helpers/mock-set-timeout";
+import { msw } from "../helpers/msw";
 import { runInTempDir } from "../helpers/run-in-tmp";
 import { runWrangler } from "../helpers/run-wrangler";
 import writeWranglerToml from "../helpers/write-wrangler-toml";
 
 describe("migrate", () => {
 	runInTempDir();
 	mockConsoleMethods();
+	mockSetTimeout();
+
 	const { setIsTTY } = useMockIsTTY();
 
 	describe("create", () => {
@@ -28,6 +35,9 @@ describe("migrate", () => {
 	});
 
 	describe("apply", () => {
+		mockAccountId({ accountId: null });
+		mockApiToken();
+		const { mockOAuthServerCallback } = mockOAuthFlow();
 		it("should not attempt to login in local mode", async () => {
 			setIsTTY(false);
 			writeWranglerToml({
@@ -77,6 +87,96 @@ describe("migrate", () => {
 				runWrangler("d1 migrations apply --local db --preview")
 			).rejects.toThrowError(`Error: can't use --preview with --local`);
 		});
+
+		it("multiple accounts: should throw when trying to apply migrations without an account_id in config", async () => {
+			setIsTTY(false);
+
+			writeWranglerToml({
+				d1_databases: [
+					{
+						binding: "DATABASE",
+						database_name: "db",
+						database_id: "xxxx",
+						migrations_dir: "/tmp/my-migrations-go-here",
+					},
+				],
+			});
+			mockOAuthServerCallback();
+			mockGetMemberships([
+				{ id: "IG-88", account: { id: "1701", name: "enterprise" } },
+				{ id: "R2-D2", account: { id: "nx01", name: "enterprise-nx" } },
+			]);
+			mockConfirm({
+				text: `No migrations folder found.
+Ok to create /tmp/my-migrations-go-here?`,
+				result: true,
+			});
+			await runWrangler("d1 migrations create db test");
+			mockConfirm({
+				text: `About to apply 1 migration(s)
+Your database may not be available to serve requests during the migration, continue?`,
+				result: true,
+			});
+			await expect(runWrangler("d1 migrations apply db")).rejects.toThrowError(
+				`More than one account available but unable to select one in non-interactive mode.`
+			);
+		});
+		it("multiple accounts: should let the user apply migrations with an account_id in config", async () => {
+			setIsTTY(false);
+			msw.use(
+				rest.post(
+					"*/accounts/:accountId/d1/database/:databaseId/query",
+					async (req, res, ctx) => {
+						return res(
+							ctx.status(200),
+							ctx.json({
+								result: [
+									{
+										results: [],
+										success: true,
+										meta: {},
+									},
+								],
+								success: true,
+								errors: [],
+								messages: [],
+							})
+						);
+					}
+				)
+			);
+			writeWranglerToml({
+				d1_databases: [
+					{
+						binding: "DATABASE",
+						database_name: "db",
+						database_id: "xxxx",
+						migrations_dir: "/tmp/my-migrations-go-here",
+					},
+				],
+				account_id: "nx01",
+			});
+			mockOAuthServerCallback();
+			mockGetMemberships([
+				{ id: "IG-88", account: { id: "1701", name: "enterprise" } },
+				{ id: "R2-D2", account: { id: "nx01", name: "enterprise-nx" } },
+			]);
+			mockConfirm({
+				text: `No migrations folder found.
+Ok to create /tmp/my-migrations-go-here?`,
+				result: true,
+			});
+			await runWrangler("d1 migrations create db test");
+			mockConfirm({
+				text: `About to apply 1 migration(s)
+Your database may not be available to serve requests during the migration, continue?`,
+				result: true,
+			});
+			//if we get to this point, wrangler knows the account_id
+			await expect(runWrangler("d1 migrations apply db")).rejects.toThrowError(
+				`request to https://api.cloudflare.com/client/v4/accounts/nx01/d1/database/xxxx/backup failed`
+			);
+		});
 	});
 
 	describe("list", () => {

diff --git a/packages/wrangler/src/d1/backups.tsx b/packages/wrangler/src/d1/backups.tsx
@@ -24,7 +24,7 @@ export function ListOptions(yargs: CommonYargsArgv) {
 type ListHandlerOptions = StrictYargsOptionsToInterface<typeof ListOptions>;
 export const ListHandler = withConfig<ListHandlerOptions>(
 	async ({ config, name }): Promise<void> => {
-		const accountId = await requireAuth({});
+		const accountId = await requireAuth(config);
 		logger.log(d1BetaWarning);
 		const db: Database = await getDatabaseByNameOrBinding(
 			config,
@@ -87,7 +87,7 @@ type CreateHandlerOptions = StrictYargsOptionsToInterface<typeof CreateOptions>;
 
 export const CreateHandler = withConfig<CreateHandlerOptions>(
 	async ({ config, name }): Promise<void> => {
-		const accountId = await requireAuth({});
+		const accountId = await requireAuth(config);
 		logger.log(d1BetaWarning);
 		const db: Database = await getDatabaseByNameOrBinding(
 			config,
@@ -135,7 +135,7 @@ type RestoreHandlerOptions = StrictYargsOptionsToInterface<
 >;
 export const RestoreHandler = withConfig<RestoreHandlerOptions>(
 	async ({ config, name, backupId }): Promise<void> => {
-		const accountId = await requireAuth({});
+		const accountId = await requireAuth(config);
 		logger.log(d1BetaWarning);
 		const db: Database = await getDatabaseByNameOrBinding(
 			config,
@@ -183,7 +183,7 @@ type DownloadHandlerOptions = StrictYargsOptionsToInterface<
 >;
 export const DownloadHandler = withConfig<DownloadHandlerOptions>(
 	async ({ name, backupId, output, config }): Promise<void> => {
-		const accountId = await requireAuth({});
+		const accountId = await requireAuth(config);
 		logger.log(d1BetaWarning);
 		const db: Database = await getDatabaseByNameOrBinding(
 			config,

diff --git a/packages/wrangler/src/d1/create.tsx b/packages/wrangler/src/d1/create.tsx
@@ -1,6 +1,7 @@
 import { Text, Box } from "ink";
 import React from "react";
 import { fetchResult } from "../cfetch";
+import { withConfig } from "../config";
 import { logger } from "../logger";
 import { requireAuth } from "../user";
 import { renderToString } from "../utils/render";
@@ -21,48 +22,49 @@ export function Options(yargs: CommonYargsArgv) {
 		.epilogue(d1BetaWarning);
 }
 
-export async function Handler({
-	name,
-}: StrictYargsOptionsToInterface<typeof Options>): Promise<void> {
-	const accountId = await requireAuth({});
+type HandlerOptions = StrictYargsOptionsToInterface<typeof Options>;
+export const Handler = withConfig<HandlerOptions>(
+	async ({ name, config }): Promise<void> => {
+		const accountId = await requireAuth(config);
 
-	logger.log(d1BetaWarning);
+		logger.log(d1BetaWarning);
 
-	let db: Database;
-	try {
-		db = await fetchResult(`/accounts/${accountId}/d1/database`, {
-			method: "POST",
-			headers: {
-				"Content-Type": "application/json",
-			},
-			body: JSON.stringify({
-				name,
-			}),
-		});
-	} catch (e) {
-		if ((e as { code: number }).code === 7502) {
-			throw new Error("A database with that name already exists");
+		let db: Database;
+		try {
+			db = await fetchResult(`/accounts/${accountId}/d1/database`, {
+				method: "POST",
+				headers: {
+					"Content-Type": "application/json",
+				},
+				body: JSON.stringify({
+					name,
+				}),
+			});
+		} catch (e) {
+			if ((e as { code: number }).code === 7502) {
+				throw new Error("A database with that name already exists");
+			}
+			throw e;
 		}
-		throw e;
-	}
 
-	logger.log(
-		renderToString(
-			<Box flexDirection="column">
-				<Text>✅ Successfully created DB &apos;{db.name}&apos;!</Text>
-				<Text>&nbsp;</Text>
-				<Text>
-					Add the following to your wrangler.toml to connect to it from a
-					Worker:
-				</Text>
-				<Text>&nbsp;</Text>
-				<Text>[[ d1_databases ]]</Text>
-				<Text>
-					binding = &quot;DB&quot; # i.e. available in your Worker on env.DB
-				</Text>
-				<Text>database_name = &quot;{db.name}&quot;</Text>
-				<Text>database_id = &quot;{db.uuid}&quot;</Text>
-			</Box>
-		)
-	);
-}
+		logger.log(
+			renderToString(
+				<Box flexDirection="column">
+					<Text>✅ Successfully created DB &apos;{db.name}&apos;!</Text>
+					<Text>&nbsp;</Text>
+					<Text>
+						Add the following to your wrangler.toml to connect to it from a
+						Worker:
+					</Text>
+					<Text>&nbsp;</Text>
+					<Text>[[ d1_databases ]]</Text>
+					<Text>
+						binding = &quot;DB&quot; # i.e. available in your Worker on env.DB
+					</Text>
+					<Text>database_name = &quot;{db.name}&quot;</Text>
+					<Text>database_id = &quot;{db.uuid}&quot;</Text>
+				</Box>
+			)
+		);
+	}
+);
diff --git a/packages/wrangler/src/d1/delete.ts b/packages/wrangler/src/d1/delete.ts
@@ -24,7 +24,7 @@ export function Options(d1ListYargs: CommonYargsArgv) {
 type HandlerOptions = StrictYargsOptionsToInterface<typeof Options>;
 export const Handler = withConfig<HandlerOptions>(
 	async ({ name, skipConfirmation, config }): Promise<void> => {
-		const accountId = await requireAuth({});
+		const accountId = await requireAuth(config);
 		logger.log(d1BetaWarning);
 
 		const db: Database = await getDatabaseByNameOrBinding(

diff --git a/packages/wrangler/src/d1/execute.tsx b/packages/wrangler/src/d1/execute.tsx
@@ -281,7 +281,7 @@ async function executeRemotely({
 		}
 	}
 
-	const accountId = await requireAuth({});
+	const accountId = await requireAuth(config);
 	const db: Database = await getDatabaseByNameOrBinding(
 		config,
 		accountId,

diff --git a/packages/wrangler/src/d1/list.tsx b/packages/wrangler/src/d1/list.tsx
@@ -1,6 +1,7 @@
 import Table from "ink-table";
 import React from "react";
 import { fetchResult } from "../cfetch";
+import { withConfig } from "../config";
 import { logger } from "../logger";
 import { requireAuth } from "../user";
 import { renderToString } from "../utils/render";
@@ -21,19 +22,20 @@ export function Options(d1ListYargs: CommonYargsArgv) {
 		.epilogue(d1BetaWarning);
 }
 
-export async function Handler({
-	json,
-}: StrictYargsOptionsToInterface<typeof Options>): Promise<void> {
-	const accountId = await requireAuth({});
-	const dbs: Array<Database> = await listDatabases(accountId);
+type HandlerOptions = StrictYargsOptionsToInterface<typeof Options>;
+export const Handler = withConfig<HandlerOptions>(
+	async ({ json, config }): Promise<void> => {
+		const accountId = await requireAuth(config);
+		const dbs: Array<Database> = await listDatabases(accountId);
 
-	if (json) {
-		logger.log(JSON.stringify(dbs, null, 2));
-	} else {
-		logger.log(d1BetaWarning);
-		logger.log(renderToString(<Table data={dbs}></Table>));
+		if (json) {
+			logger.log(JSON.stringify(dbs, null, 2));
+		} else {
+			logger.log(d1BetaWarning);
+			logger.log(renderToString(<Table data={dbs}></Table>));
+		}
 	}
-}
+);
 
 export const listDatabases = async (
 	accountId: string

diff --git a/packages/wrangler/src/d1/migrations/apply.tsx b/packages/wrangler/src/d1/migrations/apply.tsx
@@ -122,7 +122,7 @@ Your database may not be available to serve requests during the migration, conti
 				"In non-local mode `databaseInfo` should be defined."
 			);
 			logger.log(renderToString(<Text>🕒 Creating backup...</Text>));
-			const accountId = await requireAuth({});
+			const accountId = await requireAuth(config);
 			await createBackup(accountId, databaseInfo.uuid);
 		}
 

diff --git a/packages/wrangler/src/d1/migrations/helpers.ts b/packages/wrangler/src/d1/migrations/helpers.ts
@@ -157,14 +157,11 @@ export const initMigrationsTable = async ({
 		name,
 		shouldPrompt: isInteractive() && !CI.isCI(),
 		persistTo,
-		command: `
-						CREATE TABLE IF NOT EXISTS ${migrationsTableName}
-						(
-								id         INTEGER PRIMARY KEY AUTOINCREMENT,
-								name       TEXT UNIQUE,
-								applied_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP NOT NULL
-						);
-				`,
+		command: `CREATE TABLE IF NOT EXISTS ${migrationsTableName}(
+		id         INTEGER PRIMARY KEY AUTOINCREMENT,
+		name       TEXT UNIQUE,
+		applied_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP NOT NULL
+);`,
 		file: undefined,
 		json: undefined,
 		preview,