feat(mcp): search every schema from search_schema when no schema is named (#3097)

* feat(mcp): search every schema from search_schema when no schema is named

* fix(mcp): fail search_schema on a lost connection and take the column schema from the driver

Author: datlechin

CHANGELOG.md

@@ -67,6 +67,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - **Show Previous Window Tab** and **Show Next Window Tab** for window tabs, with no default shortcut.
 - SQLite 3.53.4 built into the SQLite and libSQL drivers in place of the macOS copy.
 - One-time reset of Open Quickly's Recent query history, and of its objects on connections that switch databases.
+- Tables and views from every schema in the MCP `search_schema` tool when no schema is named. (#3048)
 
 ### Removed
 

TablePro/Core/MCP/MCPConnectionBridge+Data.swift

@@ -166,45 +166,65 @@ extension MCPConnectionBridge {
         return pagination.clampedRowCount(request.limit)
     }
 
-    func searchSchema(scope: DatabaseScope, term: String, limit: Int) async throws -> JsonValue {
-        try await ensureConnected(scope.connectionId)
-        let schema = scope.schema
-        let needle = term.lowercased()
-
-        let matches = try await DatabaseManager.shared.withMetadataDriver(
-            scope: scope,
-            workload: .bulk
-        ) { driver -> [JsonValue] in
-            let tables = MCPConnectionBridge.sortedTables(try await driver.fetchTables(schema: schema))
-            var found: [JsonValue] = []
-            for table in tables where table.name.lowercased().contains(needle) {
-                found.append(.object([
-                    "kind": .string("table"),
-                    "name": .string(table.name),
-                    "object_type": .string(table.type.rawValue),
-                    "schema": table.schema.map(JsonValue.string) ?? JsonValue.null
-                ]))
-                if found.count >= limit { return found }
-            }
-            let allColumns = (try? await driver.fetchAllColumns()) ?? [:]
-            for tableName in allColumns.keys.sorted() {
-                for column in allColumns[tableName] ?? [] where column.name.lowercased().contains(needle) {
-                    found.append(.object([
-                        "kind": .string("column"),
-                        "name": .string(column.name),
-                        "table": .string(tableName),
-                        "data_type": .string(column.dataType)
-                    ]))
-                    if found.count >= limit { return found }
-                }
-            }
-            return found
+    func searchSchema(scope: DatabaseScope, term: String, limit: Int, schemaIsNamed: Bool) async throws -> JsonValue {
+        let databaseType = try await ensureConnected(scope.connectionId)
+        let tableReach = await MainActor.run {
+            MCPSchemaSearch.tableReach(
+                schemaIsNamed: schemaIsNamed,
+                grouping: PluginManager.shared.databaseGroupingStrategy(for: databaseType),
+                systemSchemas: Set(PluginManager.shared.systemSchemaNames(for: databaseType))
+            )
         }
-        return .object([
+        let result = try await MCPSchemaSearch.run(
+            MCPSchemaSearch.Request(scope: scope, term: term, limit: limit, tableReach: tableReach),
+            metadata: DatabaseManager.shared
+        )
+        return Self.encode(search: result, term: term, scope: scope, schemaIsNamed: schemaIsNamed)
+    }
+
+    static func encode(
+        search result: MCPSchemaSearch.Result,
+        term: String,
+        scope: DatabaseScope,
+        schemaIsNamed: Bool
+    ) -> JsonValue {
+        var payload: [String: JsonValue] = [
             "term": .string(term),
-            "matches": .array(matches),
-            "is_truncated": .bool(matches.count >= limit)
-        ])
+            "database": .string(scope.database),
+            "schema": schemaIsNamed ? nullable(scope.schema) : .null,
+            "matches": .array(result.matches.map(encode(match:))),
+            "is_truncated": .bool(result.isTruncated),
+            "unlisted_schemas": .array(result.unlistedSchemas.map(JsonValue.string)),
+            "column_search": .string(result.columnSearch.outcome.rawValue)
+        ]
+        if case .searched(let schema) = result.columnSearch {
+            payload["columns_schema"] = nullable(schema)
+        }
+        return .object(payload)
+    }
+
+    static func encode(match: MCPSchemaSearch.Match) -> JsonValue {
+        switch match {
+        case .table(let name, let schema, let type):
+            return .object([
+                "kind": .string("table"),
+                "name": .string(name),
+                "schema": nullable(schema),
+                "object_type": .string(type.rawValue)
+            ])
+        case .column(let name, let table, let schema, let dataType):
+            return .object([
+                "kind": .string("column"),
+                "name": .string(name),
+                "table": .string(table),
+                "schema": nullable(schema),
+                "data_type": .string(dataType)
+            ])
+        }
+    }
+
+    private static func nullable(_ value: String?) -> JsonValue {
+        value.map(JsonValue.string) ?? .null
     }
 
     func insertRows(

TablePro/Core/MCP/MCPSchemaSearch.swift

@@ -0,0 +1,158 @@
+//
+//  MCPSchemaSearch.swift
+//  TablePro
+//
+
+import Foundation
+import os
+import TableProPluginKit
+
+internal enum MCPSchemaSearch {
+    internal enum TableReach: Equatable, Sendable {
+        case scopeSchema
+        case everySchema(excluding: Set<String>)
+    }
+
+    internal struct Request: Sendable {
+        internal let scope: DatabaseScope
+        internal let term: String
+        internal let limit: Int
+        internal let tableReach: TableReach
+    }
+
+    internal enum Match: Equatable, Sendable {
+        case table(name: String, schema: String?, type: TableInfo.TableType)
+        case column(name: String, table: String, schema: String?, dataType: String)
+    }
+
+    internal enum ColumnSearchOutcome: String, CaseIterable, Sendable {
+        case searched
+        case limitReached = "limit_reached"
+        case failed
+    }
+
+    internal enum ColumnSearch: Equatable, Sendable {
+        case searched(schema: String?)
+        case limitReached
+        case failed
+
+        internal var outcome: ColumnSearchOutcome {
+            switch self {
+            case .searched: .searched
+            case .limitReached: .limitReached
+            case .failed: .failed
+            }
+        }
+    }
+
+    internal struct Result: Equatable, Sendable {
+        internal let matches: [Match]
+        internal let isTruncated: Bool
+        internal let unlistedSchemas: [String]
+        internal let columnSearch: ColumnSearch
+    }
+
+    private struct ColumnRead: Sendable {
+        let schema: String?
+        let matches: [Match]
+    }
+
+    private static let logger = Logger(subsystem: "com.TablePro", category: "MCPSchemaSearch")
+
+    internal static func tableReach(
+        schemaIsNamed: Bool,
+        grouping: GroupingStrategy,
+        systemSchemas: Set<String>
+    ) -> TableReach {
+        guard !schemaIsNamed, DatabaseTreeMetadataService.listsTablesPerSchema(grouping) else {
+            return .scopeSchema
+        }
+        return .everySchema(excluding: systemSchemas)
+    }
+
+    internal static func run(_ request: Request, metadata: ScopedMetadataProviding) async throws -> Result {
+        let needle = request.term.lowercased()
+        let listing = try await tables(reaching: request.tableReach, in: request.scope, metadata: metadata)
+        let tableMatches = ordered(
+            listing.tables.filter { $0.name.lowercased().contains(needle) },
+            preferring: request.scope.schema
+        ).map { table in
+            Match.table(name: table.name, schema: table.schema, type: table.type)
+        }
+        let unlisted = listing.unlistedSchemas.sorted()
+
+        guard tableMatches.count <= request.limit else {
+            return Result(
+                matches: Array(tableMatches.prefix(request.limit)),
+                isTruncated: true,
+                unlistedSchemas: unlisted,
+                columnSearch: .limitReached
+            )
+        }
+
+        let room = request.limit - tableMatches.count
+        let columnRead: ColumnRead
+        do {
+            columnRead = try await columns(matching: needle, in: request.scope, metadata: metadata)
+        } catch is CancellationError {
+            throw CancellationError()
+        } catch let error as DatabaseError {
+            throw error
+        } catch {
+            logger.warning("[search] column read failed error=\(error.publicLogShape, privacy: .public)")
+            return Result(
+                matches: tableMatches,
+                isTruncated: false,
+                unlistedSchemas: unlisted,
+                columnSearch: .failed
+            )
+        }
+        return Result(
+            matches: tableMatches + columnRead.matches.prefix(room),
+            isTruncated: columnRead.matches.count > room,
+            unlistedSchemas: unlisted,
+            columnSearch: .searched(schema: columnRead.schema)
+        )
+    }
+
+    internal static func ordered(_ tables: [TableInfo], preferring schema: String?) -> [TableInfo] {
+        let sorted = MCPConnectionBridge.sortedTables(tables)
+        guard let schema else { return sorted }
+        return sorted.filter { $0.schema == schema } + sorted.filter { $0.schema != schema }
+    }
+
+    private static func tables(
+        reaching reach: TableReach,
+        in scope: DatabaseScope,
+        metadata: ScopedMetadataProviding
+    ) async throws -> CatalogTableListing.Result {
+        switch reach {
+        case .everySchema(let excluded):
+            return try await CatalogTableListing.tables(in: scope, excludingSchemas: excluded, metadata: metadata)
+        case .scopeSchema:
+            let schema = scope.schema
+            let tables = try await metadata.withMetadataDriver(scope: scope, workload: .bulk) { driver in
+                try await driver.fetchTables(schema: schema)
+            }
+            return CatalogTableListing.Result(tables: tables, unlistedSchemas: [])
+        }
+    }
+
+    private static func columns(
+        matching needle: String,
+        in scope: DatabaseScope,
+        metadata: ScopedMetadataProviding
+    ) async throws -> ColumnRead {
+        try await metadata.withMetadataDriver(scope: scope, workload: .bulk) { driver in
+            let schema = (driver as? SchemaSwitchable)?.currentSchema
+            let allColumns = try await driver.fetchAllColumns()
+            var matches: [Match] = []
+            for table in allColumns.keys.sorted() {
+                for column in allColumns[table] ?? [] where column.name.lowercased().contains(needle) {
+                    matches.append(.column(name: column.name, table: table, schema: schema, dataType: column.dataType))
+                }
+            }
+            return ColumnRead(schema: schema, matches: matches)
+        }
+    }
+}

TablePro/Core/MCP/Protocol/Tools/MCPScopeArguments.swift

@@ -17,11 +17,17 @@ enum MCPScopeArguments {
         services: MCPToolServices
     ) async throws -> DatabaseScope {
         let database = try MCPArgumentDecoder.optionalString(arguments, key: "database")
-        let schema = try MCPArgumentDecoder.optionalString(arguments, key: "schema")
         return try await services.connectionBridge.resolveScope(
             connectionId: connectionId,
             database: database,
-            schema: schema
+            schema: try namedSchema(arguments)
         )
     }
+
+    static func namedSchema(_ arguments: JsonValue) throws -> String? {
+        guard let schema = try MCPArgumentDecoder.optionalString(arguments, key: "schema"), !schema.isEmpty else {
+            return nil
+        }
+        return schema
+    }
 }

TablePro/Core/MCP/Protocol/Tools/SchemaObjectTools.swift

@@ -448,8 +448,9 @@ public struct SearchSchemaTool: MCPToolImplementation {
     public static let title: String? = String(localized: "Search Schema")
     public static let description = String(
         localized: """
-        Find tables and columns whose name contains a substring, so a column can be located without \
-        describing every table.
+        Find tables, views and columns whose name contains a substring, and the schema each one is in. \
+        Without 'schema', tables and views are searched in every schema and columns in the current \
+        one; name a schema to search only that schema, columns included.
         """
     )
     public static let requiredScopes: Set<MCPScope> = [.toolsRead]
@@ -471,16 +472,22 @@ public struct SearchSchemaTool: MCPToolImplementation {
                 maximum: 500
             ),
             "database": MCPToolSchema.database,
-            "schema": MCPToolSchema.schema
+            "schema": MCPToolSchema.string(
+                String(localized: "Schema to search, columns included. Omit to search tables and views in every schema.")
+            )
         ],
         required: ["connection_id", "term"]
     )
 
     public static let outputSchema: JsonValue? = MCPToolSchema.object(
         properties: [
             "term": MCPToolSchema.string(String(localized: "Term that was searched")),
+            "database": MCPToolSchema.string(String(localized: "Database that was searched")),
+            "schema": MCPToolSchema.nullableString(
+                String(localized: "Schema the search was narrowed to, null when none was named")
+            ),
             "matches": MCPToolSchema.array(
-                String(localized: "Matching tables first, then matching columns"),
+                String(localized: "Matching tables and views first, those in the current schema leading, then matching columns"),
                 of: MCPToolSchema.object(
                     properties: [
                         "kind": MCPToolSchema.string(
@@ -489,16 +496,31 @@ public struct SearchSchemaTool: MCPToolImplementation {
                         ),
                         "name": MCPToolSchema.string(String(localized: "Matched name")),
                         "table": MCPToolSchema.string(String(localized: "Owning table, for a column match")),
-                        "schema": MCPToolSchema.nullableString(String(localized: "Schema, for a table match")),
-                        "object_type": MCPToolSchema.string(String(localized: "Object type, for a table match")),
+                        "schema": MCPToolSchema.nullableString(
+                            String(localized: "Schema the match is in, null on an engine without schemas")
+                        ),
+                        "object_type": MCPToolSchema.string(
+                            String(localized: "Object type, such as TABLE or VIEW, for a table match")
+                        ),
                         "data_type": MCPToolSchema.string(String(localized: "Column type, for a column match"))
                     ],
-                    required: ["kind", "name"]
+                    required: ["kind", "name", "schema"]
                 )
             ),
-            "is_truncated": MCPToolSchema.boolean(String(localized: "Whether the limit clipped the matches"))
+            "is_truncated": MCPToolSchema.boolean(String(localized: "Whether the limit clipped the matches")),
+            "unlisted_schemas": MCPToolSchema.array(
+                String(localized: "Schemas whose tables could not be listed, so a match in them may be missing"),
+                of: MCPToolSchema.string(String(localized: "Schema name"))
+            ),
+            "column_search": MCPToolSchema.string(
+                String(localized: "Whether columns were searched, or left out because the table matches reached the limit or the column read failed"),
+                enumValues: MCPSchemaSearch.ColumnSearchOutcome.allCases.map(\.rawValue)
+            ),
+            "columns_schema": MCPToolSchema.nullableString(
+                String(localized: "Schema whose columns were searched, when they were")
+            )
         ],
-        required: ["term", "matches", "is_truncated"]
+        required: ["term", "database", "schema", "matches", "is_truncated", "unlisted_schemas", "column_search"]
     )
 
     public init() {}
@@ -514,11 +536,13 @@ public struct SearchSchemaTool: MCPToolImplementation {
         )
         let term = try MCPArgumentDecoder.requireNonEmptyString(arguments, key: "term")
         let limit = try MCPArgumentDecoder.optionalInt(arguments, key: "limit", range: 1...500) ?? 50
+        let namedSchema = try MCPScopeArguments.namedSchema(arguments)
         let scope = try await MCPScopeArguments.resolve(arguments, services: services)
         let payload = try await services.connectionBridge.searchSchema(
             scope: scope,
             term: term,
-            limit: limit
+            limit: limit,
+            schemaIsNamed: namedSchema != nil
         )
         return .structured(payload)
     }

TableProTests/Core/Autocomplete/SQLSchemaProviderTests.swift

@@ -193,8 +193,13 @@ final class MockDatabaseDriver: DatabaseDriver, SchemaSwitchable, @unchecked Sen
         return columnsToReturn[table.lowercased()] ?? []
     }
 
+    var fetchAllColumnsError: Error?
+
     func fetchAllColumns() async throws -> [String: [ColumnInfo]] {
         fetchAllColumnsCallCount += 1
+        if let fetchAllColumnsError {
+            throw fetchAllColumnsError
+        }
         return allColumnsToReturn
     }