From 37f120d287057c2cd2ad4b79f02a7584957b9f2e Mon Sep 17 00:00:00 2001
From: Simon Binder <oss@simonbinder.eu>
Date: Fri, 5 Apr 2024 22:39:05 +0200
Subject: [PATCH] Add setup param to schema verifier

---
 drift_dev/CHANGELOG.md                            |  3 +++
 drift_dev/lib/api/migrations.dart                 | 14 ++++++++++++--
 .../lib/src/services/schema/verifier_impl.dart    | 15 +++++++++++----
 .../test/services/schema/verifier_impl_test.dart  | 12 +++++++++++-
 4 files changed, 37 insertions(+), 7 deletions(-)

diff --git a/drift_dev/CHANGELOG.md b/drift_dev/CHANGELOG.md
index bbcfdcc0..8d6b0e3b 100644
--- a/drift_dev/CHANGELOG.md
+++ b/drift_dev/CHANGELOG.md
@@ -2,6 +2,9 @@
 
 - Fix drift using the wrong import alias in generated part files.
 - Add the `use_sql_column_name_as_json_key` builder option.
+- Add a `setup` parameter to `SchemaVerifier`. It is called when the verifier
+  creates database connections (similar to the callback on `NativeDatabase`)
+  and can be used to register custom functions.
 
 ## 2.16.0
 
diff --git a/drift_dev/lib/api/migrations.dart b/drift_dev/lib/api/migrations.dart
index a09cdf15..73047c8b 100644
--- a/drift_dev/lib/api/migrations.dart
+++ b/drift_dev/lib/api/migrations.dart
@@ -11,8 +11,18 @@ export 'package:drift_dev/src/services/schema/verifier_common.dart'
     show SchemaMismatch;
 
 abstract class SchemaVerifier {
-  factory SchemaVerifier(SchemaInstantiationHelper helper) =
-      VerifierImplementation;
+  /// Creates a schema verifier for the drift-generated [helper].
+  ///
+  /// See [tests] for more information.
+  /// The optional [setup] parameter is used internally by the verifier for
+  /// every database connection it opens. This can be used to, for instance,
+  /// register custom functions expected by your database.
+  ///
+  /// [tests]: https://drift.simonbinder.eu/docs/migrations/tests/
+  factory SchemaVerifier(
+    SchemaInstantiationHelper helper, {
+    void Function(Database raw)? setup,
+  }) = VerifierImplementation;
 
   /// Creates a [DatabaseConnection] that contains empty tables created for the
   /// known schema [version].
diff --git a/drift_dev/lib/src/services/schema/verifier_impl.dart b/drift_dev/lib/src/services/schema/verifier_impl.dart
index ba8b4941..97e54a38 100644
--- a/drift_dev/lib/src/services/schema/verifier_impl.dart
+++ b/drift_dev/lib/src/services/schema/verifier_impl.dart
@@ -13,8 +13,9 @@ Expando<List<Input>> expectedSchema = Expando();
 class VerifierImplementation implements SchemaVerifier {
   final SchemaInstantiationHelper helper;
   final Random _random = Random();
+  final void Function(Database)? setup;
 
-  VerifierImplementation(this.helper);
+  VerifierImplementation(this.helper, {this.setup});
 
   @override
   Future<void> migrateAndValidate(GeneratedDatabase db, int expectedVersion,
@@ -57,14 +58,20 @@ class VerifierImplementation implements SchemaVerifier {
     return buffer.toString();
   }
 
+  Database _setupDatabase(String uri) {
+    final database = sqlite3.open(uri, uri: true);
+    setup?.call(database);
+    return database;
+  }
+
   @override
   Future<InitializedSchema> schemaAt(int version) async {
     // Use distinct executors for setup and use, allowing us to close the helper
     // db here and avoid creating it twice.
     // https://www.sqlite.org/inmemorydb.html#sharedmemdb
     final uri = 'file:mem${_randomString()}?mode=memory&cache=shared';
-    final dbForSetup = sqlite3.open(uri, uri: true);
-    final dbForUse = sqlite3.open(uri, uri: true);
+    final dbForSetup = _setupDatabase(uri);
+    final dbForUse = _setupDatabase(uri);
 
     final executor = NativeDatabase.opened(dbForSetup);
     final db = helper.databaseForVersion(executor, version);
@@ -74,7 +81,7 @@ class VerifierImplementation implements SchemaVerifier {
     await db.close();
 
     return InitializedSchema(dbForUse, () {
-      final db = sqlite3.open(uri, uri: true);
+      final db = _setupDatabase(uri);
       return DatabaseConnection(NativeDatabase.opened(db));
     });
   }
diff --git a/drift_dev/test/services/schema/verifier_impl_test.dart b/drift_dev/test/services/schema/verifier_impl_test.dart
index b9536fcf..e65862f4 100644
--- a/drift_dev/test/services/schema/verifier_impl_test.dart
+++ b/drift_dev/test/services/schema/verifier_impl_test.dart
@@ -5,7 +5,11 @@ import 'package:drift_dev/src/services/schema/verifier_impl.dart';
 import 'package:test/test.dart';
 
 void main() {
-  final verifier = SchemaVerifier(_TestHelper());
+  final verifier = SchemaVerifier(
+    _TestHelper(),
+    setup: (rawDb) => rawDb.createFunction(
+        functionName: 'test_function', function: (args) => 1),
+  );
 
   group('startAt', () {
     test('starts at the requested version', () async {
@@ -15,6 +19,12 @@ void main() {
         expect(details.hadUpgrade, isFalse, reason: 'no upgrade expected');
       }));
     });
+
+    test('registers custom functions', () async {
+      final db = (await verifier.startAt(17)).executor;
+      await db.ensureOpen(_DelegatedUser(17, (_, details) async {}));
+      await db.runSelect('select test_function()', []);
+    });
   });
 
   group('migrateAndValidate', () {