Add getDatabaseSize() method to Sqlite-KV to retrieve size of database

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

@@ -91,5 +91,39 @@ KJ_TEST("SQLite-KV") {
   KJ_EXPECT(list(nullptr, kj::none, kj::none, F) == "");
 }
 
+KJ_TEST("SQLite-KV getDatabaseSize") {
+  auto dir = kj::newInMemoryDirectory(kj::nullClock());
+  SqliteDatabase::Vfs vfs(*dir);
+  SqliteDatabase db(vfs, kj::Path({"foo"}), kj::WriteMode::CREATE | kj::WriteMode::MODIFY);
+  SqliteKv kv(db);
+
+  // Get initial database size
+  double initialSize = kv.getDatabaseSize();
+  KJ_EXPECT(initialSize > 0, "Initial database size should be greater than 0");
+
+  // Add first large data item
+  const int largeDataSize = 100 * 1024;  // 100 KB
+  auto largeData = kj::heapArray<kj::byte>(largeDataSize);
+  for (int i = 0; i < largeDataSize; ++i) {
+    largeData[i] = static_cast<kj::byte>(i % 256);
+  }
+  kv.put("large_data_1", largeData.asPtr());
+  double sizeAfterFirstAdd = kv.getDatabaseSize();
+  KJ_EXPECT(sizeAfterFirstAdd > initialSize,
+    "Database size should increase after adding data");
+
+  // Add second large data item
+  kv.put("large_data_2", largeData.asPtr());
+  double sizeAfterSecondAdd = kv.getDatabaseSize();
+  KJ_EXPECT(sizeAfterSecondAdd > sizeAfterFirstAdd,
+    "Database size should increase after adding data");
+
+  // Delete first large data item
+  kv.delete_("large_data_1");
+  double sizeAfterDelete = kv.getDatabaseSize();
+  KJ_EXPECT(sizeAfterDelete < sizeAfterSecondAdd,
+    "Database size should not increase after deleting data");
+}
+
 }  // namespace
 }  // namespace workerd

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

@@ -53,4 +53,8 @@ uint SqliteKv::deleteAll() {
   return query.changeCount();
 }
 
+uint64_t SqliteKv::getDatabaseSize() {
+  return ensureInitialized().getDatabaseSize();
+}
+
 }  // namespace workerd

src/workerd/util/sqlite-kv.h

@@ -45,6 +45,8 @@ class SqliteKv {
 
   uint deleteAll();
 
+  uint64_t getDatabaseSize();
+
   // TODO(perf): Should we provide multi-get, multi-put, and multi-delete? It's a bit tricky to
   //   implement them as single SQL queries, while still using prepared statements. The c-array
   //   extension might help here, though it can only support arrays of NUL-terminated strings, not
@@ -57,6 +59,7 @@ class SqliteKv {
 
   struct Initialized {
     SqliteDatabase& db;
+    const uint pageSize;
 
     SqliteDatabase::Statement stmtGet = db.prepare(R"(
       SELECT value FROM _cf_KV WHERE key = ?
@@ -116,7 +119,19 @@ class SqliteKv {
       DELETE FROM _cf_KV
     )");
 
-    Initialized(SqliteDatabase& db): db(db) {}
+    uint64_t getDatabaseSize() {
+        uint64_t pages = stmtGetDatabaseSize.run().getInt64(0);
+        return pages * pageSize;
+    }
+
+    Initialized(SqliteDatabase& db): db(db),
+      pageSize(db.run("PRAGMA page_size;").getInt64(0)) {
+    }
+
+    private:
+    SqliteDatabase::Statement stmtGetDatabaseSize = db.prepare(R"(
+        SELECT (SELECT * FROM pragma_page_count) - (SELECT * FROM pragma_freelist_count)
+    )");
   };
 
   kj::OneOf<Uninitialized, Initialized> state;