diff --git a/src/workerd/util/sqlite-kv-test.c++ b/src/workerd/util/sqlite-kv-test.c++ index 28a3b4e439a2..960ae58b4cbd 100644 --- a/src/workerd/util/sqlite-kv-test.c++ +++ b/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(largeDataSize); + for (int i = 0; i < largeDataSize; ++i) { + largeData[i] = static_cast(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 diff --git a/src/workerd/util/sqlite-kv.c++ b/src/workerd/util/sqlite-kv.c++ index 4aa252044842..d5713de12725 100644 --- a/src/workerd/util/sqlite-kv.c++ +++ b/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 diff --git a/src/workerd/util/sqlite-kv.h b/src/workerd/util/sqlite-kv.h index 74cc1d74b289..e591273eb969 100644 --- a/src/workerd/util/sqlite-kv.h +++ b/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 state;