Add a callback called after a SQLite database is reset().

This can be used to re-enable WAL mode, and other things that need to happen at startup.

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<ActorSqlite>(kj::mv(db), outputGate,
                     []() -> kj::Promise<void> { return kj::READY_NOW; }, *sqliteHooks)

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,

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<void()> 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<void(SqliteDatabase&)> callback) {
+    afterResetCallback = kj::mv(callback);
+  }
+
 private:
   const Vfs& vfs;
   kj::Path path;
@@ -206,6 +216,7 @@ class SqliteDatabase {
   kj::Maybe<sqlite3_stmt&> currentStatement;
 
   kj::Maybe<kj::Function<void()>> onWriteCallback;
+  kj::Maybe<kj::Function<void(SqliteDatabase&)>> afterResetCallback;
 
   kj::List<ResetListener, &ResetListener::link> resetListeners;