From 39645f669c44d5bb68b4ec5e0d4ca8b784363619 Mon Sep 17 00:00:00 2001
From: Simon Binder <oss@simonbinder.eu>
Date: Tue, 11 Apr 2023 11:33:12 +0200
Subject: [PATCH] Allow retrying if database setup fails

---
 drift/CHANGELOG.md                            |  5 +++
 drift/lib/src/sqlite3/database.dart           | 44 +++++++++++--------
 .../platforms/ffi/native_database_test.dart   | 30 +++++++++++++
 3 files changed, 61 insertions(+), 18 deletions(-)

diff --git a/drift/CHANGELOG.md b/drift/CHANGELOG.md
index b0856783..22a437af 100644
--- a/drift/CHANGELOG.md
+++ b/drift/CHANGELOG.md
@@ -1,3 +1,8 @@
+## 2.8.0-dev
+
+- Don't keep databases in an unusable state if the `setup` callback throws an
+  exception. Instead, drift will retry the next time the database is used.
+
 ## 2.7.0
 
 - Add support for `CASE` expressions without a base in the Dart API with the
diff --git a/drift/lib/src/sqlite3/database.dart b/drift/lib/src/sqlite3/database.dart
index dc6f567d..863633e7 100644
--- a/drift/lib/src/sqlite3/database.dart
+++ b/drift/lib/src/sqlite3/database.dart
@@ -14,10 +14,12 @@ import 'native_functions.dart';
 /// through `package:js`.
 abstract class Sqlite3Delegate<DB extends CommonDatabase>
     extends DatabaseDelegate {
-  /// The underlying database instance from the `sqlite3` package.
-  late DB database;
+  DB? _database;
 
-  bool _hasCreatedDatabase = false;
+  /// The underlying database instance from the `sqlite3` package.
+  DB get database => _database!;
+
+  bool _hasInitializedDatabase = false;
   bool _isOpen = false;
 
   final void Function(DB)? _setup;
@@ -35,8 +37,7 @@ abstract class Sqlite3Delegate<DB extends CommonDatabase>
   /// A delegate using an underlying sqlite3 database object that has already
   /// been opened.
   Sqlite3Delegate.opened(
-      this.database, this._setup, this.closeUnderlyingWhenClosed)
-      : _hasCreatedDatabase = true {
+      this._database, this._setup, this.closeUnderlyingWhenClosed) {
     _initializeDatabase();
   }
 
@@ -55,26 +56,33 @@ abstract class Sqlite3Delegate<DB extends CommonDatabase>
 
   @override
   Future<void> open(QueryExecutorUser db) async {
-    if (!_hasCreatedDatabase) {
-      _createDatabase();
-      _initializeDatabase();
+    if (!_hasInitializedDatabase) {
+      assert(_database == null);
+      _database = openDatabase();
+
+      try {
+        _initializeDatabase();
+      } catch (e) {
+        // If the initialization fails, we effectively don't have a usable
+        // database, so reset
+        _database?.dispose();
+        _database = null;
+
+        rethrow;
+      }
     }
 
     _isOpen = true;
     return Future.value();
   }
 
-  void _createDatabase() {
-    assert(!_hasCreatedDatabase);
-    _hasCreatedDatabase = true;
-
-    database = openDatabase();
-  }
-
   void _initializeDatabase() {
+    assert(!_hasInitializedDatabase);
+
     database.useNativeFunctions();
     _setup?.call(database);
-    versionDelegate = _VmVersionDelegate(database);
+    versionDelegate = _SqliteVersionDelegate(database);
+    _hasInitializedDatabase = true;
   }
 
   /// Synchronously prepares and runs [statements] collected from a batch.
@@ -127,10 +135,10 @@ abstract class Sqlite3Delegate<DB extends CommonDatabase>
   }
 }
 
-class _VmVersionDelegate extends DynamicVersionDelegate {
+class _SqliteVersionDelegate extends DynamicVersionDelegate {
   final CommonDatabase database;
 
-  _VmVersionDelegate(this.database);
+  _SqliteVersionDelegate(this.database);
 
   @override
   Future<int> get schemaVersion => Future.value(database.userVersion);
diff --git a/drift/test/platforms/ffi/native_database_test.dart b/drift/test/platforms/ffi/native_database_test.dart
index e8153a24..3f15c897 100644
--- a/drift/test/platforms/ffi/native_database_test.dart
+++ b/drift/test/platforms/ffi/native_database_test.dart
@@ -90,6 +90,36 @@ void main() {
       );
     });
   });
+
+  test('calls setup twice if first invocation fails', () async {
+    const exception = 'exception';
+    var count = 0;
+    final db = NativeDatabase.memory(
+      setup: expectAsync1(
+        (_) {
+          if (count++ == 0) {
+            throw exception;
+          }
+        },
+        count: 2,
+      ),
+    );
+
+    await expectLater(db.ensureOpen(_FakeExecutorUser()), throwsA(exception));
+
+    // Should also prevent subsequent open attempts
+    await expectLater(db.ensureOpen(_FakeExecutorUser()), completes);
+  });
+
+  test('throwing in setup prevents the database from being opened', () async {
+    const exception = 'exception';
+    final db = NativeDatabase.memory(setup: (_) => throw exception);
+
+    await expectLater(db.ensureOpen(_FakeExecutorUser()), throwsA(exception));
+
+    // Should also prevent subsequent open attempts
+    await expectLater(db.ensureOpen(_FakeExecutorUser()), throwsA(exception));
+  });
 }
 
 class _FakeExecutorUser extends QueryExecutorUser {