Skip to content
This repository was archived by the owner on Mar 13, 2025. It is now read-only.

Commit 1f22fa8

Browse files
committed
Re-implement D1 gateway using new storage system
Also brings forward the changes from #533 and #544 Closes DEVX-591
1 parent dd1ce9c commit 1f22fa8

File tree

3 files changed

+49
-33
lines changed

3 files changed

+49
-33
lines changed

packages/tre/src/plugins/d1/gateway.ts

Lines changed: 32 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ import fs from "fs/promises";
33
import os from "os";
44
import path from "path";
55
import { performance } from "perf_hooks";
6-
import { Database as DatabaseType, Statement } from "better-sqlite3";
6+
import { Database as DatabaseType } from "better-sqlite3";
77
import { z } from "zod";
88
import { Response } from "../../http";
99
import { HttpError, Log } from "../../shared";
@@ -135,22 +135,8 @@ function normaliseResults(
135135
);
136136
}
137137

138-
const DOESNT_RETURN_DATA_MESSAGE =
139-
"The columns() method is only for statements that return data";
140138
const EXECUTE_RETURNS_DATA_MESSAGE =
141139
"SQL execute error: Execute returned results - did you mean to call query?";
142-
function returnsData(stmt: Statement): boolean {
143-
try {
144-
stmt.columns();
145-
return true;
146-
} catch (e) {
147-
// `columns()` fails on statements that don't return data
148-
if (e instanceof TypeError && e.message === DOESNT_RETURN_DATA_MESSAGE) {
149-
return false;
150-
}
151-
throw e;
152-
}
153-
}
154140

155141
const CHANGES_QUERY = "SELECT total_changes() AS totalChanges";
156142
const CHANGES_LAST_ROW_QUERY =
@@ -167,8 +153,30 @@ interface ChangesLastRowResult {
167153
export class D1Gateway {
168154
private readonly db: DatabaseType;
169155

170-
constructor(private readonly log: Log, private readonly storage: Storage) {
171-
this.db = storage.getSqliteDatabase();
156+
constructor(private readonly log: Log, legacyStorage: Storage) {
157+
const storage = legacyStorage.getNewStorage();
158+
this.db = storage.db;
159+
}
160+
161+
#prepareAndBind(query: D1SingleQuery) {
162+
// D1 only respects the first statement
163+
const sql = splitSqlQuery(query.sql)[0];
164+
const stmt = this.db.prepare(sql);
165+
const params = normaliseParams(query.params);
166+
if (params.length === 0) return stmt;
167+
168+
try {
169+
return stmt.bind(params);
170+
} catch (e) {
171+
// For statements using ?1, ?2, etc, we want to pass them as an array but
172+
// `better-sqlite3` expects an object with the shape:
173+
// `{ 1: params[0], 2: params[1], ... }`. Try bind like that instead.
174+
try {
175+
return stmt.bind(Object.fromEntries(params.map((v, i) => [i + 1, v])));
176+
} catch {}
177+
// If that still failed, re-throw the original error
178+
throw e;
179+
}
172180
}
173181

174182
#getTotalChanges(): number | bigint {
@@ -181,18 +189,15 @@ export class D1Gateway {
181189

182190
#query: QueryRunner = (query) => {
183191
const meta: OkMeta = { start: performance.now() };
184-
// D1 only respects the first statement
185-
const sql = splitSqlQuery(query.sql)[0];
186-
const stmt = this.db.prepare(sql);
187-
const params = normaliseParams(query.params);
192+
const stmt = this.#prepareAndBind(query);
188193
let results: Record<string, SqliteValue>[];
189-
if (returnsData(stmt)) {
194+
if (stmt.reader) {
190195
// `better-sqlite3` doesn't return `last_row_id` and `changes` from `all`.
191196
// We need to make extra queries to get them, but we only want to return
192197
// them if this `stmt` made changes. So check total changes before and
193198
// after querying `stmt`.
194199
const initialTotalChanges = this.#getTotalChanges();
195-
results = stmt.all(params);
200+
results = stmt.all();
196201
const { totalChanges, changes, lastRowId } = this.#getChangesLastRow();
197202
if (totalChanges > initialTotalChanges) {
198203
meta.last_row_id = Number(lastRowId);
@@ -201,7 +206,7 @@ export class D1Gateway {
201206
} else {
202207
// `/query` does support queries that don't return data,
203208
// returning `[]` instead of `null`
204-
const result = stmt.run(params);
209+
const result = stmt.run();
205210
results = [];
206211
meta.last_row_id = Number(result.lastInsertRowid);
207212
meta.changes = result.changes;
@@ -211,13 +216,10 @@ export class D1Gateway {
211216

212217
#execute: QueryRunner = (query) => {
213218
const meta: OkMeta = { start: performance.now() };
214-
// D1 only respects the first statement
215-
const sql = splitSqlQuery(query.sql)[0];
216-
const stmt = this.db.prepare(sql);
219+
const stmt = this.#prepareAndBind(query);
217220
// `/execute` only supports queries that don't return data
218-
if (returnsData(stmt)) throw new Error(EXECUTE_RETURNS_DATA_MESSAGE);
219-
const params = normaliseParams(query.params);
220-
const result = stmt.run(params);
221+
if (stmt.reader) throw new Error(EXECUTE_RETURNS_DATA_MESSAGE);
222+
const result = stmt.run();
221223
meta.last_row_id = Number(result.lastInsertRowid);
222224
meta.changes = result.changes;
223225
return ok(null, meta);

packages/tre/src/storage2/sql.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ export type TypedStatement<
1313
};
1414

1515
export type TypedDatabase = Omit<Database, "prepare"> & {
16-
prepare<Params, SingleResult = unknown>(
16+
prepare<Params = any[], SingleResult = unknown>(
1717
source: string
1818
): Params extends any[]
1919
? TypedStatement<Params, SingleResult>

packages/tre/test/plugins/d1/index.spec.ts

Lines changed: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -272,6 +272,18 @@ test("D1PreparedStatement: bind", async (t) => {
272272
.bind("green")
273273
.all<ColourRow>();
274274
t.is(colourResults.results?.length, 1);
275+
276+
// Check with numbered parameters (execute and query)
277+
// https://github.com/cloudflare/miniflare/issues/504
278+
await db
279+
.prepare(`INSERT INTO ${tableColours} (id, name, rgb) VALUES (?3, ?1, ?2)`)
280+
.bind("yellow", 0xffff00, 4)
281+
.run();
282+
const colourResult = await db
283+
.prepare(`SELECT * FROM ${tableColours} WHERE id = ?1`)
284+
.bind(4)
285+
.first<ColourRow>();
286+
t.deepEqual(colourResult, { id: 4, name: "yellow", rgb: 0xffff00 });
275287
});
276288

277289
// Lots of strange edge cases here...
@@ -456,8 +468,10 @@ test.serial("operations persist D1 data", async (t) => {
456468

457469
// Create new temporary file-system persistence directory
458470
const tmp = await useTmp(t);
459-
const storage = new FileStorage(path.join(tmp, "db"));
460-
const sqliteDb = storage.getSqliteDatabase();
471+
// TODO(soon): clean up this mess once we've migrated all gateways
472+
const legacyStorage = new FileStorage(path.join(tmp, "db"));
473+
const newStorage = legacyStorage.getNewStorage();
474+
const sqliteDb = newStorage.db;
461475

462476
// Set option, then reset after test
463477
await t.context.setOptions({ ...opts, d1Persist: tmp });

0 commit comments

Comments
 (0)