diff --git a/src/cloudflare/internal/d1-api.ts b/src/cloudflare/internal/d1-api.ts index 37461adc5476..357cefdecb20 100644 --- a/src/cloudflare/internal/d1-api.ts +++ b/src/cloudflare/internal/d1-api.ts @@ -16,6 +16,10 @@ type D1Result = D1Response & { results: T[] } +type D1RawOptions = { + columnNames?: boolean +} + type D1UpstreamFailure = { results?: never error: string @@ -85,12 +89,12 @@ class D1Database { public async batch( statements: D1PreparedStatement[] ): Promise[]> { - const exec = await this._sendOrThrow( + const exec = (await this._sendOrThrow( '/query', statements.map((s: D1PreparedStatement) => s.statement), statements.map((s: D1PreparedStatement) => s.params), 'ROWS_AND_COLUMNS' - ) as D1UpstreamSuccess[] + )) as D1UpstreamSuccess[] return exec.map(toArrayOfObjects) } @@ -303,7 +307,7 @@ class D1PreparedStatement { ) } - public async raw(): Promise { + public async raw(options?: D1RawOptions): Promise { const s = firstIfArray( await this.database._sendOrThrow>( '/query', @@ -319,6 +323,9 @@ class D1PreparedStatement { if (Array.isArray(s.results)) { const raw: T[] = [] for (const row of s.results) { + if (options?.columnNames && raw.length === 0) { + raw.push(Array.from(Object.keys(row)) as T) + } const entry = Object.keys(row).map((k) => { return row[k] }) @@ -327,7 +334,10 @@ class D1PreparedStatement { return raw } else { // Otherwise, data is already in the correct format - return s.results.rows as T[] + return [ + ...(options?.columnNames ? [s.results.columns as T] : []), + ...(s.results.rows as T[]), + ] } } } @@ -353,8 +363,9 @@ function toArrayOfObjects(response: D1UpstreamSuccess): D1Result { const { rows, columns } = results return { ...response, - results: rows.map((row) => - Object.fromEntries(row.map((cell, i) => [columns[i], cell])) + results: rows.map( + (row) => + Object.fromEntries(row.map((cell, i) => [columns[i], cell])) as T ), } } diff --git a/src/cloudflare/internal/test/d1/d1-api-test.js b/src/cloudflare/internal/test/d1/d1-api-test.js index 0b71fdc5ee60..3b4859344706 100644 --- a/src/cloudflare/internal/test/d1/d1-api-test.js +++ b/src/cloudflare/internal/test/d1/d1-api-test.js @@ -333,7 +333,7 @@ export const test_d1_api = test(async (DB) => { ) await itShould( - 'not lose data for duplicate columns in a join using raw', + 'not lose data for duplicate columns in a join using raw()', () => DB.prepare(`SELECT * FROM abc, cde;`).raw(), [ [1, 2, 3, 'A', 'B', 'C'], @@ -345,9 +345,39 @@ export const test_d1_api = test(async (DB) => { ] ) + await itShould( + 'add columns using .raw({ columnNames: true })', + () => DB.prepare(`SELECT * FROM abc, cde;`).raw({ columnNames: true }), + [ + ['a', 'b', 'c', 'c', 'd', 'e'], + [1, 2, 3, 'A', 'B', 'C'], + [1, 2, 3, 'D', 'E', 'F'], + [1, 2, 3, 'G', 'H', 'I'], + [4, 5, 6, 'A', 'B', 'C'], + [4, 5, 6, 'D', 'E', 'F'], + [4, 5, 6, 'G', 'H', 'I'], + ] + ) + + await itShould( + 'not add columns using .raw({ columnNames: false })', + () => DB.prepare(`SELECT * FROM abc, cde;`).raw({ columnNames: false }), + [ + [1, 2, 3, 'A', 'B', 'C'], + [1, 2, 3, 'D', 'E', 'F'], + [1, 2, 3, 'G', 'H', 'I'], + [4, 5, 6, 'A', 'B', 'C'], + [4, 5, 6, 'D', 'E', 'F'], + [4, 5, 6, 'G', 'H', 'I'], + ] + ) + await itShould( 'return 0 rows_written for IN clauses', - () => DB.prepare(`SELECT * from cde WHERE c IN ('A','B','C','X','Y','Z')`).all(), + () => + DB.prepare( + `SELECT * from cde WHERE c IN ('A','B','C','X','Y','Z')` + ).all(), { success: true, results: [{ c: 'A', d: 'B', e: 'C' }],