From b79d386088fdfe88e56af5a572e2d73e9b6454ba Mon Sep 17 00:00:00 2001
From: Glen Maddern <glen@cloudflare.com>
Date: Tue, 25 Jul 2023 09:14:43 +1000
Subject: [PATCH] Using .columnNames with .raw(), the full data can be accessed

---
 src/workerd/api/sql-test.js | 33 +++++++++++++++++++++------------
 src/workerd/api/sql.c++     | 11 +++++++++++
 src/workerd/api/sql.h       |  3 +++
 3 files changed, 35 insertions(+), 12 deletions(-)

diff --git a/src/workerd/api/sql-test.js b/src/workerd/api/sql-test.js
index 08caab653f3..c1a38aed363 100644
--- a/src/workerd/api/sql-test.js
+++ b/src/workerd/api/sql-test.js
@@ -510,20 +510,11 @@ async function test(storage) {
     sql.exec(`CREATE TABLE cde (c INT, d INT, e INT);`)
     sql.exec(`INSERT INTO abc VALUES (1,2,3),(4,5,6);`)
     sql.exec(`INSERT INTO cde VALUES (7,8,9),(1,2,3);`)
-    const fullJoin = sql.prepare(`SELECT * FROM abc, cde`)
 
-    // Raw results include both 'c' columns
-    const rawResults = Array.from(fullJoin().raw())
-    assert.equal(rawResults.length, 4)
-    assert.equal(rawResults[0].length, 6)
-    assert.equal(rawResults[1].length, 6)
-    assert.equal(rawResults[2].length, 6)
-    assert.equal(rawResults[3].length, 6)
-
-    // TODO: how to get column names from .raw() iterator?
+    const stmt = sql.prepare(`SELECT * FROM abc, cde`)
 
-    // Without .raw(), data is lost
-    const objResults = Array.from(fullJoin())
+    // In normal iteration, data is lost
+    const objResults = Array.from(stmt())
     assert.equal(Object.values(objResults[0]).length, 5) // duplicate column 'c' dropped
     assert.equal(Object.values(objResults[1]).length, 5) // duplicate column 'c' dropped
     assert.equal(Object.values(objResults[2]).length, 5) // duplicate column 'c' dropped
@@ -533,6 +524,24 @@ async function test(storage) {
     assert.equal(objResults[1].c, 1) // Value of 'c' is the second in the join
     assert.equal(objResults[2].c, 7) // Value of 'c' is the second in the join
     assert.equal(objResults[3].c, 1) // Value of 'c' is the second in the join
+
+    // Iterator has a 'columnNames' property, with .raw() that lets us get the full data
+    const iterator = stmt();
+    assert.deepEqual(iterator.columnNames, ["a","b","c","c","d","e"])
+    const rawResults = Array.from(iterator.raw())
+    assert.equal(rawResults.length, 4)
+    assert.deepEqual(rawResults[0], [1,2,3,7,8,9])
+    assert.deepEqual(rawResults[1], [1,2,3,1,2,3])
+    assert.deepEqual(rawResults[2], [4,5,6,7,8,9])
+    assert.deepEqual(rawResults[3], [4,5,6,1,2,3])
+
+    // Once an iterator is consumed, it can no longer access the columnNames.
+    assert.deepEqual(iterator.columnNames, [])
+
+    // Also works with cursors returned from .exec
+    const execIterator = sql.exec(`SELECT * FROM abc, cde`)
+    assert.deepEqual(execIterator.columnNames, ["a","b","c","c","d","e"])
+    assert.equal(Array.from(execIterator.raw())[0].length, 6);
   }
 
   await scheduler.wait(1);
diff --git a/src/workerd/api/sql.c++ b/src/workerd/api/sql.c++
index b99a2d39d93..e59e818d47e 100644
--- a/src/workerd/api/sql.c++
+++ b/src/workerd/api/sql.c++
@@ -118,6 +118,17 @@ jsg::Ref<SqlStorage::Cursor::RawIterator> SqlStorage::Cursor::raw(jsg::Lock&) {
   return jsg::alloc<RawIterator>(JSG_THIS);
 }
 
+kj::Array<jsg::V8Ref<v8::String>> SqlStorage::Cursor::getColumnNames(jsg::Lock& js) {
+  KJ_IF_MAYBE(s, state) {
+    cachedColumnNames.ensureInitialized(js, (*s)->query);
+    return KJ_MAP(name, this->cachedColumnNames.get()) {
+      return name.addRef(js);
+    };
+  } else {
+    return kj::Array<jsg::V8Ref<v8::String>>();
+  }
+}
+
 kj::Maybe<kj::Array<SqlStorage::Cursor::Value>> SqlStorage::Cursor::rawIteratorNext(
     jsg::Lock& js, jsg::Ref<Cursor>& obj) {
   return iteratorImpl(js, obj,
diff --git a/src/workerd/api/sql.h b/src/workerd/api/sql.h
index dd2ffe7aeef..3044e651fb1 100644
--- a/src/workerd/api/sql.h
+++ b/src/workerd/api/sql.h
@@ -97,11 +97,14 @@ class SqlStorage::Cursor final: public jsg::Object {
         cachedColumnNames(cachedColumnNames) {}
   ~Cursor() noexcept(false);
 
+  kj::Array<jsg::V8Ref<v8::String>> getColumnNames(jsg::Lock& js);
   JSG_RESOURCE_TYPE(Cursor, CompatibilityFlags::Reader flags) {
     JSG_ITERABLE(rows);
     JSG_METHOD(raw);
+    JSG_READONLY_PROTOTYPE_PROPERTY(columnNames, getColumnNames);
   }
 
+
   using Value = kj::Maybe<kj::OneOf<kj::Array<byte>, kj::StringPtr, double>>;
   // One value returned from SQL. Note that we intentionally return StringPtr instead of String
   // because we know that the underlying buffer returned by SQLite will be valid long enough to be