Add service extension to run statements

drift/lib/src/runtime/devtools/devtools.dart

@@ -9,18 +9,22 @@ import 'dart:developer' as developer;
 import 'package:meta/meta.dart';
 
 import '../api/runtime_api.dart';
+import 'service_extension.dart';
 import 'shared.dart';
 
 const _releaseMode = bool.fromEnvironment('dart.vm.product');
 const _profileMode = bool.fromEnvironment('dart.vm.profile');
+
+// Avoid pulling in a bunch of unused code to describe databases and to make
+// them available through service extensions on release builds.
 const _enable = !_releaseMode && !_profileMode;
 
-void _postEvent(String type, Map<Object?, Object?> data) {
+void postEvent(String type, Map<Object?, Object?> data) {
   developer.postEvent('drift:$type', data);
 }
 
 void _postChangedEvent() {
-  _postEvent('database-list-changed', {});
+  postEvent('database-list-changed', {});
 }
 
 class TrackedDatabase {
@@ -41,6 +45,7 @@ class TrackedDatabase {
 void handleCreated(GeneratedDatabase database) {
   if (_enable) {
     TrackedDatabase(database);
+    DriftServiceExtension.registerIfNeeded();
     _postChangedEvent();
   }
 }

drift/lib/src/runtime/devtools/service_extension.dart

@@ -0,0 +1,98 @@
+import 'dart:async';
+import 'dart:convert';
+import 'dart:developer';
+
+import 'package:drift/src/remote/protocol.dart';
+
+import '../query_builder/query_builder.dart';
+import 'devtools.dart';
+
+/// A service extension making asynchronous requests on drift databases
+/// accessible via the VM service.
+///
+/// This is used by the drift DevTools extension to run queries and show their
+/// results in the DevTools view.
+class DriftServiceExtension {
+  int _subscriptionId = 0;
+  final Map<int, StreamSubscription> _activeSubscriptions = {};
+
+  Future<Object?> _handle(Map<String, String> parameters) async {
+    final action = parameters['action']!;
+    final databaseId = int.parse(parameters['db']!);
+    final tracked = TrackedDatabase.all.firstWhere((e) => e.id == databaseId);
+
+    switch (action) {
+      case 'subscribe-to-tables':
+        final stream = tracked.database.tableUpdates();
+        final id = _subscriptionId++;
+
+        stream.listen((event) {
+          postEvent('event', {
+            'subscription': id,
+            'payload':
+                _protocol.encodePayload(NotifyTablesUpdated(event.toList()))
+          });
+        });
+
+        return id.toString();
+      case 'unsubscribe-from-tables':
+        _activeSubscriptions.remove(int.parse(parameters['id']!))?.cancel();
+        return null;
+      case 'execute-query':
+        final execute = _protocol
+            .decodePayload(json.decode(parameters['query']!)) as ExecuteQuery;
+        final variables = [
+          for (final variable in execute.args) Variable(variable)
+        ];
+
+        final result = await switch (execute.method) {
+          StatementMethod.select => tracked.database
+              .customSelect(execute.sql, variables: variables)
+              .get()
+              .then((rows) => SelectResult([for (final row in rows) row.data])),
+          StatementMethod.insert =>
+            tracked.database.customInsert(execute.sql, variables: variables),
+          StatementMethod.deleteOrUpdate =>
+            tracked.database.customUpdate(execute.sql, variables: variables),
+          StatementMethod.custom => tracked.database
+              .customStatement(execute.sql, execute.args)
+              .then((_) => 0),
+        };
+
+        return _protocol.encodePayload(result);
+      default:
+        throw UnsupportedError('Method $action');
+    }
+  }
+
+  static bool _registered = false;
+
+  /// Registers the `ext.drift.database` extension if it has not yet been
+  /// registered on this isolate.
+  static void registerIfNeeded() {
+    if (!_registered) {
+      _registered = true;
+
+      final extension = DriftServiceExtension();
+      registerExtension('ext.drift.database', (method, parameters) {
+        return Future(() => extension._handle(parameters))
+            .then((value) => ServiceExtensionResponse.result(json.encode({
+                  'r': value,
+                })))
+            .onError((error, stackTrace) {
+          return ServiceExtensionResponse.error(
+            ServiceExtensionResponse.extensionErrorMin,
+            json.encode(
+              {
+                'e': error.toString(),
+                'trace': stackTrace.toString(),
+              },
+            ),
+          );
+        });
+      });
+    }
+  }
+
+  static const _protocol = DriftProtocol();
+}

extras/drift_devtools_extension/lib/main.dart

@@ -1,13 +1,23 @@
 import 'package:devtools_app_shared/ui.dart';
 import 'package:devtools_extensions/devtools_extensions.dart';
+import 'package:flutter/foundation.dart';
 import 'package:flutter/material.dart';
 import 'package:flutter_riverpod/flutter_riverpod.dart';
+import 'package:logging/logging.dart';
 
 import 'src/details.dart';
 import 'src/list.dart';
 
 void main() {
   runApp(const ProviderScope(child: DriftDevToolsExtension()));
+
+  if (kDebugMode) {
+    Logger.root.level = Level.FINE;
+    Logger.root.onRecord.listen((record) {
+      debugPrint(
+          '[${record.level.name}] ${record.loggerName}: ${record.message}');
+    });
+  }
 }
 
 class DriftDevToolsExtension extends StatelessWidget {

extras/drift_devtools_extension/lib/src/details.dart

@@ -1,9 +1,6 @@
 import 'package:devtools_app_shared/service.dart';
-import 'package:devtools_extensions/devtools_extensions.dart';
 import 'package:flutter/material.dart';
-import 'package:flutter/src/widgets/framework.dart';
 import 'package:flutter_riverpod/flutter_riverpod.dart';
-import 'package:vm_service/vm_service.dart';
 
 import 'list.dart';
 import 'remote_database.dart';
@@ -16,15 +13,23 @@ final loadedDatabase = AutoDisposeFutureProvider((ref) async {
   final isAlive = Disposable();
   ref.onDispose(isAlive.dispose);
 
-  if (selected?.database case InstanceRef dbRef) {
-    final db = await eval.safeGetInstance(dbRef, isAlive);
-
-    return await RemoteDatabase.resolve(db, eval, isAlive);
+  if (selected != null) {
+    return await RemoteDatabase.resolve(selected, eval, isAlive);
   }
 
   return null;
 });
 
+final _testQuery = AutoDisposeFutureProvider((ref) async {
+  final database = await ref.watch(loadedDatabase.future);
+
+  if (database != null) {
+    return await database.select('SELECT 1, 2, 3', []);
+  } else {
+    return null;
+  }
+});
+
 class DatabaseDetails extends ConsumerStatefulWidget {
   const DatabaseDetails({super.key});
 
@@ -46,6 +51,7 @@ class _DatabaseDetailsState extends ConsumerState<DatabaseDetails> {
   @override
   Widget build(BuildContext context) {
     final database = ref.watch(loadedDatabase);
+    final query = ref.watch(_testQuery);
 
     return database.when(
       loading: () => const Center(child: CircularProgressIndicator()),
@@ -59,6 +65,7 @@ class _DatabaseDetailsState extends ConsumerState<DatabaseDetails> {
               children: [
                 for (final entity in database.description.entities)
                   Text('${entity.name}: ${entity.type}'),
+                Text(query.toString()),
               ],
             ),
           );

extras/drift_devtools_extension/lib/src/remote_database.dart

@@ -1,30 +1,118 @@
+import 'dart:async';
 import 'dart:convert';
 
 import 'package:devtools_app_shared/service.dart';
+import 'package:devtools_extensions/devtools_extensions.dart';
+import 'package:drift/drift.dart';
 // ignore: invalid_use_of_internal_member, implementation_imports
 import 'package:drift/src/runtime/devtools/shared.dart';
-import 'package:vm_service/vm_service.dart';
+// ignore: implementation_imports
+import 'package:drift/src/remote/protocol.dart';
+import 'package:logging/logging.dart';
+import 'list.dart';
 
 /// Utilities to access a drift database via service extensions.
 class RemoteDatabase {
+  final TrackedDatabase db;
   final DatabaseDescription description;
 
-  RemoteDatabase({required this.description});
+  int? _remoteSubscriptionId;
+  StreamSubscription? _tableNotifications;
+  final StreamController<List<TableUpdate>> _tableUpdates =
+      StreamController.broadcast();
+
+  Stream<List<TableUpdate>> get tableUpdates => _tableUpdates.stream;
+
+  RemoteDatabase({required this.db, required this.description}) {
+    _tableUpdates
+      ..onListen = () async {
+        try {
+          final id = await _newTableSubscription();
+
+          if (_remoteSubscriptionId != null) {
+            await _unsubscribeFromTables(_remoteSubscriptionId!);
+          }
+          _remoteSubscriptionId = id;
+          _logger.fine('Received subscription $id for tables.');
+
+          _tableNotifications =
+              serviceManager.service!.onExtensionEvent.where((e) {
+            return e.extensionKind == 'drift:event' &&
+                e.extensionData?.data['subscription'] == id;
+          }).listen((event) {
+            final payload =
+                _protocol.decodePayload(event.extensionData?.data['payload'])
+                    as NotifyTablesUpdated;
+            _tableUpdates.add(payload.updates);
+          });
+        } catch (e, s) {
+          _tableUpdates
+            ..addError(e, s)
+            ..close();
+        }
+      }
+      ..onCancel = () {
+        final id = _remoteSubscriptionId;
+        _remoteSubscriptionId = null;
+        if (id != null) {
+          _unsubscribeFromTables(id);
+        }
+      };
+  }
+
+  Future<List<Map<String, Object?>>> select(
+      String sql, List<Object?> args) async {
+    final result = await _driftRequest('execute-query', payload: {
+      'query': json.encode(_protocol
+          .encodePayload(ExecuteQuery(StatementMethod.select, sql, args)))
+    });
+
+    return (_protocol.decodePayload(result) as SelectResult).rows;
+  }
+
+  Future<int> _newTableSubscription() async {
+    final result = await _driftRequest('subscribe-to-tables');
+    return result as int;
+  }
+
+  Future<void> _unsubscribeFromTables(int subId) async {
+    await _driftRequest('unsubscribe-from-tables',
+        payload: {'id': subId.toString()});
+    await _tableNotifications?.cancel();
+    _tableNotifications = null;
+  }
+
+  Future<Object?> _driftRequest(String method,
+      {Map<String, String> payload = const {}}) async {
+    final response = await serviceManager.callServiceExtensionOnMainIsolate(
+      'ext.drift.database',
+      args: {
+        'action': method,
+        'db': db.id.toString(),
+        ...payload,
+      },
+    );
+
+    return response.json!['r'];
+  }
 
   static Future<RemoteDatabase> resolve(
-    Instance database,
+    TrackedDatabase database,
     EvalOnDartLibrary eval,
     Disposable isAlive,
   ) async {
     final stringVal = await eval.evalInstance(
       'describe(db)',
       isAlive: isAlive,
-      scope: {'db': database.id!},
+      scope: {'db': database.database.id!},
     );
     final value = await eval.retrieveFullValueAsString(stringVal);
 
     final description = DatabaseDescription.fromJson(json.decode(value!));
 
-    return RemoteDatabase(description: description);
+    return RemoteDatabase(db: database, description: description);
   }
+
+  static final _logger = Logger('RemoteDatabase');
+  static const _protocol = DriftProtocol();
 }

extras/drift_devtools_extension/pubspec.yaml

@@ -43,6 +43,7 @@ dependencies:
   vm_service: ^11.10.0
   path: ^1.8.3
   drift: ^2.12.1
+  logging: ^1.2.0
 
 dev_dependencies:
   flutter_test: