diff --git a/src/workerd/server/server.c++ b/src/workerd/server/server.c++ index 7ba3e0e380ba..dbc10a29f9f5 100644 --- a/src/workerd/server/server.c++ +++ b/src/workerd/server/server.c++ @@ -1941,8 +1941,13 @@ public: kj::Path({d.uniqueKey, kj::str(idPtr, ".sqlite")}), kj::WriteMode::CREATE | kj::WriteMode::MODIFY | kj::WriteMode::CREATE_PARENT); - // Before we do anything, make sure the database is in WAL mode. - db->run("PRAGMA journal_mode=WAL;"); + // Before we do anything, make sure the database is in WAL mode. We also need to + // do this after reset() is used, so register a callback for that. + auto setWalMode = [](SqliteDatabase& db) { + db.run("PRAGMA journal_mode=WAL;"); + }; + setWalMode(*db); + db->afterReset(kj::mv(setWalMode)); return kj::heap(kj::mv(db), outputGate, []() -> kj::Promise { return kj::READY_NOW; }, *sqliteHooks) diff --git a/src/workerd/util/sqlite.c++ b/src/workerd/util/sqlite.c++ index a95d7f19da6b..c3cafba2fb29 100644 --- a/src/workerd/util/sqlite.c++ +++ b/src/workerd/util/sqlite.c++ @@ -564,6 +564,10 @@ void SqliteDatabase::executeWithRegulator( } void SqliteDatabase::reset() { + // Temporarily disable the on-write callback while resetting. + auto writeCb = kj::mv(onWriteCallback); + KJ_DEFER(onWriteCallback = kj::mv(writeCb)); + KJ_IF_SOME(db, maybeDb) { for (auto& listener: resetListeners) { listener.beforeSqliteReset(); @@ -579,6 +583,10 @@ void SqliteDatabase::reset() { KJ_ON_SCOPE_FAILURE(maybeDb = kj::none); init(); + + KJ_IF_SOME(resetCb, afterResetCallback) { + resetCb(*this); + } } bool SqliteDatabase::isAuthorized(int actionCode, diff --git a/src/workerd/util/sqlite.h b/src/workerd/util/sqlite.h index c399365fa176..1998cb57f8e2 100644 --- a/src/workerd/util/sqlite.h +++ b/src/workerd/util/sqlite.h @@ -130,8 +130,8 @@ class SqliteDatabase { // // Durable Objects uses this to automatically begin a transaction and close the output gate. // - // Note that the write callback is NOT called before a reset(). Use the `ResetListener` mechanism - // instead for that case. + // Note that the write callback is NOT called before (or at any point during) a reset(). Use the + // `ResetListener` mechanism or `afterReset()` instead for that case. void onWrite(kj::Function callback) { onWriteCallback = kj::mv(callback); } @@ -191,6 +191,16 @@ class SqliteDatabase { friend class SqliteDatabase; }; + // Registers a callback to call after a reset completes. This can be used to do basic database + // initialization, e.g. set WAL mode. (To get notified *before* a reset, use `ResetListener`.) + // + // Note that the on-write callback is disabled during reset(), including while calling the + // after-reset callback. So, queries performed by the after-reset callback will not trigger the + // on-write callback. + void afterReset(kj::Function callback) { + afterResetCallback = kj::mv(callback); + } + private: const Vfs& vfs; kj::Path path; @@ -206,6 +216,7 @@ class SqliteDatabase { kj::Maybe currentStatement; kj::Maybe> onWriteCallback; + kj::Maybe> afterResetCallback; kj::List resetListeners;