Add wasm setup APIs with custom workers (#2638)

docs/lib/snippets/platforms/web.dart

@@ -118,3 +118,30 @@ DatabaseConnection migrateFromLegacy() {
     return result.resolvedExecutor;
   }));
 }
+
+// #docregion setupAll
+void setupDatabase(CommonDatabase database) {
+  database.createFunction(
+    functionName: 'my_function',
+    function: (args) => args.length,
+  );
+}
+
+void main() {
+  WasmDatabase.workerMainForOpen(
+    setupAllDatabases: setupDatabase,
+  );
+}
+// #enddocregion setupAll
+
+void withSetup() async {
+  // #docregion setupLocal
+  final result = await WasmDatabase.open(
+    databaseName: 'my_app_db', // prefer to only use valid identifiers here
+    sqlite3Uri: Uri.parse('sqlite3.wasm'),
+    driftWorkerUri: Uri.parse('my_drift_worker.dart.js'),
+    localSetup: setupDatabase,
+  );
+  // #enddocregion setupLocal
+  print(result);
+}

docs/pages/docs/Platforms/web.md

@@ -245,6 +245,31 @@ and drift has chosen the `unsafeIndexedDb` or the `inMemory` implementation due
 persistence support, you may want to show a warning to the user explaining that they have to upgrade
 their browser.
 
+### Custom functions and database setup
+
+Constructors on `NativeDatabase` have a `setup` callback used to initialize the raw database instance,
+for instance by registering custom functions that are needed on the database.
+Unfortunately, this cannot be supported by `WasmDatabase.open` directly. Since the raw database instance
+may only exist on a web worker, we can't use Dart callback functions.
+
+However, it is possible to compile a custom worker that will setup drift databases it creates. For that,
+you can create a Dart file (usually put into `web/`) with content like:
+
+{% include "blocks/snippet" snippets = snippets name = "setupAll" %}
+
+To open the database from the application, you can then use this:
+
+{% include "blocks/snippet" snippets = snippets name = "setupLocal" %}
+
+The `setupDatabase` value is duplicated with `localSetup` in case drift determines that it needs to use
+an in-memory database or an IndexedDB-powered database that doesn't run in a worker. In that case,
+`localSetup` would get called instead. For databases running in a worker, the worker calls `setupAllDatabases`
+which creates the necessary functions.
+
+The next section explains how to compile the Dart file calling `workerMainForOpen` into a JavaScript
+worker that can be referenced with `driftWorkerUri`.
+There is no need to compile a custom `sqlite3.wasm` for this.
+
 ### Compilation
 
 Drift and the `sqlite3` Dart package provide pre-compiled versions of the worker and the WebAssembly

drift/CHANGELOG.md

@@ -1,5 +1,6 @@
 ## 2.13.0-dev
 
+- Add APIs to setup Wasm databases with custom drift workers.
 - Add `Expression.and` and `Expression.or` to create disjunctions and conjunctions
   of sub-predicates.
 

drift/lib/src/web/wasm_setup.dart

@@ -220,6 +220,7 @@ final class _ProbeResult implements WasmProbeResult {
     WasmStorageImplementation implementation,
     String name, {
     FutureOr<Uint8List?> Function()? initializeDatabase,
+    WasmDatabaseSetup? localSetup,
   }) async {
     final channel = MessageChannel();
     final initializer = initializeDatabase;
@@ -249,14 +250,17 @@ final class _ProbeResult implements WasmProbeResult {
         } else {
           // Workers seem to be broken, but we don't need them with this storage
           // mode.
-          return _hostDatabaseLocally(implementation,
-              await IndexedDbFileSystem.open(dbName: name), initializeDatabase);
+          return _hostDatabaseLocally(
+              implementation,
+              await IndexedDbFileSystem.open(dbName: name),
+              initializeDatabase,
+              localSetup);
         }
       case WasmStorageImplementation.inMemory:
         // Nothing works on this browser, so we'll fall back to an in-memory
         // database.
-        return _hostDatabaseLocally(
-            implementation, InMemoryFileSystem(), initializeDatabase);
+        return _hostDatabaseLocally(implementation, InMemoryFileSystem(),
+            initializeDatabase, localSetup);
     }
 
     initChannel?.port1.onMessage.listen((event) async {
@@ -298,24 +302,16 @@ final class _ProbeResult implements WasmProbeResult {
     WasmStorageImplementation storage,
     VirtualFileSystem vfs,
     FutureOr<Uint8List?> Function()? initializer,
+    WasmDatabaseSetup? setup,
   ) async {
-    final sqlite3 = await WasmSqlite3.loadFromUrl(opener.sqlite3WasmUri);
-    sqlite3.registerVirtualFileSystem(vfs);
-
-    if (initializer != null) {
-      final blob = await initializer();
-      if (blob != null) {
-        final (file: file, outFlags: _) =
-            vfs.xOpen(Sqlite3Filename('/database'), SqlFlag.SQLITE_OPEN_CREATE);
-        file
-          ..xWrite(blob, 0)
-          ..xClose();
-      }
-    }
-
-    return DatabaseConnection(
-      WasmDatabase(sqlite3: sqlite3, path: '/database'),
+    final database = await DriftServerController(setup).openConnection(
+      sqlite3WasmUri: opener.sqlite3WasmUri,
+      databaseName: 'database',
+      storage: storage,
+      initializer: initializer,
     );
+
+    return DatabaseConnection(database);
   }
 
   @override

drift/lib/src/web/wasm_setup/dedicated_worker.dart

@@ -3,22 +3,23 @@
 import 'dart:async';
 import 'dart:html';
 
-import 'package:drift/wasm.dart';
 import 'package:js/js_util.dart';
 import 'package:sqlite3/wasm.dart';
 
 import '../../utils/synchronized.dart';
 import 'protocol.dart';
 import 'shared.dart';
+import 'types.dart';
 
 class DedicatedDriftWorker {
   final DedicatedWorkerGlobalScope self;
   final Lock _checkCompatibility = Lock();
 
-  final DriftServerController _servers = DriftServerController();
+  final DriftServerController _servers;
   WasmCompatibility? _compatibility;
 
-  DedicatedDriftWorker(this.self);
+  DedicatedDriftWorker(this.self, WasmDatabaseSetup? setup)
+      : _servers = DriftServerController(setup);
 
   void start() {
     self.onMessage.listen((event) {

drift/lib/src/web/wasm_setup/shared.dart

@@ -1,3 +1,4 @@
+import 'dart:async';
 import 'dart:html';
 import 'dart:indexed_db';
 
@@ -165,47 +166,35 @@ Future<void> deleteDatabaseInOpfs(String databaseName) async {
 class DriftServerController {
   /// Running drift servers by the name of the database they're serving.
   final Map<String, RunningWasmServer> servers = {};
+  final WasmDatabaseSetup? _setup;
+
+  /// Creates a controller responsible for loading wasm databases and serving
+  /// them. The [_setup] callback will be invoked on created databases if set.
+  DriftServerController(this._setup);
 
   /// Serves a drift connection as requested by the [message].
   void serve(
     ServeDriftDatabase message,
   ) {
     final server = servers.putIfAbsent(message.databaseName, () {
-      final server = DriftServer(LazyDatabase(() async {
-        final sqlite3 = await WasmSqlite3.loadFromUrl(message.sqlite3WasmUri);
+      final initPort = message.initializationPort;
 
-        final vfs = await switch (message.storage) {
-          WasmStorageImplementation.opfsShared =>
-            SimpleOpfsFileSystem.loadFromStorage(
-                pathForOpfs(message.databaseName)),
-          WasmStorageImplementation.opfsLocks =>
-            _loadLockedWasmVfs(message.databaseName),
-          WasmStorageImplementation.unsafeIndexedDb ||
-          WasmStorageImplementation.sharedIndexedDb =>
-            IndexedDbFileSystem.open(dbName: message.databaseName),
-          WasmStorageImplementation.inMemory =>
-            Future.value(InMemoryFileSystem()),
-        };
+      final initializer = initPort != null
+          ? () async {
+              initPort.postMessage(true);
 
-        final initPort = message.initializationPort;
-        if (vfs.xAccess('/database', 0) == 0 && initPort != null) {
-          initPort.postMessage(true);
+              return await initPort.onMessage
+                  .map((e) => e.data as Uint8List?)
+                  .first;
+            }
+          : null;
 
-          final response =
-              await initPort.onMessage.map((e) => e.data as Uint8List?).first;
-
-          if (response != null) {
-            final (file: file, outFlags: _) = vfs.xOpen(
-                Sqlite3Filename('/database'), SqlFlag.SQLITE_OPEN_CREATE);
-            file.xWrite(response, 0);
-            file.xClose();
-          }
-        }
-
-        sqlite3.registerVirtualFileSystem(vfs, makeDefault: true);
-
-        return WasmDatabase(sqlite3: sqlite3, path: '/database');
-      }));
+      final server = DriftServer(LazyDatabase(() => openConnection(
+            sqlite3WasmUri: message.sqlite3WasmUri,
+            databaseName: message.databaseName,
+            storage: message.storage,
+            initializer: initializer,
+          )));
 
       return RunningWasmServer(message.storage, server);
     });
@@ -213,6 +202,41 @@ class DriftServerController {
     server.server.serve(message.port.channel());
   }
 
+  /// Loads a new sqlite3 WASM module, registers an appropriate VFS for [storage]
+  /// and finally opens a database, creating it if it doesn't exist.
+  Future<WasmDatabase> openConnection({
+    required Uri sqlite3WasmUri,
+    required String databaseName,
+    required WasmStorageImplementation storage,
+    required FutureOr<Uint8List?> Function()? initializer,
+  }) async {
+    final sqlite3 = await WasmSqlite3.loadFromUrl(sqlite3WasmUri);
+
+    final vfs = await switch (storage) {
+      WasmStorageImplementation.opfsShared =>
+        SimpleOpfsFileSystem.loadFromStorage(pathForOpfs(databaseName)),
+      WasmStorageImplementation.opfsLocks => _loadLockedWasmVfs(databaseName),
+      WasmStorageImplementation.unsafeIndexedDb ||
+      WasmStorageImplementation.sharedIndexedDb =>
+        IndexedDbFileSystem.open(dbName: databaseName),
+      WasmStorageImplementation.inMemory => Future.value(InMemoryFileSystem()),
+    };
+
+    if (initializer != null && vfs.xAccess('/database', 0) == 0) {
+      final response = await initializer();
+
+      if (response != null) {
+        final (file: file, outFlags: _) =
+            vfs.xOpen(Sqlite3Filename('/database'), SqlFlag.SQLITE_OPEN_CREATE);
+        file.xWrite(response, 0);
+        file.xClose();
+      }
+    }
+
+    sqlite3.registerVirtualFileSystem(vfs, makeDefault: true);
+    return WasmDatabase(sqlite3: sqlite3, path: '/database', setup: _setup);
+  }
+
   Future<WasmVfs> _loadLockedWasmVfs(String databaseName) async {
     // Create SharedArrayBuffers to synchronize requests
     final options = WasmVfs.createOptions(

drift/lib/src/web/wasm_setup/shared_worker.dart

@@ -16,9 +16,10 @@ class SharedDriftWorker {
   /// "shared-dedicated" worker hosting the database.
   Worker? _dedicatedWorker;
 
-  final DriftServerController _servers = DriftServerController();
+  final DriftServerController _servers;
 
-  SharedDriftWorker(this.self);
+  SharedDriftWorker(this.self, WasmDatabaseSetup? setup)
+      : _servers = DriftServerController(setup);
 
   void start() {
     const event = EventStreamProvider<MessageEvent>('connect');

drift/lib/src/web/wasm_setup/types.dart

@@ -7,6 +7,14 @@ library;
 import 'dart:async';
 
 import 'package:drift/drift.dart';
+import 'package:sqlite3/common.dart';
+
+/// Signature of a function that can perform setup work on a [database] before
+/// drift is fully ready.
+///
+/// This could be used to, for instance, register custom user-defined functions
+/// on the database.
+typedef WasmDatabaseSetup = void Function(CommonDatabase database);
 
 /// The storage implementation used by the `drift` and `sqlite3` packages to
 /// emulate a synchronous file system on the web, used by the sqlite3 C library
@@ -176,6 +184,7 @@ abstract interface class WasmProbeResult {
     WasmStorageImplementation implementation,
     String name, {
     FutureOr<Uint8List?> Function()? initializeDatabase,
+    WasmDatabaseSetup? localSetup,
   });
 
   /// Deletes an [ExistingDatabase] from storage.

drift/lib/wasm.dart

@@ -22,13 +22,6 @@ import 'src/web/wasm_setup/types.dart';
 
 export 'src/web/wasm_setup/types.dart';
 
-/// Signature of a function that can perform setup work on a [database] before
-/// drift is fully ready.
-///
-/// This could be used to, for instance, set encryption keys for SQLCipher
-/// implementations.
-typedef WasmDatabaseSetup = void Function(CommonDatabase database);
-
 /// An experimental, WebAssembly based implementation of a drift sqlite3
 /// database.
 ///
@@ -98,12 +91,21 @@ class WasmDatabase extends DelegatedDatabase {
   /// which you can [get here](https://github.com/simolus3/sqlite3.dart/releases),
   /// and a drift worker, which you can [get here](https://drift.simonbinder.eu/web/#worker).
   ///
+  /// [localSetup] will be called to initialize the database only if the
+  /// database will be opened directly in this JavaScript context. It is likely
+  /// that the database will actually be opened in a web worker, with drift
+  /// using communication mechanisms to access the database. As there is no way
+  /// to send the database over to the main context, [localSetup] would not be
+  /// called in that case. Instead, you'd have to compile a custom drift worker
+  /// with a setup function - see [workerMainForOpen] for additional information.
+  ///
   /// For more detailed information, see https://drift.simonbinder.eu/web.
   static Future<WasmDatabaseResult> open({
     required String databaseName,
     required Uri sqlite3Uri,
     required Uri driftWorkerUri,
     FutureOr<Uint8List?> Function()? initializeDatabase,
+    WasmDatabaseSetup? localSetup,
   }) async {
     final probed = await probe(
       sqlite3Uri: sqlite3Uri,
@@ -147,7 +149,8 @@ class WasmDatabase extends DelegatedDatabase {
 
     final bestImplementation = availableImplementations.firstOrNull ??
         WasmStorageImplementation.inMemory;
-    final connection = await probed.open(bestImplementation, databaseName);
+    final connection = await probed.open(bestImplementation, databaseName,
+        localSetup: localSetup);
 
     return WasmDatabaseResult(
         connection, bestImplementation, probed.missingFeatures);
@@ -196,13 +199,18 @@ class WasmDatabase extends DelegatedDatabase {
   /// If you prefer to compile the worker yourself, write a simple Dart program
   /// that calls this method in its `main()` function and compile that with
   /// `dart2js`.
-  static void workerMainForOpen() {
+  /// This is particularly useful when using [setupAllDatabases], a callback
+  /// that will be invoked on every new [CommonDatabase] created by the web
+  /// worker. This is a suitable place to register custom functions.
+  static void workerMainForOpen({
+    WasmDatabaseSetup? setupAllDatabases,
+  }) {
     final self = WorkerGlobalScope.instance;
 
     if (self is DedicatedWorkerGlobalScope) {
-      DedicatedDriftWorker(self).start();
+      DedicatedDriftWorker(self, setupAllDatabases).start();
     } else if (self is SharedWorkerGlobalScope) {
-      SharedDriftWorker(self).start();
+      SharedDriftWorker(self, setupAllDatabases).start();
     }
   }
 }

drift_dev/pubspec.yaml

@@ -30,7 +30,7 @@ dependencies:
   io: ^1.0.3
 
   # Drift-specific analysis and apis
-  drift: '>=2.12.0 <2.13.0'
+  drift: '>=2.13.0 <2.14.0'
   sqlite3: '>=0.1.6 <3.0.0'
   sqlparser: '^0.31.2'
 

extras/integration_tests/web_wasm/test/drift_wasm_test.dart

@@ -142,7 +142,7 @@ void main() {
           }
 
           group(
-            'initialization from ',
+            'initialization from',
             () {
               test('static blob', () async {
                 await driver.enableInitialization(InitializationMode.loadAsset);

extras/integration_tests/web_wasm/web/main.dart

@@ -141,6 +141,15 @@ Future<void> _open(String? implementationName) async {
       sqlite3Uri: sqlite3WasmUri,
       driftWorkerUri: driftWorkerUri,
       initializeDatabase: _initializeDatabase,
+      localSetup: (db) {
+        // The worker has a similar setup call that will make database_host
+        // return `worker` instead.
+        db.createFunction(
+          functionName: 'database_host',
+          function: (args) => 'document',
+          argumentCount: const AllowedArgumentCount(1),
+        );
+      },
     );
 
     connection = result.resolvedExecutor;
@@ -149,7 +158,7 @@ Future<void> _open(String? implementationName) async {
   final db = openedDatabase = TestDatabase(connection);
 
   // Make sure it works!
-  await db.customSelect('SELECT 1').get();
+  await db.customSelect('SELECT database_host()').get();
 
   tableUpdates = StreamQueue(db.testTable.all().watch());
   await tableUpdates!.next;

extras/integration_tests/web_wasm/web/worker.dart

@@ -1,5 +1,12 @@
 import 'package:drift/wasm.dart';
+import 'package:sqlite3/wasm.dart';
 
 void main() {
-  WasmDatabase.workerMainForOpen();
+  WasmDatabase.workerMainForOpen(setupAllDatabases: (db) {
+    db.createFunction(
+      functionName: 'database_host',
+      function: (args) => 'worker',
+      argumentCount: const AllowedArgumentCount(1),
+    );
+  });
 }