Adding and testing .raw({columnNames: true})

src/cloudflare/internal/d1-api.ts

@@ -16,6 +16,10 @@ type D1Result<T = unknown> = D1Response & {
   results: T[]
 }
 
+type D1RawOptions = {
+  columnNames?: boolean
+}
+
 type D1UpstreamFailure = {
   results?: never
   error: string
@@ -85,12 +89,12 @@ class D1Database {
   public async batch<T = unknown>(
     statements: D1PreparedStatement[]
   ): Promise<D1Result<T>[]> {
-    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<T>[]
+    )) as D1UpstreamSuccess<T>[]
     return exec.map(toArrayOfObjects)
   }
 
@@ -303,7 +307,7 @@ class D1PreparedStatement {
     )
   }
 
-  public async raw<T = unknown[]>(): Promise<T[]> {
+  public async raw<T = unknown[]>(options?: D1RawOptions): Promise<T[]> {
     const s = firstIfArray(
       await this.database._sendOrThrow<Record<string, unknown>>(
         '/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<T>(response: D1UpstreamSuccess<T>): D1Result<T> {
     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
       ),
     }
   }

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' }],