perf(catalog): ERD neighbourhood and MCP lookups read only the tables they need (#1572)

lib/core/erd/erd_catalog.dart

@@ -26,7 +26,17 @@ class ErdCatalog {
 
   /// Columns of [schema] (PostgreSQL) or of the connected database. [schema]
   /// `null` means the current schema.
-  static String columnsSql(SqlDialect d, {String? schema}) => switch (d) {
+  ///
+  /// [tables] limits it to those bare table names (of [schema]): a diagram
+  /// around one table or an MCP lookup of one table then reads those tables
+  /// and not the whole schema, which also kept the row limit from cutting a
+  /// large schema before the tables that were wanted (#1572).
+  static String columnsSql(
+    SqlDialect d, {
+    String? schema,
+    Iterable<String>? tables,
+  }) =>
+      switch (d) {
         // pg_catalog is readable by every role, unlike information_schema,
         // which hides keys of tables the current role has only SELECT on.
         SqlDialect.postgres => '''
@@ -61,7 +71,7 @@ LEFT JOIN pg_catalog.pg_constraint pk
 LEFT JOIN pg_catalog.pg_attrdef df
   ON df.adrelid = c.oid AND df.adnum = a.attnum
 WHERE ns.nspname = ${schema == null ? 'current_schema()' : "'${_lit(schema)}'"}
-  AND c.relkind IN ('r', 'p')
+  AND c.relkind IN ('r', 'p')${_nameFilter('c.relname', tables)}
   AND a.attnum > 0 AND NOT a.attisdropped
 ORDER BY table_name, a.attnum''',
         SqlDialect.mysql => '''
@@ -77,7 +87,7 @@ FROM information_schema.columns c
 JOIN information_schema.tables t
   ON t.table_schema = c.table_schema AND t.table_name = c.table_name
   AND t.table_type = 'BASE TABLE'
-WHERE c.table_schema = DATABASE()
+WHERE c.table_schema = DATABASE()${_nameFilter('c.table_name', tables)}
 ORDER BY c.table_name, c.ordinal_position''',
         SqlDialect.sqlite => '''
 SELECT m.name AS table_name, p.name AS column_name, p.type AS data_type,
@@ -96,14 +106,16 @@ SELECT m.name AS table_name, p.name AS column_name, p.type AS data_type,
   NULL AS column_comment,
   NULL AS table_comment
 FROM sqlite_master m JOIN pragma_table_info(m.name) p
-WHERE m.type = 'table' AND m.name NOT LIKE 'sqlite_%'
+WHERE m.type = 'table' AND m.name NOT LIKE 'sqlite_%'${_nameFilter('m.name', tables)}
 ORDER BY m.name, p.cid''',
       };
 
   /// Columns of the views (and PostgreSQL materialized views) of the current
   /// schema or database, in the shape of [columnsSql]: no primary key, every
-  /// column nullable. MCP lists them next to the tables.
-  static String viewColumnsSql(SqlDialect d) => switch (d) {
+  /// column nullable. MCP lists them next to the tables. [tables] limits it
+  /// to those names, as in [columnsSql].
+  static String viewColumnsSql(SqlDialect d, {Iterable<String>? tables}) =>
+      switch (d) {
         SqlDialect.postgres => '''
 SELECT ${_pgQualified('c', 'ns')} AS table_name, a.attname AS column_name,
   format_type(a.atttypid, a.atttypmod) AS data_type,
@@ -112,7 +124,7 @@ FROM pg_catalog.pg_attribute a
 JOIN pg_catalog.pg_class c ON c.oid = a.attrelid
 JOIN pg_catalog.pg_namespace ns ON ns.oid = c.relnamespace
 WHERE ns.nspname = current_schema()
-  AND c.relkind IN ('v', 'm')
+  AND c.relkind IN ('v', 'm')${_nameFilter('c.relname', tables)}
   AND a.attnum > 0 AND NOT a.attisdropped
 ORDER BY table_name, a.attnum''',
         SqlDialect.mysql => '''
@@ -121,13 +133,13 @@ FROM information_schema.columns c
 JOIN information_schema.tables t
   ON t.table_schema = c.table_schema AND t.table_name = c.table_name
   AND t.table_type = 'VIEW'
-WHERE c.table_schema = DATABASE()
+WHERE c.table_schema = DATABASE()${_nameFilter('c.table_name', tables)}
 ORDER BY c.table_name, c.ordinal_position''',
         SqlDialect.sqlite => '''
 SELECT m.name AS table_name, p.name AS column_name, p.type AS data_type,
   0 AS is_pk, 1 AS is_nullable
 FROM sqlite_master m JOIN pragma_table_info(m.name) p
-WHERE m.type = 'view' AND m.name NOT LIKE 'sqlite_%'
+WHERE m.type = 'view' AND m.name NOT LIKE 'sqlite_%'${_nameFilter('m.name', tables)}
 ORDER BY m.name, p.cid''',
       };
 
@@ -165,8 +177,37 @@ WHERE m.type = 'table' AND m.name NOT LIKE 'sqlite_%'
 ''',
       };
 
+  /// Names of the tables and views of the current schema or database, and
+  /// nothing else: what a lookup by name needs, without reading columns.
+  static String tableNamesSql(SqlDialect d) => switch (d) {
+        SqlDialect.postgres => '''
+SELECT c.relname AS object_name
+FROM pg_catalog.pg_class c
+JOIN pg_catalog.pg_namespace ns ON ns.oid = c.relnamespace
+WHERE ns.nspname = current_schema() AND c.relkind IN ('r', 'p', 'v', 'm')
+ORDER BY 1''',
+        SqlDialect.mysql => '''
+SELECT table_name AS object_name
+FROM information_schema.tables
+WHERE table_schema = DATABASE() AND table_type IN ('BASE TABLE', 'VIEW')
+ORDER BY 1''',
+        SqlDialect.sqlite => '''
+SELECT name AS object_name
+FROM sqlite_master
+WHERE type IN ('table', 'view') AND name NOT LIKE 'sqlite_%'
+ORDER BY 1''',
+      };
+
   static String _lit(String s) => s.replaceAll("'", "''");
 
+  /// `AND [column] IN ('a', 'b')` for [names], or nothing when [names] is
+  /// null; an empty list matches no row.
+  static String _nameFilter(String column, Iterable<String>? names) {
+    if (names == null) return '';
+    if (names.isEmpty) return '\n  AND 1 = 0';
+    return "\n  AND $column IN (${names.map((n) => "'${_lit(n)}'").join(', ')})";
+  }
+
   /// Schema part of a diagram name: `sales` for `sales.orders`, `null` for a
   /// table of the current schema.
   static String? schemaOf(String name) {
@@ -217,11 +258,18 @@ WHERE m.type = 'table' AND m.name NOT LIKE 'sqlite_%'
     }
     var names = neighbourhood(fks, focus, depth);
 
+    // Only the columns of these tables, grouped by PostgreSQL schema (#1572).
+    // `schema.table` asks its schema for `table`, which also brings the bare
+    // name the fallback below looks for when that schema is the current one.
+    final wanted = <String?, Set<String>>{};
+    for (final n in names) {
+      final s = dialect == SqlDialect.postgres ? schemaOf(n) : null;
+      (wanted[s] ??= <String>{}).add(s == null ? n : n.substring(s.length + 1));
+    }
     final fetched = <List<String>>[];
-    final schemas = {for (final n in names) schemaOf(n)};
-    for (final s in schemas) {
+    for (final MapEntry(key: s, value: tables) in wanted.entries) {
       final rows = await delegate.executeQuery(
-        columnsSql(dialect, schema: dialect == SqlDialect.postgres ? s : null),
+        columnsSql(dialect, schema: s, tables: tables),
         limit: catalogRowLimit,
       );
       truncated = truncated || rows.isTruncated;

lib/core/mcp/mcp_query_service.dart

@@ -107,9 +107,10 @@ class McpQueryService {
     String table,
   ) {
     return _withDelegate(connectionId, (delegate, dialect) async {
-      final cat = await _loadCatalog(delegate, dialect);
+      final name = await _resolveTable(delegate, dialect, table);
+      final cat = await _loadCatalog(delegate, dialect, tables: [name]);
       final schema = cat.schema;
-      final t = _findTable(schema, table);
+      final t = _findTable(schema, name);
       if (cat.views.contains(t.name)) return _describe(schema, t, const []);
       final indexes = await _run(
         () => delegate.executeQuery(
@@ -162,11 +163,10 @@ class McpQueryService {
   }) {
     final n = rows.clamp(1, maxSampleRows);
     return _withDelegate(connectionId, (delegate, dialect) async {
-      final t =
-          _findTable((await _loadCatalog(delegate, dialect)).schema, table);
+      final name = await _resolveTable(delegate, dialect, table);
       final result = await _run(
         () => delegate.executeQuery(
-          'SELECT * FROM ${quoteIdentifier(t.name, dialect)}',
+          'SELECT * FROM ${quoteIdentifier(name, dialect)}',
           limit: n,
           timeout: timeout,
         ),
@@ -242,14 +242,19 @@ class McpQueryService {
 
   /// Tables and keys, then the views: a name that is already a table stays a
   /// table. Views have no keys and no indexes.
+  /// The catalog of the current schema or database; with [tables], of those
+  /// tables only (#1572).
   Future<({ErdSchema schema, Set<String> views})> _loadCatalog(
     SqlQueryRunner delegate,
-    SqlDialect dialect,
-  ) async {
-    final base = await _loadSchema(delegate, dialect);
+    SqlDialect dialect, {
+    Iterable<String>? tables,
+  }) async {
+    final base = await _loadSchema(delegate, dialect, tables: tables);
     final viewRows = await _run(
-      () => delegate.executeQuery(ErdCatalog.viewColumnsSql(dialect),
-          limit: _catalogRows, timeout: timeout),
+      () => delegate.executeQuery(
+          ErdCatalog.viewColumnsSql(dialect, tables: tables),
+          limit: _catalogRows,
+          timeout: timeout),
       dialect,
     );
     final known = {for (final t in base.tables) t.name};
@@ -271,21 +276,59 @@ class McpQueryService {
 
   Future<ErdSchema> _loadSchema(
     SqlQueryRunner delegate,
-    SqlDialect dialect,
-  ) async {
-    final cols = await _run(
-      () => delegate.executeQuery(ErdCatalog.columnsSql(dialect),
-          limit: _catalogRows, timeout: timeout),
-      dialect,
-    );
+    SqlDialect dialect, {
+    Iterable<String>? tables,
+  }) async {
     final fks = await _run(
       () => delegate.executeQuery(ErdCatalog.foreignKeysSql(dialect),
           limit: _catalogRows, timeout: timeout),
       dialect,
     );
+    // With [tables], also the tables their keys point to: ErdSchema drops a
+    // relation to a table it has no columns for, and a description lists
+    // each column's reference.
+    final wanted = tables == null
+        ? null
+        : {
+            ...tables,
+            for (final r in fks.rows)
+              if (r.length >= 4 && tables.contains(r[0])) r[2],
+          };
+    final cols = await _run(
+      () => delegate.executeQuery(
+          ErdCatalog.columnsSql(dialect, tables: wanted),
+          limit: _catalogRows,
+          timeout: timeout),
+      dialect,
+    );
     return ErdSchema.fromCatalog(columnRows: cols.rows, fkRows: fks.rows);
   }
 
+  /// The catalog's name for [name]: the exact one, else the only one that
+  /// matches ignoring case. Reads the names only, not the whole catalog: one
+  /// table's description or sample no longer loads every column (#1572).
+  Future<String> _resolveTable(
+    SqlQueryRunner delegate,
+    SqlDialect dialect,
+    String name,
+  ) async {
+    final result = await _run(
+      () => delegate.executeQuery(ErdCatalog.tableNamesSql(dialect),
+          limit: _catalogRows, timeout: timeout),
+      dialect,
+    );
+    final names = [
+      for (final r in result.rows)
+        if (r.isNotEmpty) r.first,
+    ];
+    if (names.contains(name)) return name;
+    final lower = name.toLowerCase();
+    final loose = names.where((n) => n.toLowerCase() == lower).toList();
+    if (loose.length == 1) return loose.single;
+    throw McpToolException(
+        'Table "$name" not found. Call list_tables to see the available tables.');
+  }
+
   /// Exact name first, then a unique case-insensitive match. The table must
   /// exist in the catalog, so model input never reaches SQL unchecked.
   ErdTable _findTable(ErdSchema schema, String name) {

test/core/mcp/mcp_query_service_test.dart

@@ -3,7 +3,6 @@ import 'dart:convert';
 
 import 'package:flutter_test/flutter_test.dart';
 import 'package:querya_desktop/core/database/table_mutation_engine.dart';
-import 'package:querya_desktop/core/erd/erd_catalog.dart';
 import 'package:querya_desktop/core/mcp/mcp_access_store.dart';
 import 'package:querya_desktop/core/mcp/mcp_query_service.dart';
 import 'package:querya_desktop/core/security/connection_environment.dart';
@@ -55,6 +54,12 @@ ConnectionRow _row(int id, String type, {String? env}) {
 }
 
 SqlExecutionResult _catalog(String sql) {
+  if (sql.contains('object_name')) {
+    return const SqlExecutionResult(columns: ['n'], rows: [
+      ['orders'],
+      ['users'],
+    ]);
+  }
   if (sql.contains('pg_indexes')) {
     return const SqlExecutionResult(columns: ['n', 'd'], rows: [
       ['users_pkey', 'CREATE UNIQUE INDEX users_pkey ON users (id)'],
@@ -226,7 +231,14 @@ void main() {
     group('views', () {
       setUp(() {
         db.onExecute = (sql) {
-          if (sql == ErdCatalog.viewColumnsSql(SqlDialect.postgres)) {
+          if (sql.contains('object_name')) {
+            return const SqlExecutionResult(columns: ['n'], rows: [
+              ['active_users'],
+              ['orders'],
+              ['users'],
+            ]);
+          }
+          if (sql.contains("relkind IN ('v', 'm')")) {
             return const SqlExecutionResult(columns: ['t', 'c', 'ty', 'pk'], rows: [
               ['active_users', 'id', 'integer', '0'],
               ['active_users', 'name', 'text', '0'],
@@ -272,6 +284,23 @@ void main() {
       expect(orders.columns.last.references, 'users.id');
     });
 
+    test('describeTable and sampleRows read only the table asked for (#1572)',
+        () async {
+      await service().describeTable(1, 'users');
+      // Column and view-column queries (the key query has no format_type).
+      final columnQueries =
+          db.executed.where((s) => s.contains('format_type')).toList();
+      expect(columnQueries, isNotEmpty);
+      for (final q in columnQueries) {
+        expect(q, contains("AND c.relname IN ('users')"));
+      }
+
+      db.executed.clear();
+      await service().sampleRows(1, 'USERS', rows: 3);
+      expect(db.executed.where((s) => s.contains('format_type')), isEmpty);
+      expect(db.executed.last, 'SELECT * FROM "users"');
+    });
+
     test('an unknown table is an error and runs no table query', () async {
       await expectLater(service().describeTable(1, 'users; DROP TABLE x'),
           throwsA(isA<McpToolException>()));

test/features/erd/erd_layout_persistence_test.dart

@@ -30,7 +30,7 @@ const _key = ErdLayoutKey(connectionId: 1, scope: 'shop');
 
 FakeSqlExecutionDelegate _catalog({bool withPayments = false}) =>
     FakeSqlExecutionDelegate(onExecute: (sql) {
-      if (sql == ErdCatalog.columnsSql(SqlDialect.sqlite)) {
+      if (_unfiltered(sql) == ErdCatalog.columnsSql(SqlDialect.sqlite)) {
         return SqlExecutionResult(rows: [
           const ['users', 'id', 'INTEGER', '1'],
           const ['orders', 'user_id', 'INTEGER', '0'],
@@ -45,6 +45,10 @@ FakeSqlExecutionDelegate _catalog({bool withPayments = false}) =>
 Finder _card(String table) =>
     find.byKey(material.ValueKey('erd_table_$table'));
 
+/// [sql] without the table-name filter `columnsSql(tables: ...)` adds (#1572).
+String _unfiltered(String sql) => sql.replaceAll(
+    RegExp(r"\n  AND (c\.relname|c\.table_name|m\.name) IN \([^)]*\)"), '');
+
 void main() {
   Future<void> pumpDiagram(
     WidgetTester t,