Fix connection pool causing unhandled exceptions

simolus3 · Feb 27, 2023 · b137b40 · b137b40
1 parent 46390a2
commit b137b40
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 27 deletions.
diff --git a/drift/lib/src/runtime/executor/connection_pool.dart b/drift/lib/src/runtime/executor/connection_pool.dart
@@ -29,23 +29,13 @@ abstract class MultiExecutor extends QueryExecutor {
   MultiExecutor._();
 }
 
-class _ExecutorCompleter {
-  _ExecutorCompleter(this.statement, this.args)
-      : _completer = Completer<List<Map<String, Object?>>>();
+class _PendingSelect {
+  _PendingSelect(this.statement, this.args)
+      : completer = Completer<List<Map<String, Object?>>>();
 
   final String statement;
   final List<Object?> args;
-  final Completer<List<Map<String, Object?>>> _completer;
-
-  Future<List<Map<String, Object?>>> get future => _completer.future;
-
-  void complete([FutureOr<List<Map<String, Object?>>>? value]) {
-    _completer.complete(value);
-  }
-
-  void completeError(Object error, [StackTrace? stackTrace]) {
-    _completer.completeError(error, stackTrace);
-  }
+  final Completer<List<Map<String, Object?>>> completer;
 }
 
 class _QueryExecutorPool {
@@ -54,8 +44,8 @@ class _QueryExecutorPool {
   final List<QueryExecutor> _executors;
   final List<QueryExecutor> _idleExecutors;
 
-  final List<_ExecutorCompleter> _queue = [];
-  final List<_ExecutorCompleter> _running = [];
+  final List<_PendingSelect> _queue = [];
+  final List<_PendingSelect> _running = [];
 
   Future<bool> ensureOpen(QueryExecutorUser user) async {
     final result = await Future.wait(
@@ -72,10 +62,10 @@ class _QueryExecutorPool {
       return _executors.single.runSelect(statement, args);
     }
 
-    final executorCompleter = _ExecutorCompleter(statement, args);
+    final executorCompleter = _PendingSelect(statement, args);
     _queue.add(executorCompleter);
     _run();
-    return executorCompleter.future;
+    return executorCompleter.completer.future;
   }
 
   void _run() {
@@ -87,15 +77,15 @@ class _QueryExecutorPool {
 
     _running.add(completer);
 
-    completer.future.whenComplete(() {
-      _running.remove(completer);
-      _idleExecutors.add(executor);
-      _run();
-    });
-
-    executor
-        .runSelect(completer.statement, completer.args)
-        .then(completer.complete, onError: completer.completeError);
+    completer.completer.complete(Future.sync(() async {
+      try {
+        return await executor.runSelect(completer.statement, completer.args);
+      } finally {
+        _running.remove(completer);
+        _idleExecutors.add(executor);
+        _run();
+      }
+    }));
   }
 }
 

diff --git a/drift/test/engines/connection_pool_test.dart b/drift/test/engines/connection_pool_test.dart
@@ -115,4 +115,17 @@ void main() {
     verify(write.transactions.ensureOpen(any));
     verify(write.transactions.runSelect('select', []));
   });
+
+  test('select failure does not cause an unhandled exception', () async {
+    // https://github.com/simolus3/drift/issues/2323
+    final read2 = MockExecutor();
+    final multi =
+        MultiExecutor.withReadPool(reads: [read2, read], write: write);
+
+    when(read2.runSelect(any, any)).thenThrow('bang!');
+
+    await multi.ensureOpen(db);
+
+    expect(multi.runSelect('select 1', []), throwsA(isA<String>()));
+  });
 }