SQLite: Don't bill for internal queries.

We bill for:
- SQL queries issued via the JavaScript SQL API.
- DO old-style key/value storage operations.

Anything else, we assume is an internal query, and not billed.

Author: kentonv

src/workerd/api/sql.c++

@@ -111,6 +111,11 @@ bool SqlStorage::allowTransactions() const {
       "write coalescing.");
 }
 
+bool SqlStorage::shouldAddQueryStats() const {
+  // Bill for queries executed from JavaScript.
+  return true;
+}
+
 SqlStorage::StatementCache::~StatementCache() noexcept(false) {
   for (auto& entry: lru) {
     lru.remove(entry);

src/workerd/api/sql.h

@@ -69,6 +69,7 @@ class SqlStorage final: public jsg::Object, private SqliteDatabase::Regulator {
   bool isAllowedTrigger(kj::StringPtr name) const override;
   void onError(kj::Maybe<int> sqliteErrorCode, kj::StringPtr message) const override;
   bool allowTransactions() const override;
+  bool shouldAddQueryStats() const override;
 
   SqliteDatabase& getDb(jsg::Lock& js) {
     return storage->getSqliteDb(js);

src/workerd/util/sqlite-kv-test.c++

@@ -10,16 +10,32 @@ namespace workerd {
 namespace {
 
 KJ_TEST("SQLite-KV") {
+  class TestSqliteObserver: public SqliteObserver {
+   public:
+    void addQueryStats(uint64_t read, uint64_t written) override {
+      rowsRead += read;
+      rowsWritten += written;
+    }
+
+    uint64_t rowsRead = 0;
+    uint64_t rowsWritten = 0;
+  };
+
   auto dir = kj::newInMemoryDirectory(kj::nullClock());
   SqliteDatabase::Vfs vfs(*dir);
-  SqliteDatabase db(vfs, kj::Path({"foo"}), kj::WriteMode::CREATE | kj::WriteMode::MODIFY);
+  TestSqliteObserver sqliteObserver;
+  SqliteDatabase db(
+      vfs, kj::Path({"foo"}), kj::WriteMode::CREATE | kj::WriteMode::MODIFY, sqliteObserver);
   SqliteKv kv(db);
 
   kv.put("foo", "abc"_kj.asBytes());
   kv.put("bar", "def"_kj.asBytes());
   kv.put("baz", "123"_kj.asBytes());
   kv.put("qux", "321"_kj.asBytes());
 
+  KJ_EXPECT(sqliteObserver.rowsWritten == 4);
+  KJ_EXPECT(sqliteObserver.rowsRead == 0);
+
   {
     bool called = false;
     KJ_EXPECT(kv.get("foo", [&](kj::ArrayPtr<const byte> value) {
@@ -29,6 +45,9 @@ KJ_TEST("SQLite-KV") {
     KJ_EXPECT(called);
   }
 
+  KJ_EXPECT(sqliteObserver.rowsWritten == 4);
+  KJ_EXPECT(sqliteObserver.rowsRead == 1);
+
   {
     bool called = false;
     KJ_EXPECT(kv.get("bar", [&](kj::ArrayPtr<const byte> value) {
@@ -38,10 +57,16 @@ KJ_TEST("SQLite-KV") {
     KJ_EXPECT(called);
   }
 
+  KJ_EXPECT(sqliteObserver.rowsWritten == 4);
+  KJ_EXPECT(sqliteObserver.rowsRead == 2);
+
   KJ_EXPECT(!kv.get("corge", [&](kj::ArrayPtr<const byte> value) {
     KJ_FAIL_EXPECT("should not call callback when no match", value.asChars());
   }));
 
+  KJ_EXPECT(sqliteObserver.rowsWritten == 4);
+  KJ_EXPECT(sqliteObserver.rowsRead == 2);
+
   auto list = [&](auto&&... params) {
     kj::Vector<kj::String> results;
     auto callback = [&](kj::StringPtr key, kj::ArrayPtr<const byte> value) {

src/workerd/util/sqlite-kv.h

@@ -15,6 +15,11 @@ namespace workerd {
 // of KJ exceptions which become internal errors.
 class SqliteKvRegulator: public SqliteDatabase::Regulator {
   void onError(kj::Maybe<int> sqliteErrorCode, kj::StringPtr message) const override;
+
+  // We bill for KV operations as rows read/written.
+  virtual bool shouldAddQueryStats() const override {
+    return true;
+  }
 };
 
 // Class which implements KV storage on top of SQLite. This is intended to be used for Durable
@@ -123,9 +128,7 @@ class SqliteKv: private SqliteDatabase::ResetListener {
       ORDER BY key DESC
       LIMIT ?
     )");
-    // We don't pass in the regulator here because this query doesn't take user input, so
-    // it failing is always our fault.
-    SqliteDatabase::Statement stmtCountKeys = db.prepare(R"(
+    SqliteDatabase::Statement stmtCountKeys = db.prepare(regulator, R"(
       SELECT count(*) FROM _cf_KV
     )");
 

src/workerd/util/sqlite-test.c++

@@ -834,9 +834,17 @@ KJ_TEST("SQLite observer addQueryStats") {
     uint64_t rowsWritten = 0;
   };
 
+  class TestQueryStatsRegulator: public SqliteDatabase::Regulator {
+   public:
+    bool shouldAddQueryStats() const override {
+      return true;
+    }
+  };
+
   TempDirOnDisk dir;
   SqliteDatabase::Vfs vfs(*dir);
   TestSqliteObserver sqliteObserver = TestSqliteObserver();
+  TestQueryStatsRegulator regulator;
   SqliteDatabase db(
       vfs, kj::Path({"foo"}), kj::WriteMode::CREATE | kj::WriteMode::MODIFY, sqliteObserver);
 
@@ -851,17 +859,17 @@ KJ_TEST("SQLite observer addQueryStats") {
   int rowsWrittenBefore = sqliteObserver.rowsWritten;
   constexpr int dbRowCount = 3;
   {
-    db.run("INSERT INTO things (id) VALUES (10)");
-    db.run("INSERT INTO things (id) VALUES (11)");
-    db.run("INSERT INTO things (id) VALUES (12)");
+    db.run(regulator, "INSERT INTO things (id) VALUES (10)");
+    db.run(regulator, "INSERT INTO things (id) VALUES (11)");
+    db.run(regulator, "INSERT INTO things (id) VALUES (12)");
   }
   KJ_EXPECT(sqliteObserver.rowsRead - rowsReadBefore == dbRowCount);
   KJ_EXPECT(sqliteObserver.rowsWritten - rowsWrittenBefore == dbRowCount);
 
   rowsReadBefore = sqliteObserver.rowsRead;
   rowsWrittenBefore = sqliteObserver.rowsWritten;
   {
-    auto getCount = db.prepare("SELECT COUNT(*) FROM things");
+    auto getCount = db.prepare(regulator, "SELECT COUNT(*) FROM things");
     KJ_EXPECT(getCount.run().getInt(0) == dbRowCount);
   }
   KJ_EXPECT(sqliteObserver.rowsRead - rowsReadBefore == dbRowCount);
@@ -871,7 +879,7 @@ KJ_TEST("SQLite observer addQueryStats") {
   rowsReadBefore = sqliteObserver.rowsRead;
   rowsWrittenBefore = sqliteObserver.rowsWritten;
   {
-    auto stmt = db.prepare("SELECT * FROM things");
+    auto stmt = db.prepare(regulator, "SELECT * FROM things");
     auto query = stmt.run();
     KJ_ASSERT(!query.isDone());
     while (!query.isDone()) {
@@ -881,11 +889,24 @@ KJ_TEST("SQLite observer addQueryStats") {
   KJ_EXPECT(sqliteObserver.rowsRead - rowsReadBefore == dbRowCount);
   KJ_EXPECT(sqliteObserver.rowsWritten - rowsWrittenBefore == 0);
 
+  // Verify system queries don't affect stats.
+  rowsReadBefore = sqliteObserver.rowsRead;
+  rowsWrittenBefore = sqliteObserver.rowsWritten;
+  db.run("INSERT INTO things (id) VALUES (13)");
+  {
+    auto query = db.run("SELECT * FROM things");
+    while (!query.isDone()) {
+      query.nextRow();
+    }
+  }
+  KJ_EXPECT(sqliteObserver.rowsRead == rowsReadBefore);
+  KJ_EXPECT(sqliteObserver.rowsWritten == rowsWrittenBefore);
+
   // Verify addQueryStats works correctly when db is reset
   rowsReadBefore = sqliteObserver.rowsRead;
   rowsWrittenBefore = sqliteObserver.rowsWritten;
   {
-    auto query = db.run("INSERT INTO things (id) VALUES (100)");
+    auto query = db.run(regulator, "INSERT INTO things (id) VALUES (100)");
     db.reset();
   }
   KJ_EXPECT(sqliteObserver.rowsRead - rowsReadBefore == 1);

src/workerd/util/sqlite.c++

@@ -1291,8 +1291,10 @@ SqliteDatabase::Query::~Query() noexcept(false) {
 }
 
 void SqliteDatabase::Query::destroy() {
-  //Update the db stats that we have collected for the query
-  db.sqliteObserver.addQueryStats(rowsRead, rowsWritten);
+  if (regulator.shouldAddQueryStats()) {
+    //Update the db stats that we have collected for the query
+    db.sqliteObserver.addQueryStats(rowsRead, rowsWritten);
+  }
 
   // We only need to reset the statement if we don't own it. If we own it, it's about to be
   // destroyed anyway.

src/workerd/util/sqlite.h

@@ -118,6 +118,14 @@ class SqliteDatabase {
     virtual bool allowTransactions() const {
       return true;
     }
+
+    // Whether or not this query's rows read and written should be recorded to the SqliteObserver
+    // when the query is done. (In other words, determines whether this query is billed.)
+    //
+    // We don't bill for TRUSTED queries since they are used internally by the system.
+    virtual bool shouldAddQueryStats() const {
+      return false;
+    }
   };
 
   // Use as the `Regulator&` for queries that are fully trusted. As a general rule, this should