Return row counts for SQL ingestion (storage.sql.ingest()).

D1 bills based on rows read and written, but sql.ingest() didn't
return those. Now it does.

This required changing the return type of sql.ingest(), but it's an
experimental API so that should be okay.

  let sqlCode = "INSERT INTO tbl (col) VALUES (123); INSERT I";
  let result = sql.ingest(sqlCode);

Old:

  result == " INSERT I"

New:

  result.remainder == " INSERT I"
  result.meta.rowsRead == 0
  result.meta.rowsWritten == 1

This "meta" attribute also gives us a convenient spot to stick
additional information. For example, if we wanted to count how many
SQL statements ingest() executed, we could easily add
"result.meta.statementsExecuted" without breaking any existing users.

Implementation-wise, there might be an optimization opportunity around
SqlStorage::IngestResult::getMeta(). Right now, it allocates a new
Meta every time it's called, so if you do the obvious thing like so:

  totalRead += result.meta.rowsRead
  totalWritten += result.meta.rowsWritten
  totalRedFish += result.meta.redFish  // and so on

then that might end up allocating and discarding multiple Meta
objects, and some caching could fix that. On the other hand, I'm not
very familiar[1] with JSG internals, so maybe that's already happening
somehow. I have no idea.

[1] This is a terrible understatement.

Author: smerritt

src/workerd/api/sql-test.js

@@ -62,23 +62,23 @@ async function test(storage) {
 
 
   // Test partial query ingestion
-  assert.deepEqual(sql.ingest(`SELECT 123; SELECT 456;    `), '    ')
-  assert.deepEqual(sql.ingest(`SELECT 123; SELECT 456;`), '')
-  assert.deepEqual(sql.ingest(`SELECT 123; SELECT 456`), ' SELECT 456')
-  assert.deepEqual(sql.ingest(`SELECT 123; SELECT 45`), ' SELECT 45')
-  assert.deepEqual(sql.ingest(`SELECT 123; SELECT 4`), ' SELECT 4')
-  assert.deepEqual(sql.ingest(`SELECT 123; SELECT `), ' SELECT ')
-  assert.deepEqual(sql.ingest(`SELECT 123; SELECT`), ' SELECT')
-  assert.deepEqual(sql.ingest(`SELECT 123; SELEC`), ' SELEC')
-  assert.deepEqual(sql.ingest(`SELECT 123; SELE`), ' SELE')
-  assert.deepEqual(sql.ingest(`SELECT 123; SEL`), ' SEL')
-  assert.deepEqual(sql.ingest(`SELECT 123; SE`), ' SE')
-  assert.deepEqual(sql.ingest(`SELECT 123; S`), ' S')
-  assert.deepEqual(sql.ingest(`SELECT 123; `), ' ')
-  assert.deepEqual(sql.ingest(`SELECT 123;`), '')
-  assert.deepEqual(sql.ingest(`SELECT 123`), 'SELECT 123')
-  assert.deepEqual(sql.ingest(`SELECT 12`), 'SELECT 12')
-  assert.deepEqual(sql.ingest(`SELECT 1`), 'SELECT 1')
+  assert.deepEqual(sql.ingest(`SELECT 123; SELECT 456;    `).remainder, '    ')
+  assert.deepEqual(sql.ingest(`SELECT 123; SELECT 456;`).remainder, '')
+  assert.deepEqual(sql.ingest(`SELECT 123; SELECT 456`).remainder, ' SELECT 456')
+  assert.deepEqual(sql.ingest(`SELECT 123; SELECT 45`).remainder, ' SELECT 45')
+  assert.deepEqual(sql.ingest(`SELECT 123; SELECT 4`).remainder, ' SELECT 4')
+  assert.deepEqual(sql.ingest(`SELECT 123; SELECT `).remainder, ' SELECT ')
+  assert.deepEqual(sql.ingest(`SELECT 123; SELECT`).remainder, ' SELECT')
+  assert.deepEqual(sql.ingest(`SELECT 123; SELEC`).remainder, ' SELEC')
+  assert.deepEqual(sql.ingest(`SELECT 123; SELE`).remainder, ' SELE')
+  assert.deepEqual(sql.ingest(`SELECT 123; SEL`).remainder, ' SEL')
+  assert.deepEqual(sql.ingest(`SELECT 123; SE`).remainder, ' SE')
+  assert.deepEqual(sql.ingest(`SELECT 123; S`).remainder, ' S')
+  assert.deepEqual(sql.ingest(`SELECT 123; `).remainder, ' ')
+  assert.deepEqual(sql.ingest(`SELECT 123;`).remainder, '')
+  assert.deepEqual(sql.ingest(`SELECT 123`).remainder, 'SELECT 123')
+  assert.deepEqual(sql.ingest(`SELECT 12`).remainder, 'SELECT 12')
+  assert.deepEqual(sql.ingest(`SELECT 1`).remainder, 'SELECT 1')
 
   // Exec throws with trailing comments
   assert.throws(
@@ -87,7 +87,7 @@ async function test(storage) {
   )
   // Ingest does not
   assert.deepEqual(
-    sql.ingest(`SELECT 123; SELECT 456; -- trailing comment`),
+    sql.ingest(`SELECT 123; SELECT 456; -- trailing comment`).remainder,
     ' -- trailing comment'
   )
 
@@ -105,8 +105,10 @@ async function test(storage) {
     const inputBytes = new TextEncoder().encode(INSERT_36_ROWS)
     const decoder = new TextDecoder()
 
+
     // Use a chunk size 1, 3, 9, 27, 81, ... bytes
     for (let length = 1; length < inputBytes.length; length = length * 3) {
+      let totalRowsWritten = 0;
       let buffer = ''
       for (let offset = 0; offset < inputBytes.length; offset += length) {
         // Simulate a single "chunk" arriving
@@ -116,7 +118,9 @@ async function test(storage) {
         buffer += decoder.decode(chunk, { stream: true })
 
         // Ingest any complete statements and snip those chars off the buffer
-        buffer = sql.ingest(buffer)
+        let result = sql.ingest(buffer);
+        buffer = result.remainder;
+        totalRowsWritten += result.meta.rowsWritten;
 
         // Simulate awaiting next chunk
         await scheduler.wait(1)
@@ -139,6 +143,10 @@ async function test(storage) {
           { val: 'f: 🔥😎🔥' },
         ]
       )
+
+      // Verify that all 36 rows we inserted were accounted for.
+      assert.equal(totalRowsWritten, 36);
+
       sql.exec(`DELETE FROM streaming`)
       await scheduler.wait(1)
     }
@@ -1030,9 +1038,10 @@ async function testStreamingIngestion(request, storage) {
       buffer += chunk
 
       // Ingest any complete statements and snip those chars off the buffer
-      buffer = sql.ingest(buffer)
+      buffer = sql.ingest(buffer).remainder
     }
   })
+
   // Verify exactly 36 rows were added
   assert.deepEqual(Array.from(sql.exec(`SELECT count(*) FROM streaming`)), [
     { 'count(*)': 36 },

src/workerd/api/sql.c++

@@ -19,9 +19,10 @@ jsg::Ref<SqlStorage::Cursor> SqlStorage::exec(jsg::Lock& js, kj::String querySql
   return jsg::alloc<Cursor>(*sqlite, regulator, querySql, kj::mv(bindings));
 }
 
-kj::String SqlStorage::ingest(jsg::Lock& js, kj::String querySql) {
+jsg::Ref<SqlStorage::IngestResult> SqlStorage::ingest(jsg::Lock& js, kj::String querySql) {
   SqliteDatabase::Regulator& regulator = *this;
-  return kj::str(sqlite->ingestSql(regulator, querySql));
+  auto result = sqlite->ingestSql(regulator, querySql);
+  return jsg::alloc<IngestResult>(kj::str(result.remainder), result.rowsRead, result.rowsWritten);
 }
 
 jsg::Ref<SqlStorage::Statement> SqlStorage::prepare(jsg::Lock& js, kj::String query) {
@@ -295,4 +296,18 @@ void SqlStorage::visitForMemoryInfo(jsg::MemoryTracker& tracker) const {
   }
 }
 
+SqlStorage::IngestResult::IngestResult(kj::String remainder, uint64_t rowsRead, uint64_t rowsWritten) : remainder(kj::mv(remainder)), rowsRead(rowsRead), rowsWritten(rowsWritten) {}
+
+kj::StringPtr SqlStorage::IngestResult::getRemainder() { return remainder; }
+
+jsg::Ref<SqlStorage::IngestResult::Meta>  SqlStorage::IngestResult::getMeta() { return jsg::alloc<Meta>(rowsRead, rowsWritten); }
+
+
+SqlStorage::IngestResult::Meta::Meta(uint64_t rowsRead, uint64_t rowsWritten)  : rowsRead(rowsRead), rowsWritten(rowsWritten) {}
+
+double SqlStorage::IngestResult::Meta::getRowsRead() { return rowsRead; }
+
+double SqlStorage::IngestResult::Meta::getRowsWritten() { return rowsWritten; }
+
+
 }  // namespace workerd::api

src/workerd/api/sql.h

@@ -22,9 +22,10 @@ class SqlStorage final: public jsg::Object, private SqliteDatabase::Regulator {
 
   class Cursor;
   class Statement;
+  class IngestResult;
 
   jsg::Ref<Cursor> exec(jsg::Lock& js, kj::String query, jsg::Arguments<BindingValue> bindings);
-  kj::String ingest(jsg::Lock& js, kj::String query);
+  jsg::Ref<IngestResult> ingest(jsg::Lock& js, kj::String query);
 
   jsg::Ref<Statement> prepare(jsg::Lock& js, kj::String query);
 
@@ -249,10 +250,52 @@ class SqlStorage::Statement final: public jsg::Object {
   friend class Cursor;
 };
 
+
+
+class SqlStorage::IngestResult final : public jsg::Object {
+public:
+  class Meta;
+
+  IngestResult(kj::String remainder, uint64_t rowsRead, uint64_t rowsWritten);
+
+  JSG_RESOURCE_TYPE(IngestResult) {
+    JSG_READONLY_PROTOTYPE_PROPERTY(meta, getMeta);
+    JSG_READONLY_PROTOTYPE_PROPERTY(remainder, getRemainder);
+    JSG_NESTED_TYPE(Meta);
+  }
+
+  jsg::Ref<Meta> getMeta();
+  kj::StringPtr getRemainder();
+
+private:
+  kj::String remainder;
+  uint64_t rowsRead;
+  uint64_t rowsWritten;
+};
+
+class SqlStorage::IngestResult::Meta final : public jsg::Object {
+public:
+  Meta(uint64_t rowsRead, uint64_t rowsWritten);
+
+  JSG_RESOURCE_TYPE(Meta) {
+    JSG_READONLY_PROTOTYPE_PROPERTY(rowsRead, getRowsRead);
+    JSG_READONLY_PROTOTYPE_PROPERTY(rowsWritten, getRowsWritten);
+  }
+
+  double getRowsRead();
+  double getRowsWritten();
+
+private:
+  uint64_t rowsRead;
+  uint64_t rowsWritten;
+};
+
 #define EW_SQL_ISOLATE_TYPES                    \
   api::SqlStorage,                              \
   api::SqlStorage::Statement,                   \
   api::SqlStorage::Cursor,                      \
+  api::SqlStorage::IngestResult,                \
+  api::SqlStorage::IngestResult::Meta,          \
   api::SqlStorage::Cursor::RowIterator,         \
   api::SqlStorage::Cursor::RowIterator::Next,   \
   api::SqlStorage::Cursor::RawIterator,         \

src/workerd/util/sqlite.c++

@@ -446,7 +446,10 @@ kj::Own<sqlite3_stmt> SqliteDatabase::prepareSql(
   }
 }
 
-kj::StringPtr SqliteDatabase::ingestSql(Regulator& regulator, kj::StringPtr sqlCode) {
+SqliteDatabase::IngestResult SqliteDatabase::ingestSql(Regulator& regulator, kj::StringPtr sqlCode) {
+  uint64_t rowsRead = 0;
+  uint64_t rowsWritten = 0;
+
   // While there's still some input SQL to process
   while (sqlCode.begin() != sqlCode.end()) {
     // And there are still valid statements:
@@ -456,13 +459,15 @@ kj::StringPtr SqliteDatabase::ingestSql(Regulator& regulator, kj::StringPtr sqlC
     // Slice off the next valid statement SQL
     auto nextStatement = kj::str(sqlCode.slice(0, statementLength));
     // Create a Query object, which will prepare & execute it
-    Query(*this, regulator, nextStatement);
+    auto q = Query(*this, regulator, nextStatement);
 
+    rowsRead += q.getRowsRead();
+    rowsWritten += q.getRowsWritten();
     sqlCode = sqlCode.slice(statementLength);
   }
 
   // Return the leftover buffer
-  return sqlCode;
+  return {.remainder = sqlCode, .rowsRead = rowsRead, .rowsWritten = rowsWritten};
 }
 
 bool SqliteDatabase::isAuthorized(int actionCode,

src/workerd/util/sqlite.h

@@ -36,6 +36,12 @@ class SqliteDatabase {
   class Regulator;
   struct VfsOptions;
 
+  struct IngestResult {
+    kj::StringPtr remainder;
+    uint64_t rowsRead;
+    uint64_t rowsWritten;
+  };
+
   SqliteDatabase(const Vfs& vfs, kj::PathPtr path, kj::Maybe<kj::WriteMode> maybeMode = kj::none);
   ~SqliteDatabase() noexcept(false);
   KJ_DISALLOW_COPY_AND_MOVE(SqliteDatabase);
@@ -92,7 +98,7 @@ class SqliteDatabase {
   // Helper to execute a chunk of SQL that may not be complete.
   // Executes every valid statement provided, and returns the remaining portion of the input
   // that was not processed. This is used for streaming SQL ingestion.
-  kj::StringPtr ingestSql(Regulator& regulator, kj::StringPtr sqlCode);
+  IngestResult ingestSql(Regulator& regulator, kj::StringPtr sqlCode);
 
 private:
   sqlite3* db;