Improvements

getsentry · Nov 7, 2023 · df27999 · df27999
1 parent d0a18eb
commit df27999
Show file tree

Hide file tree

Showing 4 changed files with 36 additions and 16 deletions.
diff --git a/drift/lib/src/sentry_query_executor.dart b/drift/lib/src/sentry_query_executor.dart
@@ -67,11 +67,13 @@ class SentryQueryExecutor extends QueryExecutor {
   @override
   TransactionExecutor beginTransaction() {
     final transactionExecutor = _queryExecutor.beginTransaction();
-    return SentryTransactionExecutor(
+    final sentryTransactionExecutor = SentryTransactionExecutor(
       transactionExecutor,
       _hub,
       dbName: _dbName,
     );
+    sentryTransactionExecutor.beginTransaction();
+    return sentryTransactionExecutor;
   }
 
   @override

diff --git a/drift/lib/src/sentry_span_helper.dart b/drift/lib/src/sentry_span_helper.dart
@@ -27,8 +27,12 @@ class SentrySpanHelper {
     String description,
     Future<T> Function() execute, {
     String? dbName,
+    bool useTransactionSpan = false,
   }) async {
-    final currentSpan = _hub.getSpan();
+    ISentrySpan? currentSpan = _hub.getSpan();
+    if (useTransactionSpan) {
+      currentSpan = transactionSpan;
+    }
     final span = currentSpan?.startChild(
       SentryQueryExecutor.dbOp,
       description: description,

diff --git a/drift/lib/src/sentry_transaction_executor.dart b/drift/lib/src/sentry_transaction_executor.dart
@@ -22,7 +22,6 @@ class SentryTransactionExecutor extends TransactionExecutor {
       : _hub = hub,
         _dbName = dbName {
     _spanHelper.setHub(_hub);
-    beginTransaction();
   }
 
   @override
@@ -81,7 +80,14 @@ class SentryTransactionExecutor extends TransactionExecutor {
 
   @override
   Future<int> runInsert(String statement, List<Object?> args) {
-    return _executor.runInsert(statement, args);
+    return _spanHelper.asyncWrapInSpan(
+      'within transaction: $statement',
+      () async {
+        return _executor.runInsert(statement, args);
+      },
+      dbName: _dbName,
+      useTransactionSpan: true,
+    );
   }
 
   @override

diff --git a/drift/test/sentry_database_test.dart b/drift/test/sentry_database_test.dart
@@ -175,10 +175,16 @@ void main() {
       });
 
       verifySpan(
-        'transaction',
+        'within transaction: $expectedInsertStatement',
         fixture.getCreatedSpan(),
         origin: SentryTraceOrigins.autoDbDriftTransactionExecutor,
       );
+
+      verifySpan(
+        'transaction',
+        fixture.getCreatedSpanByDescription('transaction'),
+        origin: SentryTraceOrigins.autoDbDriftTransactionExecutor,
+      );
     });
 
     test('transaction rollback adds span', () async {
@@ -379,19 +385,21 @@ void main() {
       when(mockTransactionExecutor.beginTransaction())
           .thenThrow(fixture.exception);
 
-      try {
-        // We need to move it inside the try/catch becaue SentryTransactionExecutor
-        // starts beginTransaction() directly after init
-        final SentryTransactionExecutor transactionExecutor =
-            SentryTransactionExecutor(
-          mockTransactionExecutor,
-          fixture.hub,
-          dbName: Fixture.dbName,
-        );
+      // We need to move it inside the try/catch becaue SentryTransactionExecutor
+      // starts beginTransaction() directly after init
+      final SentryTransactionExecutor transactionExecutor =
+      SentryTransactionExecutor(
+        mockTransactionExecutor,
+        fixture.hub,
+        dbName: Fixture.dbName,
+      );
 
-        when(fixture.mockLazyDatabase.beginTransaction())
-            .thenReturn(transactionExecutor);
+      when(fixture.mockLazyDatabase.beginTransaction())
+          .thenReturn(transactionExecutor);
 
+      when(fixture.mockLazyDatabase.runInsert(any, any)).thenAnswer((realInvocation) => Future.value(1));
+
+      try {
         await fixture.sut.batch((batch) async {
           await insertRow(fixture.sut);
         });