[miniflare] Allow URLs to be passed in hyperdrives option

Importantly, this change makes the Zod output type of `hyperdrives` assignable to the input type. This means we can pass the output of parsing options schemas back to the `new Miniflare()` constructor. Also adds a few more `hyperdrives` tests to verify this doesn't change behaviour.
cloudflare · Feb 7, 2024 · 8a1ea9e · 8a1ea9e
1 parent 97d4461
commit 8a1ea9e
Show file tree

Hide file tree

Showing 4 changed files with 164 additions and 78 deletions.
diff --git a/.changeset/proud-plums-refuse.md b/.changeset/proud-plums-refuse.md
@@ -0,0 +1,7 @@
+---
+"miniflare": minor
+---
+
+feat: allow `URL`s to be passed in `hyperdrives`
+
+Previously, the `hyperdrives` option only accepted `string`s as connection strings. This change allows `URL` objects to be passed too.
diff --git a/packages/miniflare/src/plugins/hyperdrive/index.ts b/packages/miniflare/src/plugins/hyperdrive/index.ts
@@ -1,24 +1,31 @@
+import assert from "node:assert";
 import { z } from "zod";
 import { Service, Worker_Binding } from "../../runtime";
 import { Plugin } from "../shared";
 
 export const HYPERDRIVE_PLUGIN_NAME = "hyperdrive";
 
+function hasPostgresProtocol(url: URL) {
+	return url.protocol === "postgresql:" || url.protocol === "postgres:";
+}
+
+function getPort(url: URL) {
+	if (url.port !== "") return url.port;
+	if (hasPostgresProtocol(url)) return "5432";
+	// Validated in `HyperdriveSchema`
+	assert.fail(`Expected known protocol, got ${url.protocol}`);
+}
+
 export const HyperdriveSchema = z
-	.string()
-	.url()
-	.transform((urlString, ctx) => {
-		const url = new URL(urlString);
+	.union([z.string().url(), z.instanceof(URL)])
+	.transform((url, ctx) => {
+		if (typeof url === "string") url = new URL(url);
 		if (url.protocol === "") {
 			ctx.addIssue({
 				code: z.ZodIssueCode.custom,
 				message: "You must specify the database protocol - e.g. 'postgresql'.",
 			});
-		} else if (
-			url.protocol !== "postgresql:" &&
-			url.protocol !== "postgres:" &&
-			url.protocol !== ""
-		) {
+		} else if (!hasPostgresProtocol(url)) {
 			ctx.addIssue({
 				code: z.ZodIssueCode.custom,
 				message:
@@ -32,21 +39,6 @@ export const HyperdriveSchema = z
 					"You must provide a hostname or IP address in your connection string - e.g. 'user:[email protected]:5432/databasename",
 			});
 		}
-		let port: string | undefined;
-		if (
-			url.port === "" &&
-			(url.protocol === "postgresql:" || url.protocol == "postgres:")
-		) {
-			port = "5432";
-		} else if (url.port !== "") {
-			port = url.port;
-		} else {
-			ctx.addIssue({
-				code: z.ZodIssueCode.custom,
-				message:
-					"You must provide a port number - e.g. 'user:[email protected]:port/databasename",
-			});
-		}
 		if (url.pathname === "") {
 			ctx.addIssue({
 				code: z.ZodIssueCode.custom,
@@ -68,14 +60,8 @@ export const HyperdriveSchema = z
 					"You must provide a password - e.g. 'user:[email protected]:port/databasename' ",
 			});
 		}
-		return {
-			database: url.pathname.replace("/", ""),
-			user: url.username,
-			password: url.password,
-			scheme: url.protocol.replace(":", ""),
-			host: url.hostname,
-			port: port,
-		};
+
+		return url;
 	});
 
 export const HyperdriveInputOptionsSchema = z.object({
@@ -86,29 +72,33 @@ export const HYPERDRIVE_PLUGIN: Plugin<typeof HyperdriveInputOptionsSchema> = {
 	options: HyperdriveInputOptionsSchema,
 	getBindings(options) {
 		return Object.entries(options.hyperdrives ?? {}).map<Worker_Binding>(
-			([name, config]) => ({
-				name,
-				hyperdrive: {
-					designator: {
-						name: `${HYPERDRIVE_PLUGIN_NAME}:${name}`,
+			([name, url]) => {
+				const database = url.pathname.replace("/", "");
+				const scheme = url.protocol.replace(":", "");
+				return {
+					name,
+					hyperdrive: {
+						designator: {
+							name: `${HYPERDRIVE_PLUGIN_NAME}:${name}`,
+						},
+						database,
+						user: url.username,
+						password: url.password,
+						scheme,
 					},
-					database: config.database,
-					user: config.user,
-					password: config.password,
-					scheme: config.scheme,
-				},
-			})
+				};
+			}
 		);
 	},
 	getNodeBindings() {
 		return {};
 	},
 	async getServices({ options }) {
 		return Object.entries(options.hyperdrives ?? {}).map<Service>(
-			([name, config]) => ({
+			([name, url]) => ({
 				name: `${HYPERDRIVE_PLUGIN_NAME}:${name}`,
 				external: {
-					address: `${config.host}:${config.port}`,
+					address: `${url.hostname}:${getPort(url)}`,
 					tcp: {},
 				},
 			})

diff --git a/packages/miniflare/test/plugins/hyperdrive/index.spec.ts b/packages/miniflare/test/plugins/hyperdrive/index.spec.ts
@@ -1,36 +1,126 @@
-import { Hyperdrive } from "@cloudflare/workers-types/experimental";
-import { MiniflareOptions } from "miniflare";
-import { MiniflareTestContext, miniflareTest } from "../../test-shared";
-
-const TEST_CONN_STRING = `postgresql://user:password@localhost:5432/database`;
-
-const opts: Partial<MiniflareOptions> = {
-	hyperdrives: {
-		hyperdrive: TEST_CONN_STRING,
-	},
-};
-
-const test = miniflareTest<{ hyperdrive: Hyperdrive }, MiniflareTestContext>(
-	opts,
-	async (global, _, env) => {
-		return global.Response.json({
-			connectionString: env.hyperdrive.connectionString,
-			user: env.hyperdrive.user,
-			password: env.hyperdrive.password,
-			database: env.hyperdrive.database,
-			host: env.hyperdrive.host,
-		});
-	}
-);
-
-test("configuration: fields match expected", async (t) => {
-	const hyperdriveResp = await t.context.mf.dispatchFetch("http://localhost/");
-	const hyperdrive: any = await hyperdriveResp.json();
+import test from "ava";
+import { Miniflare, MiniflareOptions } from "miniflare";
+
+test("fields match expected", async (t) => {
+	const connectionString = `postgresql://user:password@localhost:5432/database`;
+	const mf = new Miniflare({
+		modules: true,
+		script: `export default {
+			fetch(request, env) {
+				return Response.json({
+					connectionString: env.HYPERDRIVE.connectionString,
+					user: env.HYPERDRIVE.user,
+					password: env.HYPERDRIVE.password,
+					database: env.HYPERDRIVE.database,
+					host: env.HYPERDRIVE.host,
+					port: env.HYPERDRIVE.port,
+				});
+			}
+		}`,
+		hyperdrives: {
+			HYPERDRIVE: connectionString,
+		},
+	});
+	t.teardown(() => mf.dispose());
+	const res = await mf.dispatchFetch("http://localhost/");
+	const hyperdrive = (await res.json()) as Record<string, unknown>;
 	// Since the host is random, this connectionString should be different
-	t.not(hyperdrive.connectionString, TEST_CONN_STRING);
+	t.not(hyperdrive.connectionString, connectionString);
 	t.is(hyperdrive.user, "user");
 	t.is(hyperdrive.password, "password");
 	t.is(hyperdrive.database, "database");
 	// Random host should not be the same as the original
 	t.not(hyperdrive.host, "localhost");
+	t.is(hyperdrive.port, 5432);
+});
+
+test("validates config", async (t) => {
+	const opts: MiniflareOptions = { modules: true, script: "" };
+	const mf = new Miniflare(opts);
+	t.teardown(() => mf.dispose());
+
+	// Check requires Postgres protocol
+	await t.throwsAsync(
+		mf.setOptions({
+			...opts,
+			hyperdrives: {
+				HYPERDRIVE: "mariadb://user:password@localhost:3306/database",
+			},
+		}),
+		{
+			message:
+				/Only PostgreSQL or PostgreSQL compatible databases are currently supported/,
+		}
+	);
+
+	// Check requires host
+	await t.throwsAsync(
+		mf.setOptions({
+			...opts,
+			hyperdrives: { HYPERDRIVE: "postgres:///database" },
+		}),
+		{
+			message:
+				/You must provide a hostname or IP address in your connection string/,
+		}
+	);
+
+	// Check requires database name
+	await t.throwsAsync(
+		mf.setOptions({
+			...opts,
+			hyperdrives: { HYPERDRIVE: "postgres://user:password@localhost:5432" },
+		}),
+		{
+			message: /You must provide a database name as the path component/,
+		}
+	);
+
+	// Check requires username
+	await t.throwsAsync(
+		mf.setOptions({
+			...opts,
+			hyperdrives: { HYPERDRIVE: "postgres://localhost:5432/database" },
+		}),
+		{
+			message: /You must provide a username/,
+		}
+	);
+
+	// Check requires password
+	await t.throwsAsync(
+		mf.setOptions({
+			...opts,
+			hyperdrives: { HYPERDRIVE: "postgres://user@localhost:5432/database" },
+		}),
+		{
+			message: /You must provide a password/,
+		}
+	);
+});
+
+test("sets default port based on protocol", async (t) => {
+	// Check defaults port to 5432 for Postgres
+	const opts = {
+		modules: true,
+		script: `export default {
+			fetch(request, env) {
+				return new Response(env.HYPERDRIVE.port);
+			}
+		}`,
+		hyperdrives: {
+			HYPERDRIVE: "postgresql://user:password@localhost/database" as string | URL,
+		},
+	} satisfies MiniflareOptions;
+	const mf = new Miniflare(opts);
+	t.teardown(() => mf.dispose());
+
+	let res = await mf.dispatchFetch("http://localhost/");
+	t.is(await res.text(), "5432");
+
+	// Check `URL` accepted too
+	opts.hyperdrives.HYPERDRIVE = new URL("postgres://user:password@localhost/database");
+	await mf.setOptions(opts);
+	res = await mf.dispatchFetch("http://localhost/");
+	t.is(await res.text(), "5432");
 });
diff --git a/packages/vitest-pool-workers/test/basic/vitest.config.ts b/packages/vitest-pool-workers/test/basic/vitest.config.ts
@@ -14,10 +14,9 @@ export default defineConfig({
 					bindings: { KEY: "value" },
 					// This doesn't actually do anything in tests
 					upstream: `http://localhost:${inject("port")}`,
-					// TODO(soon): allow object for hyperdrive bindings
-					// hyperdrives: {
-					// 	DATABASE: `postgres://user:[email protected]:${inject("port")}/db`,
-					// },
+					hyperdrives: {
+						DATABASE: `postgres://user:[email protected]:${inject("port")}/db`,
+					},
 				},
 			})),
 		},