Update

drift/lib/src/sentry_drift_database.dart

@@ -9,18 +9,21 @@ import 'sentry_span_helper.dart';
 /// Signature of a function that opens a database connection when instructed to.
 typedef DatabaseOpener = FutureOr<QueryExecutor> Function();
 
-/// Sentry Drift Database Executor based on [LazyDatabase].
+/// The Sentry Drift Database Executor.
 ///
-/// A special database executor that delegates work to another [QueryExecutor].
-/// The other executor is lazily opened by a [DatabaseOpener].
-class SentryDriftDatabase extends LazyDatabase {
+/// If queryExecutor is not null, it will be used instead of the default one.
+/// The default is [LazyDatabase].
+class SentryDriftDatabase extends QueryExecutor {
   final Hub _hub;
+
   final _spanHelper = SentrySpanHelper(
     // ignore: invalid_use_of_internal_member
     SentryTraceOrigins.autoDbDriftDatabaseExecutor,
   );
 
-  String? dbName = 'people-drift-impl';
+  final QueryExecutor _queryExecutor;
+
+  final String? _dbName;
 
   @internal
   // ignore: public_member_api_docs
@@ -41,61 +44,76 @@ class SentryDriftDatabase extends LazyDatabase {
   /// Declares a [SentryDriftDatabase] that will run [opener] when the database is
   /// first requested to be opened. You must specify the same [dialect] as the
   /// underlying database has
-  SentryDriftDatabase(
-    super.opener, {
-    super.dialect,
-    @internal Hub? hub,
-  })  : _hub = hub ?? HubAdapter() {
+  SentryDriftDatabase(DatabaseOpener opener,
+      {@internal Hub? hub,
+      @internal QueryExecutor? queryExecutor,
+      @internal String? dbName})
+      : _hub = hub ?? HubAdapter(),
+        _dbName = dbName,
+        _queryExecutor = queryExecutor ?? LazyDatabase(opener) {
     _spanHelper.setHub(_hub);
   }
 
+  @override
+  TransactionExecutor beginTransaction() {
+    return _queryExecutor.beginTransaction();
+  }
+
+  @override
+  Future<void> runBatched(BatchedStatements statements) {
+    return _queryExecutor.runBatched(statements);
+  }
+
   @override
   Future<bool> ensureOpen(QueryExecutorUser user) {
     return _spanHelper.asyncWrapInSpan('open', () async {
-      return await super.ensureOpen(user);
-    });
+      return await _queryExecutor.ensureOpen(user);
+    }, dbName: _dbName);
   }
 
   @override
   Future<void> runCustom(String statement, [List<Object?>? args]) {
     return _spanHelper.asyncWrapInSpan('custom', () async {
-      return await super.runCustom(statement, args);
-    }, dbName: dbName);
+      return await _queryExecutor.runCustom(statement, args);
+    }, dbName: _dbName);
   }
 
   @override
   Future<int> runDelete(String statement, List<Object?> args) {
     return _spanHelper.asyncWrapInSpan('delete', () async {
-      return await super.runDelete(statement, args);
-    }, dbName: dbName);
+      return await _queryExecutor.runDelete(statement, args);
+    }, dbName: _dbName);
   }
 
   @override
   Future<int> runInsert(String statement, List<Object?> args) {
     return _spanHelper.asyncWrapInSpan('insert', () async {
-      return await super.runInsert(statement, args);
-    }, dbName: dbName);
+      return await _queryExecutor.runInsert(statement, args);
+    }, dbName: _dbName);
   }
 
   @override
   Future<List<Map<String, Object?>>> runSelect(
       String statement, List<Object?> args) {
     return _spanHelper.asyncWrapInSpan('select', () async {
-      return await super.runSelect(statement, args);
-    }, dbName: dbName);
+      return await _queryExecutor.runSelect(statement, args);
+    }, dbName: _dbName);
   }
 
   @override
   Future<int> runUpdate(String statement, List<Object?> args) {
     return _spanHelper.asyncWrapInSpan('update', () async {
-      return await super.runUpdate(statement, args);
-    }, dbName: dbName);
+      return await _queryExecutor.runUpdate(statement, args);
+    }, dbName: _dbName);
   }
 
   @override
   Future<void> close() {
     return _spanHelper.asyncWrapInSpan('close', () async {
-      return await super.close();
-    }, dbName: dbName);
+      return await _queryExecutor.close();
+    }, dbName: _dbName);
   }
+
+  @override
+  SqlDialect get dialect => _queryExecutor.dialect;
 }

drift/lib/src/sentry_span_helper.dart

@@ -45,7 +45,6 @@ class SentrySpanHelper {
 
     try {
       final result = await execute();
-
       span?.status = SpanStatus.ok();
 
       return result;

drift/test/sentry_database_test.dart

@@ -23,11 +23,9 @@ void main() {
     expect(span?.data[SentryDriftDatabase.dbNameKey], Fixture.dbName);
   }
 
-  void verifyErrorSpan(
-    String description,
-    Exception exception,
-    SentrySpan? span,
-  ) {
+  void verifyErrorSpan(String description,
+      Exception exception,
+      SentrySpan? span,) {
     expect(span?.context.operation, SentryDriftDatabase.dbOp);
     expect(span?.context.description, description);
     expect(span?.status, SpanStatus.internalError());
@@ -56,26 +54,26 @@ void main() {
     });
 
     test('insert adds span', () async {
-      final sut = fixture.getSut();
+      final sut = fixture.sut;
 
       await sut.into(sut.todoItems).insert(TodoItemsCompanion.insert(
-            title: 'todo: finish drift setup',
-            content: 'We can now write queries and define our own tables.',
-          ));
+        title: 'todo: finish drift setup',
+        content: 'We can now write queries and define our own tables.',
+      ));
 
       verifySpan('insert', fixture.getCreatedSpan());
     });
 
     test('update adds span', () async {
-      final sut = fixture.getSut();
+      final sut = fixture.sut;
 
       await sut.into(sut.todoItems).insert(TodoItemsCompanion.insert(
-            title: 'todo: finish drift setup',
-            content: 'We can now write queries and define our own tables.',
-          ));
+        title: 'todo: finish drift setup',
+        content: 'We can now write queries and define our own tables.',
+      ));
 
       await (sut.update(sut.todoItems)
-            ..where((tbl) => tbl.title.equals('todo: finish drift setup')))
+        ..where((tbl) => tbl.title.equals('todo: finish drift setup')))
           .write(TodoItemsCompanion(
         title: Value('after update'),
         content: Value('We can now write queries and define our own tables.'),
@@ -85,93 +83,64 @@ void main() {
     });
 
     test('custom adds span', () async {
-      final sut = fixture.getSut();
+      final sut = fixture.sut;
 
       await sut.customStatement('SELECT * FROM todo_items');
 
       verifySpan('custom', fixture.getCreatedSpan());
     });
 
     test('batch adds span', () async {
-      final sut = fixture.getSut();
+      final sut = fixture.sut;
 
       await sut.transaction(() async {
         await sut.into(sut.todoItems).insert(TodoItemsCompanion.insert(
-              title: 'todo: finish drift setup',
-              content: 'We can now write queries and define our own tables.',
-            ));
+          title: 'todo: finish drift setup',
+          content: 'We can now write queries and define our own tables.',
+        ));
+      });
+
+      await sut.batch((batch) async {
+        await sut.into(sut.todoItems).insert(TodoItemsCompanion.insert(
+          title: 'todo: finish drift setup',
+          content: 'We can now write queries and define our own tables.',
+        ));
       });
 
       // TODO
     });
 
     test('close adds span', () async {
-      final sut = fixture.getSut();
+      final sut = fixture.sut;
 
       await sut.close();
 
       verifySpan('close', fixture.getCreatedSpan());
     });
 
     test('open adds span', () async {
-      final sut = fixture.getSut();
+      final sut = fixture.sut;
 
-      // SentryDriftDatabase is lazily opened by default so it won't create a span
-      // until it is actually used.
+      // SentryDriftDatabase is by default lazily opened by default so it won't
+      // create a span until it is actually used.
       await sut.select(sut.todoItems).get();
 
       verifySpan('open', fixture.getSpanByDescription('open'));
     });
 
     test('will not add open span if db is not used', () async {
-      fixture.getSut();
+      fixture.sut;
 
       expect(fixture.tracer.children.isEmpty, true);
     });
   });
-
-  group('adds error span', () {
-    late Fixture fixture;
-
-    setUp(() async {
-      fixture = Fixture();
-      await fixture.setUp();
-
-      when(fixture.hub.options).thenReturn(fixture.options);
-      when(fixture.hub.getSpan()).thenReturn(fixture.tracer);
-    });
-
-    tearDown(() async {
-      await fixture.tearDown();
-    });
-
-    // test('throwing insert adds error span', () async {
-    //   final sut = fixture.getSut();
-    //   Exception? exception;
-    //
-    //   try {
-    //     await sut.into(fixture.sut.todoItems).insert(TodoItemsCompanion.insert(
-    //       title: 'AAAaaaaaa',
-    //       content: 'aaaaaaaaaa',
-    //     ));
-    //   } catch (error) {
-    //     expect(error is InvalidDataException, true);
-    //     exception = error as InvalidDataException;
-    //   }
-    //
-    //
-    //
-    //   //verifyErrorSpan('insert', exception!, fixture.getCreatedSpan());
-    // });
-  });
 }
 
 class Fixture {
   final options = SentryOptions();
   final hub = MockHub();
   static final dbName = 'people-drift-impl';
   final exception = Exception('fixture-exception');
-
   final _context = SentryTransactionContext('name', 'operation');
   late final tracer = SentryTracer(_context, hub);
   late AppDatabase sut;
@@ -184,21 +153,18 @@ class Fixture {
     await sut.close();
   }
 
-  AppDatabase getSut() {
-    return sut;
-  }
-
   SentrySpan? getCreatedSpan() {
     return tracer.children.last;
   }
-  
+
   SentrySpan? getSpanByDescription(String description) {
-    return tracer.children.firstWhere((element) => element.context.description == description);
+    return tracer.children
+        .firstWhere((element) => element.context.description == description);
   }
 
   SentryDriftDatabase openConnection() {
     return SentryDriftDatabase(() {
       return NativeDatabase.memory();
-    }, hub: hub);
+    }, hub: hub, dbName: dbName);
   }
 }