Partially fix #3160 by properly handling resultsFormat query (#5917)

* Handle resultsFormat in D1 worker

This PR fixes #3160 caused by ignoring the `resultsFormat` query that `workerd` sends.

The commits add proper handling of the query. It makes `raw` preserve columns with the same name by handling `resultsFormat=ROWS_AND_COLUMNS` and removes `results` from `run` as it sends `resultsFormat=NONE` (see https://github.com/cloudflare/workerd/blob/1d89f3b8e9cdcd898ea486656d72d9551e79f4a3/src/cloudflare/internal/d1-api.ts#L295)

* Handle resultsFormat in D1 worker

This PR fixes #3160 caused by ignoring the `resultsFormat` query that `workerd` sends.

The commits add proper handling of the query. It makes `raw` preserve columns with the same name by handling `resultsFormat=ROWS_AND_COLUMNS` and removes `results` from `run` as it sends `resultsFormat=NONE` (see https://github.com/cloudflare/workerd/blob/1d89f3b8e9cdcd898ea486656d72d9551e79f4a3/src/cloudflare/internal/d1-api.ts#L295)

* Update packages/miniflare/test/plugins/d1/test.ts

* Ignore NONE format behaviour in D1 worker

See comment: #5917 (comment)

* Add passing test for new bindings

* Update .changeset/little-jars-train.md

* Update .changeset/little-jars-train.md

* Update .changeset/little-jars-train.md

* Update .changeset/little-jars-train.md

* Update little-jars-train.md

---------

Co-authored-by: Max Rozen <3822106+rozenmd@users.noreply.github.com>

.changeset/little-jars-train.md

@@ -0,0 +1,7 @@
+---
+"miniflare": minor
+---
+
+fix: D1's JOIN behaviour when selecting columns with the same name.
+
+Properly handle the `resultsFormat` query that `workerd` sends. This partially fixes [the JOIN bug](https://github.com/cloudflare/workers-sdk/issues/3160) and makes the behaviour of `raw` consistent with the `workerd` behaviour.

packages/miniflare/src/workers/d1/database.worker.ts

@@ -1,6 +1,5 @@
 import assert from "node:assert";
 import {
-	all,
 	get,
 	HttpError,
 	MiniflareDurableObject,
@@ -28,10 +27,23 @@ const D1QuerySchema = z.object({
 type D1Query = z.infer<typeof D1QuerySchema>;
 const D1QueriesSchema = z.union([D1QuerySchema, z.array(D1QuerySchema)]);
 
+const D1ResultsFormatSchema = z
+	.enum(["ARRAY_OF_OBJECTS", "ROWS_AND_COLUMNS", "NONE"])
+	.catch("ARRAY_OF_OBJECTS");
+
+type D1ResultsFormat = z.infer<typeof D1ResultsFormatSchema>;
+
+interface D1RowsAndColumns {
+	columns: string[];
+	rows: D1Value[][];
+}
+
+type D1Results = D1RowsAndColumns | Record<string, D1Value>[] | undefined;
+
 const served_by = "miniflare.db";
 interface D1SuccessResponse {
 	success: true;
-	results: Record<string, D1Value>[];
+	results: D1Results;
 	meta: {
 		served_by: string;
 		duration: number;
@@ -72,22 +84,28 @@ function convertParams(params: D1Query["params"]): TypedValue[] {
 		Array.isArray(param) ? viewToBuffer(new Uint8Array(param)) : param
 	);
 }
-function convertResults(
-	rows: Record<string, TypedValue>[]
+
+function convertRows(rows: TypedValue[][]): D1Value[][] {
+	return rows.map((row) =>
+		row.map((value) => {
+			let normalised: D1Value;
+			if (value instanceof ArrayBuffer) {
+				// If `value` is an array, convert it to a regular numeric array
+				normalised = Array.from(new Uint8Array(value));
+			} else {
+				normalised = value;
+			}
+			return normalised;
+		})
+	);
+}
+
+function rowsToObjects(
+	columns: string[],
+	rows: D1Value[][]
 ): Record<string, D1Value>[] {
 	return rows.map((row) =>
-		Object.fromEntries(
-			Object.entries(row).map(([key, value]) => {
-				let normalised: D1Value;
-				if (value instanceof ArrayBuffer) {
-					// If `value` is an array, convert it to a regular numeric array
-					normalised = Array.from(new Uint8Array(value));
-				} else {
-					normalised = value;
-				}
-				return [key, normalised];
-			})
-		)
+		Object.fromEntries(columns.map((name, i) => [name, row[i]]))
 	);
 }
 
@@ -113,15 +131,22 @@ export class D1DatabaseObject extends MiniflareDurableObject {
 		return changes;
 	}
 
-	#query = (query: D1Query): D1SuccessResponse => {
+	#query = (format: D1ResultsFormat, query: D1Query): D1SuccessResponse => {
 		const beforeTime = performance.now();
 
 		const beforeSize = this.state.storage.sql.databaseSize;
 		const beforeChanges = this.#changes();
 
 		const params = convertParams(query.params ?? []);
 		const cursor = this.db.prepare(query.sql)(...params);
-		const results = convertResults(all(cursor));
+		const columns = cursor.columnNames;
+		const rows = convertRows(Array.from(cursor.raw()));
+
+		let results = undefined;
+		if (format === "ROWS_AND_COLUMNS") results = { columns, rows };
+		else results = rowsToObjects(columns, rows);
+		// Note that the "NONE" format behaviour here is inconsistent with workerd.
+		// See comment: https://github.com/cloudflare/workers-sdk/pull/5917#issuecomment-2133313156
 
 		const afterTime = performance.now();
 		const afterSize = this.state.storage.sql.databaseSize;
@@ -151,7 +176,7 @@ export class D1DatabaseObject extends MiniflareDurableObject {
 		};
 	};
 
-	#txn(queries: D1Query[]): D1SuccessResponse[] {
+	#txn(queries: D1Query[], format: D1ResultsFormat): D1SuccessResponse[] {
 		// Filter out queries that are just comments
 		queries = queries.filter(
 			(query) => query.sql.replace(/^\s+--.*/gm, "").trim().length > 0
@@ -162,7 +187,9 @@ export class D1DatabaseObject extends MiniflareDurableObject {
 		}
 
 		try {
-			return this.state.storage.transactionSync(() => queries.map(this.#query));
+			return this.state.storage.transactionSync(() =>
+				queries.map(this.#query.bind(this, format))
+			);
 		} catch (e) {
 			throw new D1Error(e);
 		}
@@ -173,6 +200,12 @@ export class D1DatabaseObject extends MiniflareDurableObject {
 	queryExecute: RouteHandler = async (req) => {
 		let queries = D1QueriesSchema.parse(await req.json());
 		if (!Array.isArray(queries)) queries = [queries];
-		return Response.json(this.#txn(queries));
+
+		const { searchParams } = new URL(req.url);
+		const resultsFormat = D1ResultsFormatSchema.parse(
+			searchParams.get("resultsFormat")
+		);
+
+		return Response.json(this.#txn(queries, resultsFormat));
 	};
 }

packages/miniflare/test/plugins/d1/suite.ts

@@ -4,12 +4,18 @@ import { Miniflare, MiniflareOptions } from "miniflare";
 import { useTmp, utf8Encode } from "../../test-shared";
 import { binding, getDatabase, opts, test } from "./test";
 
-export const SCHEMA = (tableColours: string, tableKitchenSink: string) => `
+export const SCHEMA = (
+	tableColours: string,
+	tableKitchenSink: string,
+	tablePalettes: string
+) => `
 CREATE TABLE ${tableColours} (id INTEGER PRIMARY KEY, name TEXT NOT NULL, rgb INTEGER NOT NULL);
 CREATE TABLE ${tableKitchenSink} (id INTEGER PRIMARY KEY, int INTEGER, real REAL, text TEXT, blob BLOB);
+CREATE TABLE ${tablePalettes} (id INTEGER PRIMARY KEY, name TEXT NOT NULL, colour_id INTEGER NOT NULL, FOREIGN KEY (colour_id) REFERENCES ${tableColours}(id));
 INSERT INTO ${tableColours} (id, name, rgb) VALUES (1, 'red', 0xff0000);
 INSERT INTO ${tableColours} (id, name, rgb) VALUES (2, 'green', 0x00ff00);
 INSERT INTO ${tableColours} (id, name, rgb) VALUES (3, 'blue', 0x0000ff);
+INSERT INTO ${tablePalettes} (id, name, colour_id) VALUES (1, 'Night', 3);
 `;
 
 export interface ColourRow {
@@ -32,13 +38,18 @@ test.beforeEach(async (t) => {
 	)}`;
 	const tableColours = `colours_${ns}`;
 	const tableKitchenSink = `kitchen_sink_${ns}`;
+	const tablePalettes = `palettes_${ns}`;
 
 	const db = await getDatabase(t.context.mf);
-	await db.exec(SCHEMA(tableColours, tableKitchenSink));
+	const bindings = await t.context.mf.getBindings();
 
+	await db.exec(SCHEMA(tableColours, tableKitchenSink, tablePalettes));
+
+	t.context.bindings = bindings;
 	t.context.db = db;
 	t.context.tableColours = tableColours;
 	t.context.tableKitchenSink = tableKitchenSink;
+	t.context.tablePalettes = tablePalettes;
 });
 
 function throwCause<T>(promise: Promise<T>): Promise<T> {
@@ -392,10 +403,29 @@ test("D1PreparedStatement: raw", async (t) => {
 		.bind("yellow")
 		.first("id");
 	t.is(id, 4);
+
+	// Check whether workerd raw test case passes here too
+	// Note that this test did not pass with the old binding
+	if (!t.context.bindings["__D1_BETA__DB"]) {
+		await db.prepare(`CREATE TABLE abc (a INT, b INT, c INT);`).run();
+		await db.prepare(`CREATE TABLE cde (c INT, d INT, e INT);`).run();
+		await db.prepare(`INSERT INTO abc VALUES (1,2,3),(4,5,6);`).run();
+		await db.prepare(`INSERT INTO cde VALUES (7,8,9),(1,2,3);`).run();
+		const rawPromise = await db
+			.prepare(`SELECT * FROM abc, cde;`)
+			.raw({ columnNames: true });
+		t.deepEqual(rawPromise, [
+			["a", "b", "c", "c", "d", "e"],
+			[1, 2, 3, 7, 8, 9],
+			[1, 2, 3, 1, 2, 3],
+			[4, 5, 6, 7, 8, 9],
+			[4, 5, 6, 1, 2, 3],
+		]);
+	}
 });
 
 test("operations persist D1 data", async (t) => {
-	const { tableColours, tableKitchenSink } = t.context;
+	const { tableColours, tableKitchenSink, tablePalettes } = t.context;
 
 	// Create new temporary file-system persistence directory
 	const tmp = await useTmp(t);
@@ -405,7 +435,7 @@ test("operations persist D1 data", async (t) => {
 	let db = await getDatabase(mf);
 
 	// Check execute respects persist
-	await db.exec(SCHEMA(tableColours, tableKitchenSink));
+	await db.exec(SCHEMA(tableColours, tableKitchenSink, tablePalettes));
 	await db
 		.prepare(
 			`INSERT INTO ${tableColours} (id, name, rgb) VALUES (4, 'purple', 0xff00ff);`
@@ -431,7 +461,7 @@ test("operations persist D1 data", async (t) => {
 });
 
 test.serial("operations permit strange database names", async (t) => {
-	const { tableColours, tableKitchenSink } = t.context;
+	const { tableColours, tableKitchenSink, tablePalettes } = t.context;
 
 	// Set option, then reset after test
 	const id = "my/ Database";
@@ -441,7 +471,7 @@ test.serial("operations permit strange database names", async (t) => {
 
 	// Check basic operations work
 
-	await db.exec(SCHEMA(tableColours, tableKitchenSink));
+	await db.exec(SCHEMA(tableColours, tableKitchenSink, tablePalettes));
 
 	await db
 		.prepare(
@@ -453,3 +483,23 @@ test.serial("operations permit strange database names", async (t) => {
 		.first<Pick<ColourRow, "name">>();
 	t.deepEqual(result, { name: "pink" });
 });
+
+test("it properly handles ROWS_AND_COLUMNS results format", async (t) => {
+	const { tableColours, tablePalettes } = t.context;
+	const db = await getDatabase(t.context.mf);
+
+	const results = await db
+		.prepare(
+			`SELECT ${tableColours}.name, ${tablePalettes}.name FROM ${tableColours} JOIN ${tablePalettes} ON ${tableColours}.id = ${tablePalettes}.colour_id`
+		)
+		.raw();
+
+	let expectedResults;
+	// Note that this test did not pass with the old binding
+	if (!t.context.bindings["__D1_BETA__DB"]) {
+		expectedResults = [["blue", "Night"]];
+	} else {
+		expectedResults = [["Night"]];
+	}
+	t.deepEqual(results, expectedResults);
+});

packages/miniflare/test/plugins/d1/test.ts

@@ -18,6 +18,8 @@ export interface Context extends MiniflareTestContext {
 	db: D1Database;
 	tableColours: string;
 	tableKitchenSink: string;
+	tablePalettes: string;
+	bindings: Record<string, unknown>;
 }
 
 export let binding: string;