diff --git a/src/workerd/api/sql.c++ b/src/workerd/api/sql.c++ index baf21d6a81f..0eeba4011ad 100644 --- a/src/workerd/api/sql.c++ +++ b/src/workerd/api/sql.c++ @@ -7,6 +7,7 @@ #include "actor-state.h" #include +#include #if _WIN32 #define strncasecmp _strnicmp @@ -72,6 +73,8 @@ jsg::Ref SqlStorage::exec( }; } + auto queryId = db.generateQueryId(); + if (slot->isShared()) { // Oops, this CachedStatement is currently in-use (presumably by a Cursor). // @@ -81,10 +84,12 @@ jsg::Ref SqlStorage::exec( // In theory we could try to cache multiple copies of the statement, but as this is probably // exceedingly rare, it is not worth the added code complexity. return js.alloc(js, kj::mv(doneCallback), db, - SqliteDatabase::StaticRegulator(regulator), js.toString(querySql), kj::mv(bindings)); + SqliteDatabase::StaticRegulator(regulator), js.toString(querySql), kj::mv(bindings), + queryId); } - auto result = js.alloc(js, kj::mv(doneCallback), slot.addRef(), kj::mv(bindings)); + auto result = + js.alloc(js, kj::mv(doneCallback), slot.addRef(), kj::mv(bindings), queryId); // If the statement cache grew too big, drop the least-recently-used entry. while (statementCache.totalSize > SQL_STATEMENT_CACHE_MAX_SIZE) { @@ -195,15 +200,19 @@ jsg::JsValue SqlStorage::wrapSqlValue(jsg::Lock& js, SqlValue value) { SqlStorage::Cursor::State::State(SqliteDatabase& db, SqliteDatabase::StaticRegulator regulator, kj::StringPtr sqlCode, - kj::Array bindingsParam) + kj::Array bindingsParam, + kj::Maybe queryId) : bindings(kj::mv(bindingsParam)), - query(db.run({.regulator = regulator}, sqlCode, mapBindings(bindings).asPtr())) {} + query(db.run( + {.regulator = regulator, .queryId = queryId}, sqlCode, mapBindings(bindings).asPtr())) {} -SqlStorage::Cursor::State::State( - kj::Rc cachedStatementParam, kj::Array bindingsParam) +SqlStorage::Cursor::State::State(kj::Rc cachedStatementParam, + kj::Array bindingsParam, + kj::Maybe queryId) : bindings(kj::mv(bindingsParam)), query(cachedStatement.emplace(kj::mv(cachedStatementParam)) - ->statement.run(mapBindings(bindings).asPtr())) {} + ->statement.run(SqliteDatabase::Statement::StatementOptions{.queryId = queryId}, + mapBindings(bindings).asPtr())) {} SqlStorage::Cursor::~Cursor() noexcept(false) { // If this Cursor was created from a Statement, clear the Statement's currentCursor weak ref. diff --git a/src/workerd/api/sql.h b/src/workerd/api/sql.h index 36aee44395f..4745b9039a3 100644 --- a/src/workerd/api/sql.h +++ b/src/workerd/api/sql.h @@ -257,9 +257,12 @@ class SqlStorage::Cursor final: public jsg::Object { State(SqliteDatabase& db, SqliteDatabase::StaticRegulator regulator, kj::StringPtr sqlCode, - kj::Array bindings); + kj::Array bindings, + kj::Maybe queryId); - State(kj::Rc cachedStatement, kj::Array bindings); + State(kj::Rc cachedStatement, + kj::Array bindings, + kj::Maybe queryId); }; // Nulled out when query is done or canceled. diff --git a/src/workerd/util/BUILD.bazel b/src/workerd/util/BUILD.bazel index e2bb295d90b..1d56cbe0eca 100644 --- a/src/workerd/util/BUILD.bazel +++ b/src/workerd/util/BUILD.bazel @@ -244,6 +244,7 @@ wd_cc_library( ], implementation_deps = [ ":autogate", + ":entropy", ":strings", "//src/workerd/jsg:exception", "@sqlite3", diff --git a/src/workerd/util/sqlite-kv.c++ b/src/workerd/util/sqlite-kv.c++ index 7644878de9d..6dd044b8b30 100644 --- a/src/workerd/util/sqlite-kv.c++ +++ b/src/workerd/util/sqlite-kv.c++ @@ -95,37 +95,43 @@ kj::Own SqliteKv::list( if (!tableCreated) return kj::heap(nullptr); auto& stmts = KJ_UNWRAP_OR(state.tryGet(), return kj::heap(nullptr)); + // Mint one id for this list so every statement it runs shares it. + SqliteDatabase::Statement::StatementOptions runOptions{.queryId = db.generateQueryId()}; + if (order == Order::FORWARD) { KJ_IF_SOME(e, end) { KJ_IF_SOME(l, limit) { - return kj::heap(kj::Badge(), *this, stmts.stmtListEndLimit, begin, e, - static_cast(l)); + return kj::heap(kj::Badge(), *this, stmts.stmtListEndLimit, + runOptions, begin, e, static_cast(l)); } else { - return kj::heap(kj::Badge(), *this, stmts.stmtListEnd, begin, e); + return kj::heap( + kj::Badge(), *this, stmts.stmtListEnd, runOptions, begin, e); } } else { KJ_IF_SOME(l, limit) { - return kj::heap( - kj::Badge(), *this, stmts.stmtListLimit, begin, static_cast(l)); + return kj::heap(kj::Badge(), *this, stmts.stmtListLimit, runOptions, + begin, static_cast(l)); } else { - return kj::heap(kj::Badge(), *this, stmts.stmtList, begin); + return kj::heap( + kj::Badge(), *this, stmts.stmtList, runOptions, begin); } } } else { KJ_IF_SOME(e, end) { KJ_IF_SOME(l, limit) { return kj::heap(kj::Badge(), *this, stmts.stmtListEndLimitReverse, - begin, e, static_cast(l)); + runOptions, begin, e, static_cast(l)); } else { return kj::heap( - kj::Badge(), *this, stmts.stmtListEndReverse, begin, e); + kj::Badge(), *this, stmts.stmtListEndReverse, runOptions, begin, e); } } else { KJ_IF_SOME(l, limit) { - return kj::heap(kj::Badge(), *this, stmts.stmtListLimitReverse, begin, - static_cast(l)); + return kj::heap(kj::Badge(), *this, stmts.stmtListLimitReverse, + runOptions, begin, static_cast(l)); } else { - return kj::heap(kj::Badge(), *this, stmts.stmtListReverse, begin); + return kj::heap( + kj::Badge(), *this, stmts.stmtListReverse, runOptions, begin); } } } @@ -151,9 +157,10 @@ void SqliteKv::put(KeyPtr key, ValuePtr value) { } void SqliteKv::put(KeyPtr key, ValuePtr value, WriteOptions options) { + auto queryId = resolveQueryId(options.queryId); ensureInitialized(options.allowUnconfirmed) - .stmtPut.run({.allowUnconfirmed = options.allowUnconfirmed}, key, value); - clearExternalsIfPresent(key); + .stmtPut.run({.allowUnconfirmed = options.allowUnconfirmed, .queryId = queryId}, key, value); + clearExternalsIfPresent(key, queryId); } bool SqliteKv::delete_(KeyPtr key) { @@ -161,18 +168,20 @@ bool SqliteKv::delete_(KeyPtr key) { } bool SqliteKv::delete_(KeyPtr key, WriteOptions options) { + auto queryId = resolveQueryId(options.queryId); bool result; { auto query = ensureInitialized(options.allowUnconfirmed) - .stmtDelete.run({.allowUnconfirmed = options.allowUnconfirmed}, key); + .stmtDelete.run( + {.allowUnconfirmed = options.allowUnconfirmed, .queryId = queryId}, key); result = query.changeCount() > 0; } - clearExternalsIfPresent(key); + clearExternalsIfPresent(key, queryId); return result; } -void SqliteKv::clearExternalsIfPresent(KeyPtr key) { +void SqliteKv::clearExternalsIfPresent(KeyPtr key, SqliteQueryId queryId) { // If the externals table hasn't been created yet, there's nothing to clear. We deliberately // avoid creating it here -- it should only be created when externals are actually being // written. @@ -184,14 +193,19 @@ void SqliteKv::clearExternalsIfPresent(KeyPtr key) { // Externals writes are always paired with a regular KV write, and that paired write decides // whether the operation is allowed to be unconfirmed. So we always pass allowUnconfirmed here. - stmts.stmtDeleteExternals.run({.allowUnconfirmed = true}, key); + stmts.stmtDeleteExternals.run({.allowUnconfirmed = true, .queryId = queryId}, key); } uint SqliteKv::deleteAll() { // TODO(perf): Consider introducing a compatibility flag that causes deleteAll() to always return // 1. Apps almost certainly don't care about the return value but historically we returned the // count of keys deleted, so now we're stuck counting the table size for no good reason. - uint count = tableCreated ? ensureInitialized(false).stmtCountKeys.run().getInt(0) : 0; + uint count = tableCreated + ? ensureInitialized(false) + .stmtCountKeys + .run(SqliteDatabase::Statement::StatementOptions{.queryId = db.generateQueryId()}) + .getInt(0) + : 0; db.reset(); return count; } @@ -233,7 +247,8 @@ kj::Array> SqliteKv::getExternals(kj::StringPtr key) { if (!externalsTableCreated) return nullptr; auto& stmts = KJ_UNWRAP_OR(externalsState.tryGet(), return nullptr); - auto query = stmts.stmtGetExternals.run(key); + auto query = stmts.stmtGetExternals.run( + SqliteDatabase::Statement::StatementOptions{.queryId = db.generateQueryId()}, key); kj::Vector> result; while (!query.isDone()) { @@ -249,12 +264,15 @@ void SqliteKv::putExternals(kj::StringPtr key, kj::Array> tokens // Externals writes are always paired with a regular KV write, and that paired write decides // whether the operation is allowed to be unconfirmed. So we always pass allowUnconfirmed here. + // Group every statement of this putExternals under one id. + auto queryId = db.generateQueryId(); + // Replace any existing tokens for this key with the new set. - stmts.stmtDeleteExternals.run({.allowUnconfirmed = true}, key); + stmts.stmtDeleteExternals.run({.allowUnconfirmed = true, .queryId = queryId}, key); for (auto i: kj::indices(tokens)) { - stmts.stmtPutExternal.run( - {.allowUnconfirmed = true}, key, static_cast(i), tokens[i].asPtr()); + stmts.stmtPutExternal.run({.allowUnconfirmed = true, .queryId = queryId}, key, + static_cast(i), tokens[i].asPtr()); } } @@ -267,13 +285,14 @@ void SqliteKv::beforeSqliteReset() { void SqliteKv::rollbackMultiPut(Initialized& stmts, WriteOptions options) { KJ_IF_SOME(e, kj::runCatchingExceptions([&]() { // This should be rare, so we don't prepare a statement for it. - stmts.db.run( - {.regulator = Initialized::regulator, .allowUnconfirmed = options.allowUnconfirmed}, + stmts.db.run({.regulator = Initialized::regulator, + .allowUnconfirmed = options.allowUnconfirmed, + .queryId = options.queryId}, kj::str("ROLLBACK TO _cf_put_multiple_savepoint")); - stmts.stmtMultiPutRelease.run({.allowUnconfirmed = options.allowUnconfirmed}); + stmts.stmtMultiPutRelease.run( + {.allowUnconfirmed = options.allowUnconfirmed, .queryId = options.queryId}); })) { KJ_LOG(WARNING, "silencing exception encountered while rolling back multi-put", e); } } - } // namespace workerd diff --git a/src/workerd/util/sqlite-kv.h b/src/workerd/util/sqlite-kv.h index 5292c990c77..a2e85670a41 100644 --- a/src/workerd/util/sqlite-kv.h +++ b/src/workerd/util/sqlite-kv.h @@ -59,6 +59,10 @@ class SqliteKv: private SqliteDatabase::ResetListener { struct WriteOptions { bool allowUnconfirmed = false; + + // Groups all statements of one top-level KV operation under a single id. Left null by callers; + // minted by the operation and threaded into any statements/sub-operations it spawns. + kj::Maybe queryId = kj::none; }; // Store a value into the table. @@ -220,10 +224,20 @@ class SqliteKv: private SqliteDatabase::ResetListener { // Delete all externals associated with `key`, if the _cf_EXTERNALS table exists. No-op if the // table has never been created. Called automatically from put() and delete_(). Always runs // with `allowUnconfirmed = true` since the paired KV write decides confirmation semantics. - void clearExternalsIfPresent(KeyPtr key); + void clearExternalsIfPresent(KeyPtr key, SqliteQueryId queryId); void beforeSqliteReset() override; + // Returns `queryId` if set, otherwise mints a fresh one. Used so every statement of a single + // top-level KV operation shares one grouping id, while recursive operations (e.g. multi-put) + // reuse the id minted by the outer operation. + SqliteQueryId resolveQueryId(kj::Maybe queryId) { + KJ_IF_SOME(id, queryId) { + return id; + } + return db.generateQueryId(); + } + // Helper function that rolls back a multi-put statement and swallows any exceptions that may // occur during the rollback. void rollbackMultiPut(Initialized& stmts, WriteOptions options); @@ -300,7 +314,8 @@ bool SqliteKv::get(KeyPtr key, Func&& callback) { if (!tableCreated) return false; auto& stmts = KJ_UNWRAP_OR(state.tryGet(), return false); - auto query = stmts.stmtGet.run(key); + auto query = stmts.stmtGet.run( + SqliteDatabase::Statement::StatementOptions{.queryId = db.generateQueryId()}, key); if (query.isDone()) { return false; @@ -320,18 +335,24 @@ template void SqliteKv::put(ArrayOfKeyValuePair& pairs, WriteOptions options) { // TODO(cleanup): This code is very similar to DurableObjectStorage::transactionSync. Perhaps the // general structure can be shared somehow? + // Mint one id for the whole multi-put so every statement it issues (savepoint, each put, and + // release) shares it. + options.queryId = resolveQueryId(options.queryId); auto& stmts = ensureInitialized(options.allowUnconfirmed); - stmts.stmtMultiPutSavepoint.run({.allowUnconfirmed = options.allowUnconfirmed}); + stmts.stmtMultiPutSavepoint.run( + {.allowUnconfirmed = options.allowUnconfirmed, .queryId = options.queryId}); { // If any of the puts throw an exception, rollback the transaction and re-throw the exception // from the put that failed. KJ_ON_SCOPE_FAILURE(rollbackMultiPut(stmts, options)); for (const auto& pair: pairs) { - put(pair.key, pair.value, {.allowUnconfirmed = options.allowUnconfirmed}); + put(pair.key, pair.value, + {.allowUnconfirmed = options.allowUnconfirmed, .queryId = options.queryId}); } } - stmts.stmtMultiPutRelease.run({.allowUnconfirmed = options.allowUnconfirmed}); + stmts.stmtMultiPutRelease.run( + {.allowUnconfirmed = options.allowUnconfirmed, .queryId = options.queryId}); } } // namespace workerd diff --git a/src/workerd/util/sqlite-test.c++ b/src/workerd/util/sqlite-test.c++ index fa50abc5e3b..71aa49b857a 100644 --- a/src/workerd/util/sqlite-test.c++ +++ b/src/workerd/util/sqlite-test.c++ @@ -118,6 +118,7 @@ class MockSqliteObserver: public SqliteObserver { int extendedErrorCode; bool isInternalQuery; kj::Maybe queryErrorDescription; + kj::Maybe queryId; }; uint64_t rowsRead = 0; @@ -137,7 +138,8 @@ class MockSqliteObserver: public SqliteObserver { int queryError, int extendedErrorCode, bool isInternalQuery, - kj::Maybe queryErrorDescription) override { + kj::Maybe queryErrorDescription, + kj::Maybe queryId) override { events.add(Event{ .queryStatement = kj::mv(queryStatement), .queryRowsRead = queryRowsRead, @@ -148,6 +150,7 @@ class MockSqliteObserver: public SqliteObserver { .extendedErrorCode = extendedErrorCode, .isInternalQuery = isInternalQuery, .queryErrorDescription = kj::mv(queryErrorDescription), + .queryId = queryId, }); } }; @@ -983,6 +986,27 @@ KJ_TEST("SQLite observer addQueryStats") { } KJ_TEST("SQLite observer reportQueryEvent") { + class TestSqliteObserver: public SqliteObserver { + public: + int capturedEvents = 0; + + void reportQueryEvent(kj::Maybe queryStatement, + uint64_t queryRowsRead, + uint64_t queryRowsWritten, + kj::Duration, + uint64_t dbWalBytesWritten, + int queryError, + int extendedErrorCode, + bool isInternalQuery, + kj::Maybe queryErrorDescription, + kj::Maybe queryId) override { + KJ_IF_SOME(err, queryErrorDescription) { + KJ_ASSERT(err.contains("query canceled because reset()")); + } + capturedEvents++; + } + }; + auto dir = kj::newInMemoryDirectory(kj::nullClock()); SqliteDatabase::Vfs vfs(*dir); MockSqliteObserver sqliteObserver; @@ -1097,6 +1121,45 @@ KJ_TEST("SQLite failed statement reset") { KJ_EXPECT(db.run("SELECT COUNT(*) FROM things").getInt(0) == 4); } +KJ_TEST("SQLite observer reportQueryEvent groups statements by queryId") { + auto dir = kj::newInMemoryDirectory(kj::nullClock()); + SqliteDatabase::Vfs vfs(*dir); + MockSqliteObserver sqliteObserver; + SqliteDatabase db(vfs, kj::Path({"foo"}), kj::WriteMode::CREATE | kj::WriteMode::MODIFY, + /*sqliteMaxMemoryBytes=*/kj::maxValue, /*sqliteMaxMemoryPerProcessBytes=*/kj::maxValue, + sqliteObserver); + + auto& regulator = DEFAULT_REGULATOR_FOR_TEST; + + // A multi-statement operation stamps every statement it runs with the one id it was given. The + // leading statements are reported from prepareSql(); the final statement is reported when the + // Query is destroyed. + auto id1 = db.generateQueryId(); + { + db.run({.regulator = regulator, .queryId = id1}, + "CREATE TABLE t (a INTEGER); INSERT INTO t VALUES (1); SELECT * FROM t"_kj); + } + + KJ_ASSERT(sqliteObserver.events.size() == 3); + for (auto& event: sqliteObserver.events) { + KJ_EXPECT(KJ_ASSERT_NONNULL(event.queryId) == id1); + } + + // A second top-level operation gets a distinct id, and nextQueryId() keeps advancing. + auto id2 = db.generateQueryId(); + KJ_EXPECT(id2 != id1); + size_t before = sqliteObserver.events.size(); + { db.run({.regulator = regulator, .queryId = id2}, "SELECT * FROM t"_kj); } + KJ_ASSERT(sqliteObserver.events.size() == before + 1); + KJ_EXPECT(KJ_ASSERT_NONNULL(sqliteObserver.events[before].queryId) == id2); + + // An operation run without an id leaves queryId unset. + before = sqliteObserver.events.size(); + { db.run({.regulator = regulator}, "SELECT * FROM t"_kj); } + KJ_ASSERT(sqliteObserver.events.size() == before + 1); + KJ_EXPECT(sqliteObserver.events[before].queryId == kj::none); +} + KJ_TEST("SQLite extended error codes in messages") { // Verify that error messages include named extended error codes when they differ from the // primary error code. diff --git a/src/workerd/util/sqlite.c++ b/src/workerd/util/sqlite.c++ index 0c0ab26e748..30041bd062f 100644 --- a/src/workerd/util/sqlite.c++ +++ b/src/workerd/util/sqlite.c++ @@ -8,6 +8,7 @@ #include #include +#include #include #include @@ -831,13 +832,23 @@ void SqliteDatabase::applyChange(const StateChange& change) { } } +SqliteQueryId SqliteDatabase::generateQueryId() { + SqliteQueryId id = 0; + // We wanna avoid zero as 0 can be used as an unset value + while (id == 0) { + getEntropy(kj::asBytes(id)); + } + return id; +} + // Set up the regulator that will be used for authorizer callbacks while preparing this // statement. SqliteDatabase::StatementAndEffect SqliteDatabase::prepareSql(StaticRegulator regulator, kj::StringPtr sqlCode, uint prepFlags, Multi multi, - kj::Maybe&> prelude) { + kj::Maybe&> prelude, + kj::Maybe queryId) { sqlite3* db = &KJ_ASSERT_NONNULL(maybeDb, "previous reset() failed"); ParseContext parseContext; @@ -884,7 +895,7 @@ SqliteDatabase::StatementAndEffect SqliteDatabase::prepareSql(StaticRegulator re sqliteObserver.reportQueryEvent( kj::heapString(sqlCode.slice(0, kj::min(RA_MAX_METRICS_QUERY_SIZE, sqlCode.size()))), 0, 0, /* Reports queryLatency to be 0 explicitly */ 0 * kj::NANOSECONDS, 0, prepareResult, - extendedErrorCode, !regulator->shouldAddQueryStats(), kj::none); + extendedErrorCode, !regulator->shouldAddQueryStats(), kj::none, queryId); } // If we had an auth error specifically, check if we recorded a better error message during @@ -964,7 +975,7 @@ SqliteDatabase::StatementAndEffect SqliteDatabase::prepareSql(StaticRegulator re // Report queryEvent for this statement sqliteObserver.reportQueryEvent(kj::mv(queryStatement), rowsRead, rowsWritten, queryLatency, dbWalBytesWritten, err, extendedCode, - regulator->shouldAddQueryStats(), kj::mv(queryErrorDescription)); + regulator->shouldAddQueryStats(), kj::mv(queryErrorDescription), queryId); if (err == SQLITE_DONE) { // good @@ -1501,9 +1512,10 @@ SqliteDatabase::Statement SqliteDatabase::prepare( *this, regulator, prepareSql(regulator, sqlCode, SQLITE_PREPARE_PERSISTENT, SINGLE)); } -SqliteDatabase::StatementAndEffect& SqliteDatabase::Statement::prepareForExecution() { +SqliteDatabase::StatementAndEffect& SqliteDatabase::Statement::prepareForExecution( + kj::Maybe queryId) { for (auto& stmt: prelude) { - stmt.run(); + stmt.run(StatementOptions{.queryId = queryId}); } KJ_IF_SOME(sqlCode, stmt.tryGet()) { @@ -1518,7 +1530,8 @@ SqliteDatabase::StatementAndEffect& SqliteDatabase::Statement::prepareForExecuti // We use the MULTI flag here in case this Statement was created by prepareMulti(). If multiple // statements are parsed, they'll be added to our `prelude`, and also executed immediately. - stmt = db.prepareSql(regulator, sqlCodeToPrepare, SQLITE_PREPARE_PERSISTENT, MULTI, prelude); + stmt = db.prepareSql( + regulator, sqlCodeToPrepare, SQLITE_PREPARE_PERSISTENT, MULTI, prelude, queryId); } return KJ_ASSERT_NONNULL(stmt.tryGet()); @@ -1545,10 +1558,11 @@ SqliteDatabase::Query::Query(SqliteDatabase& db, kj::ArrayPtr bindings) : ResetListener(db), regulator(options.regulator), - maybeStatement(statement.prepareForExecution()), + maybeStatement(statement.prepareForExecution(options.queryId)), queryEvent(this->db.sqliteObserver) { // If we throw from the constructor, the destructor won't run. Need to call destroy() explicitly. KJ_ON_SCOPE_FAILURE(destroy()); + queryEvent.setQueryId(options.queryId); init(bindings); } @@ -1558,11 +1572,12 @@ SqliteDatabase::Query::Query(SqliteDatabase& db, kj::ArrayPtr bindings) : ResetListener(db), regulator(options.regulator), - ownStatement(db.prepareSql(regulator, sqlCode, 0, MULTI)), + ownStatement(db.prepareSql(regulator, sqlCode, 0, MULTI, kj::none, options.queryId)), maybeStatement(ownStatement), queryEvent(this->db.sqliteObserver) { // If we throw from the constructor, the destructor won't run. Need to call destroy() explicitly. KJ_ON_SCOPE_FAILURE(destroy()); + queryEvent.setQueryId(options.queryId); init(bindings); } diff --git a/src/workerd/util/sqlite.h b/src/workerd/util/sqlite.h index 42ca9793a3c..bbe1e538948 100644 --- a/src/workerd/util/sqlite.h +++ b/src/workerd/util/sqlite.h @@ -26,6 +26,11 @@ namespace workerd { using kj::byte; using kj::uint; +// Identifies the top-level operation (i.e. a single `exec()` or KV op) that a query belongs to, +// so related statements can be grouped together in analytics. Minted by callers at the operation +// boundary and threaded down into every Query the operation spawns. +using SqliteQueryId = uint64_t; + // Used to collect periodic metrics about queries and size of sqlite db class SqliteObserver { public: @@ -50,7 +55,8 @@ class SqliteObserver { int queryResult, int extendedErrorCode, bool isInternalQuery, - kj::Maybe queryErrorDescription) {} + kj::Maybe queryErrorDescription, + kj::Maybe queryId) {} static SqliteObserver DEFAULT; @@ -105,6 +111,7 @@ class SqliteDatabase { struct QueryOptions { StaticRegulator regulator; bool allowUnconfirmed = false; + kj::Maybe queryId = kj::none; }; struct IngestResult { @@ -356,6 +363,10 @@ class SqliteDatabase { } } + // Generates a randomly generated fresh id for a new top-level operation so related statements + // can be grouped in analytics. + SqliteQueryId generateQueryId(); + private: const Vfs& vfs; kj::Path path; @@ -467,7 +478,8 @@ class SqliteDatabase { kj::StringPtr sqlCode, uint prepFlags, Multi multi, - kj::Maybe&> prelude = kj::none); + kj::Maybe&> prelude = kj::none, + kj::Maybe queryId = kj::none); // Implements SQLite authorizer callback, see sqlite3_set_authorizer(). bool isAuthorized(int actionCode, @@ -523,6 +535,7 @@ class SqliteDatabase::Statement final: private ResetListener { struct StatementOptions { bool allowUnconfirmed = false; + kj::Maybe queryId = kj::none; }; template @@ -551,7 +564,7 @@ class SqliteDatabase::Statement final: private ResetListener { // Get the underlying StatementAndEffect, which the caller will then execute. If `prelude` is // non-empty, prepareForExecution() actually executes the prelude. - StatementAndEffect& prepareForExecution(); + StatementAndEffect& prepareForExecution(kj::Maybe queryId = kj::none); friend class SqliteDatabase; }; @@ -675,7 +688,7 @@ class SqliteDatabase::Query final: private ResetListener { observer.reportQueryEvent(kj::mv(queryStatement), rowsRead, rowsWritten, queryLatency, dbWalBytesWritten, queryResult, extendedErrorCode, isInternalQuery, - kj::mv(queryErrorDescription)); + kj::mv(queryErrorDescription), queryId); } void setQueryEventStats(uint64_t rowsRead, uint64_t rowsWritten, bool isInternalQuery) { @@ -700,6 +713,10 @@ class SqliteDatabase::Query final: private ResetListener { extendedErrorCode = res; } + void setQueryId(kj::Maybe queryId) { + this->queryId = queryId; + } + private: SqliteObserver& observer; kj::Maybe queryStatement = kj::none; @@ -710,6 +727,7 @@ class SqliteDatabase::Query final: private ResetListener { uint64_t rowsWritten = 0; int queryResult = 0; int extendedErrorCode = 0; + kj::Maybe queryId = kj::none; kj::Maybe queryErrorDescription = kj::none; }; @@ -742,25 +760,27 @@ class SqliteDatabase::Query final: private ResetListener { Query(SqliteDatabase& db, QueryOptions options, Statement& statement, Params&&... bindings) : ResetListener(db), regulator(options.regulator), - maybeStatement(statement.prepareForExecution()), + maybeStatement(statement.prepareForExecution(options.queryId)), queryEvent(this->db.sqliteObserver), allowUnconfirmed(options.allowUnconfirmed) { // If we throw from the constructor, the destructor won't run. Need to call destroy() // explicitly. KJ_ON_SCOPE_FAILURE(destroy()); + queryEvent.setQueryId(options.queryId); bindAll(std::index_sequence_for(), kj::fwd(bindings)...); } template Query(SqliteDatabase& db, QueryOptions options, kj::StringPtr sqlCode, Params&&... bindings) : ResetListener(db), regulator(options.regulator), - ownStatement(db.prepareSql(regulator, sqlCode, 0, MULTI)), + ownStatement(db.prepareSql(regulator, sqlCode, 0, MULTI, kj::none, options.queryId)), maybeStatement(ownStatement), queryEvent(this->db.sqliteObserver), allowUnconfirmed(options.allowUnconfirmed) { // If we throw from the constructor, the destructor won't run. Need to call destroy() // explicitly. KJ_ON_SCOPE_FAILURE(destroy()); + queryEvent.setQueryId(options.queryId); bindAll(std::index_sequence_for(), kj::fwd(bindings)...); } @@ -1070,8 +1090,11 @@ SqliteDatabase::Query SqliteDatabase::Statement::run(Params&&... params) { template SqliteDatabase::Query SqliteDatabase::Statement::run(StatementOptions options, Params&&... params) { - return Query(db, {.regulator = regulator, .allowUnconfirmed = options.allowUnconfirmed}, *this, - kj::fwd(params)...); + return Query(db, + {.regulator = regulator, + .allowUnconfirmed = options.allowUnconfirmed, + .queryId = options.queryId}, + *this, kj::fwd(params)...); } template