Fine-grained query intercepting

docs/docs/Examples/tracing.md

@@ -5,8 +5,6 @@ description: Using the `QueryInterceptor` API to log details about database oper
 
 ---
 
-
-
 Drift provides the relatively simple `logStatements` option to print the statements it
 executes.
 The `QueryInterceptor` API can be used to extend this logging to provide more information,
@@ -19,6 +17,11 @@ Interceptors can be applied with the `interceptWith` extension on `QueryExecutor
 
 {{ load_snippet('use','lib/snippets/log_interceptor.dart.excerpt.json') }}
 
+If you only want to apply an interceptor on a certain block instead of on the whole database,
+that's possible too:
+
+{{ load_snippet('runWithInterceptor','lib/snippets/log_interceptor.dart.excerpt.json') }}
+
 The `QueryInterceptor` class is pretty powerful, as it allows you to fully control the underlying
 database connection. You could also use it to retry some failing statements or to aggregate
 statistics about query times to an external monitoring service.

docs/lib/snippets/log_interceptor.dart

@@ -4,6 +4,8 @@ import 'dart:io';
 import 'package:drift/drift.dart';
 import 'package:drift/native.dart';
 
+import 'dart_api/manager.dart';
+
 // #docregion class
 class LogInterceptor extends QueryInterceptor {
   Future<T> _run<T>(
@@ -90,3 +92,14 @@ void use() {
   ).interceptWith(LogInterceptor());
   // #enddocregion use
 }
+
+void useScoped() async {
+  final database = AppDatabase(NativeDatabase.memory());
+
+  // #docregion runWithInterceptor
+  final interceptor = LogInterceptor();
+  await database.runWithInterceptor(interceptor: interceptor, () async {
+    // Only database operations in this block will reach the interceptor.
+  });
+  // #enddocregion runWithInterceptor
+}

drift/CHANGELOG.md

@@ -5,6 +5,8 @@
 - Deprecate `TypeConverter.json` utility in favor of `TypeConverter.json2`. The
   new method avoids encoding values twice when mapping drift row classes to
   JSON.
+- Add `runWithInterceptor` method to databases to only apply interceptors in
+  a restricted block.
 
 ## 2.23.1
 

drift/lib/src/runtime/api/connection_user.dart

@@ -613,6 +613,19 @@ abstract class DatabaseConnectionUser {
     }
   }
 
+  /// Executes [action] with calls intercepted by the given [interceptor]
+  ///
+  /// This can be used to, for instance, write a custom statement logger or to
+  /// retry failing statements automatically.
+  Future<T> runWithInterceptor<T>(Future<T> Function() action,
+      {required QueryInterceptor interceptor}) async {
+    return await resolvedEngine.doWhenOpened((executor) {
+      final inner = _ExclusiveExecutor(this,
+          executor: executor.interceptWith(interceptor));
+      return _runConnectionZoned(inner, action);
+    });
+  }
+
   /// Runs [calculation] in a forked [Zone] that has its [resolvedEngine] set
   /// to the [user].
   @protected

drift/test/engines/interceptor_test.dart

@@ -29,6 +29,35 @@ void main() {
     await database.users.delete().go();
     expect(events, ['insert', 'update', 'delete']);
   });
+
+  test('calls interceptor methods in runWithInterceptor', () async {
+    final interceptor = EmittingInterceptor();
+    final events = <String>[];
+    interceptor.events.stream.listen(events.add);
+
+    final database = TodoDb(testInMemoryDatabase());
+    expect(await database.categories.select().get(), isEmpty);
+
+    await database.runWithInterceptor(
+      () => database.batch((batch) {
+        batch.insert(database.categories,
+            CategoriesCompanion.insert(description: 'from batch'));
+      }),
+      interceptor: interceptor,
+    );
+    expect(events, ['begin', 'batched', 'commit']);
+    events.clear();
+
+    await database.users.insertOne(
+        UsersCompanion.insert(name: 'Simon B', profilePicture: Uint8List(0)));
+    await database.runWithInterceptor(
+      () =>
+          database.users.update().write(UsersCompanion(isAwesome: Value(true))),
+      interceptor: interceptor,
+    );
+    await database.users.delete().go();
+    expect(events, ['update']);
+  });
 }
 
 class EmittingInterceptor extends QueryInterceptor {

drift/test/generated/todos.mocks.dart

@@ -1478,6 +1478,57 @@ class MockTodoDb extends _i1.Mock implements _i3.TodoDb {
         returnValueForMissingStub: _i5.Future<void>.value(),
       ) as _i5.Future<void>);
 
+  @override
+  _i5.Future<T> runWithInterceptor<T>(
+    _i5.Future<T> Function()? action, {
+    required _i2.QueryInterceptor? interceptor,
+  }) =>
+      (super.noSuchMethod(
+        Invocation.method(
+          #runWithInterceptor,
+          [action],
+          {#interceptor: interceptor},
+        ),
+        returnValue: _i6.ifNotNull(
+              _i6.dummyValueOrNull<T>(
+                this,
+                Invocation.method(
+                  #runWithInterceptor,
+                  [action],
+                  {#interceptor: interceptor},
+                ),
+              ),
+              (T v) => _i5.Future<T>.value(v),
+            ) ??
+            _FakeFuture_26<T>(
+              this,
+              Invocation.method(
+                #runWithInterceptor,
+                [action],
+                {#interceptor: interceptor},
+              ),
+            ),
+        returnValueForMissingStub: _i6.ifNotNull(
+              _i6.dummyValueOrNull<T>(
+                this,
+                Invocation.method(
+                  #runWithInterceptor,
+                  [action],
+                  {#interceptor: interceptor},
+                ),
+              ),
+              (T v) => _i5.Future<T>.value(v),
+            ) ??
+            _FakeFuture_26<T>(
+              this,
+              Invocation.method(
+                #runWithInterceptor,
+                [action],
+                {#interceptor: interceptor},
+              ),
+            ),
+      ) as _i5.Future<T>);
+
   @override
   _i2.GenerationContext $write(
     _i2.Component? component, {