feat(plugin-mongodb): nested field paths, write options, and MQL formatting

CHANGELOG.md

@@ -18,6 +18,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Autocomplete for MongoDB queries. Typing `db.` lists collections, `db.users.` lists the driver methods, and inside a query you get field names plus the operators that are valid in that spot: query operators in a filter, update operators in an update, stage names in a pipeline, and expression operators inside a stage. (#2095)
 - PostgreSQL autocomplete now knows the operators, including `::`, the JSON ones (`->`, `->>`, `#>`, `@>`, `?`, `?|`, `?&`), array and range containment, regex matching and full-text search. Each one shows what it does and which types it works on. Typing `::` offers the type names. (#2095)
 - PostgreSQL autocomplete covers about 400 built-in functions and the multi-word syntax people actually type, such as `ON CONFLICT DO UPDATE SET`, `GENERATED ALWAYS AS IDENTITY` and window frame clauses. (#2095)
+- MongoDB field suggestions include nested paths. A document with `address: { city }` now suggests `address.city`, not just `address`. (#2095)
+- MongoDB updates accept an options argument, so `db.users.updateOne({...}, {...}, {upsert: true})` upserts instead of ignoring the option. `arrayFilters` and `hint` are passed through too. (#2095)
+- Format Query follows the editor language. On a MongoDB tab it lays out filters and pipelines by nesting depth instead of running the SQL formatter over them. (#2095)
 
 ### Added
 

Plugins/MongoDBDriverPlugin/BsonDocumentFlattener.swift

@@ -259,6 +259,59 @@ struct BsonDocumentFlattener {
         return value > 0 ? "Infinity" : "-Infinity"
     }
 
+    // MARK: - Nested Field Paths
+
+    /// Dotted paths across the sampled documents, including paths inside nested objects and
+    /// inside the objects an array holds. This is deliberately separate from `unionColumns`,
+    /// which stays flat because the grid renders a nested object as one JSON column.
+    static func fieldPaths(
+        from documents: [[String: Any]],
+        representation: MongoDBUuidRepresentation,
+        maxDepth: Int = 4
+    ) -> [PluginFieldPath] {
+        var kinds: [String: [BsonValueKind: Int]] = [:]
+        var depths: [String: Int] = [:]
+        var order: [String] = []
+
+        func visit(_ document: [String: Any], prefix: String, depth: Int) {
+            guard depth <= maxDepth else { return }
+
+            for key in document.keys.sorted() {
+                guard let value = document[key], !(value is NSNull) else { continue }
+                let path = prefix.isEmpty ? key : "\(prefix).\(key)"
+
+                if depths[path] == nil {
+                    depths[path] = depth
+                    order.append(path)
+                }
+                kinds[path, default: [:]][valueKind(for: value, representation: representation), default: 0] += 1
+
+                if let nested = value as? [String: Any] {
+                    visit(nested, prefix: path, depth: depth + 1)
+                } else if let array = value as? [Any] {
+                    for element in array.prefix(20) {
+                        guard let nested = element as? [String: Any] else { continue }
+                        visit(nested, prefix: path, depth: depth + 1)
+                    }
+                }
+            }
+        }
+
+        for document in documents {
+            visit(document, prefix: "", depth: 1)
+        }
+
+        return order.compactMap { path in
+            guard let winner = kinds[path]?.max(by: { $0.value < $1.value })?.key,
+                  let depth = depths[path] else { return nil }
+            return PluginFieldPath(
+                path: path,
+                typeName: typeName(for: winner, representation: representation),
+                depth: depth
+            )
+        }
+    }
+
     // MARK: - Type Inference
 
     private static func inferValueKind(

Plugins/MongoDBDriverPlugin/MongoDBPluginDriver.swift

@@ -215,6 +215,61 @@ final class MongoDBPluginDriver: PluginDatabaseDriver, @unchecked Sendable {
             .map { PluginTableInfo(name: $0, type: "table", rowCount: nil) }
     }
 
+    private func executeWrite(
+        kind: MongoWriteKind,
+        collection: String,
+        filter: String,
+        document: String,
+        options: MongoWriteOptions,
+        conn: MongoDBConnection,
+        db: String,
+        startTime: Date
+    ) async throws -> PluginQueryResult {
+        var fields: [String] = []
+        if options.upsert { fields.append("\"upsert\": true") }
+        if let arrayFilters = options.arrayFilters { fields.append("\"arrayFilters\": \(arrayFilters)") }
+        if let hint = options.hint { fields.append("\"hint\": \(hint)") }
+        let extras = fields.isEmpty ? "" : ", " + fields.joined(separator: ", ")
+
+        if kind == .findOneAndUpdate {
+            let command = """
+                {"findAndModify": "\(escapeJsonString(collection))", "query": \(filter), \
+                "update": \(document), "new": true\(extras)}
+                """
+            let docs = try await conn.runCommand(command, database: db)
+            return buildPluginResult(from: docs.isEmpty ? [] : [docs[0]], startTime: startTime)
+        }
+
+        let command = """
+            {"update": "\(escapeJsonString(collection))", \
+            "updates": [{"q": \(filter), "u": \(document), "multi": \(kind == .updateMany)\(extras)}]}
+            """
+        let result = try await conn.runCommand(command, database: db)
+        let modified = (result.first?["nModified"] as? Int64)
+            ?? (result.first?["nModified"] as? Int).map(Int64.init) ?? 0
+        let upserted = (result.first?["upserted"] as? [Any])?.count ?? 0
+        let affected = Int(modified) + upserted
+
+        return PluginQueryResult(
+            columns: ["modifiedCount", "upsertedCount"], columnTypeNames: ["Int64", "Int64"],
+            rows: [[.text(String(modified)), .text(String(upserted))]], rowsAffected: affected,
+            executionTime: Date().timeIntervalSince(startTime)
+        )
+    }
+
+    func sampleFieldPaths(table: String, schema: String?, limit: Int) async throws -> [PluginFieldPath] {
+        guard let conn = mongoConnection else {
+            throw MongoDBPluginError.notConnected
+        }
+
+        let docs = try await conn.find(
+            database: currentDb, collection: table,
+            filter: "{}", sort: nil, projection: nil, skip: 0, limit: max(1, limit)
+        ).docs
+
+        return BsonDocumentFlattener.fieldPaths(from: docs, representation: uuidRepresentation)
+    }
+
     func fetchColumns(table: String, schema: String?) async throws -> [PluginColumnInfo] {
         guard let conn = mongoConnection else {
             throw MongoDBPluginError.notConnected
@@ -637,6 +692,14 @@ final class MongoDBPluginDriver: PluginDatabaseDriver, @unchecked Sendable {
             let cmd = "\"findAndModify\": \"\(escapeJsonString(collection))\", \"query\": \(filter), \"update\": \(update)"
             return "db.runCommand({\"explain\": {\(cmd)}, \"verbosity\": \"executionStats\"})"
 
+        case .write(let kind, let collection, let filter, let document, _):
+            let multi = kind == .updateMany
+            if kind == .findOneAndUpdate {
+                let cmd = "\"findAndModify\": \"\(escapeJsonString(collection))\", \"query\": \(filter), \"update\": \(document)"
+                return "db.runCommand({\"explain\": {\(cmd)}, \"verbosity\": \"executionStats\"})"
+            }
+            return "db.runCommand({\"explain\": {\"update\": \"\(escapeJsonString(collection))\", \"updates\": [{\"q\": \(filter), \"u\": \(document), \"multi\": \(multi)}]}, \"verbosity\": \"executionStats\"})"
+
         case .findOneAndReplace(let collection, let filter, let replacement):
             let cmd = "\"findAndModify\": \"\(escapeJsonString(collection))\", \"query\": \(filter), \"update\": \(replacement)"
             return "db.runCommand({\"explain\": {\(cmd)}, \"verbosity\": \"executionStats\"})"
@@ -923,6 +986,12 @@ final class MongoDBPluginDriver: PluginDatabaseDriver, @unchecked Sendable {
             let docs = try await conn.runCommand(cmd, database: db)
             return buildPluginResult(from: docs.isEmpty ? [] : [docs[0]], startTime: startTime)
 
+        case .write(let kind, let collection, let filter, let document, let options):
+            return try await executeWrite(
+                kind: kind, collection: collection, filter: filter, document: document,
+                options: options, conn: conn, db: db, startTime: startTime
+            )
+
         case .findOneAndReplace(let collection, let filter, let replacement):
             let cmd = "{\"findAndModify\": \"\(escapeJsonString(collection))\", \"query\": \(filter), \"update\": \(replacement), \"new\": true}"
             let docs = try await conn.runCommand(cmd, database: db)

Plugins/TableProPluginKit/MongoShellParser.swift

@@ -32,6 +32,37 @@ public enum MongoOperation {
     case listCollections
     case listDatabases
     case ping
+    case write(
+        kind: MongoWriteKind,
+        collection: String,
+        filter: String,
+        document: String,
+        options: MongoWriteOptions
+    )
+}
+
+/// Which write a `.write` operation performs. Only produced when the shell call carries an
+/// options argument; the plain two-argument forms keep their original operation cases.
+public enum MongoWriteKind: String, Sendable {
+    case updateOne
+    case updateMany
+    case replaceOne
+    case findOneAndUpdate
+}
+
+/// The third argument of an update, replace or findOneAndUpdate call.
+public struct MongoWriteOptions: Sendable {
+    public var upsert: Bool
+    public var arrayFilters: String?
+    public var hint: String?
+    public var raw: String
+
+    public init(upsert: Bool = false, arrayFilters: String? = nil, hint: String? = nil, raw: String = "{}") {
+        self.upsert = upsert
+        self.arrayFilters = arrayFilters
+        self.hint = hint
+        self.raw = raw
+    }
 }
 
 /// Options for a find operation parsed from chained methods
@@ -297,16 +328,16 @@ public struct MongoShellParser {
             operation = .insertMany(collection: collection, documents: arg)
 
         case "updateOne":
-            let (filter, update) = try parseTwoArgs(arg, method: "updateOne")
-            operation = .updateOne(collection: collection, filter: filter, update: update)
+            operation = makeWrite(.updateOne, collection: collection,
+                                  args: try parseWriteArgs(arg, method: "updateOne"))
 
         case "updateMany":
-            let (filter, update) = try parseTwoArgs(arg, method: "updateMany")
-            operation = .updateMany(collection: collection, filter: filter, update: update)
+            operation = makeWrite(.updateMany, collection: collection,
+                                  args: try parseWriteArgs(arg, method: "updateMany"))
 
         case "replaceOne":
-            let (filter, replacement) = try parseTwoArgs(arg, method: "replaceOne")
-            operation = .replaceOne(collection: collection, filter: filter, replacement: replacement)
+            operation = makeWrite(.replaceOne, collection: collection,
+                                  args: try parseWriteArgs(arg, method: "replaceOne"))
 
         case "deleteOne":
             let filter = arg.isEmpty ? "{}" : arg
@@ -324,8 +355,8 @@ public struct MongoShellParser {
             operation = .dropIndex(collection: collection, indexName: arg)
 
         case "findOneAndUpdate":
-            let (filter, update) = try parseTwoArgs(arg, method: "findOneAndUpdate")
-            operation = .findOneAndUpdate(collection: collection, filter: filter, update: update)
+            operation = makeWrite(.findOneAndUpdate, collection: collection,
+                                  args: try parseWriteArgs(arg, method: "findOneAndUpdate"))
 
         case "findOneAndReplace":
             let (filter, replacement) = try parseTwoArgs(arg, method: "findOneAndReplace")
@@ -587,6 +618,78 @@ public struct MongoShellParser {
         return (parts[0], parts[1])
     }
 
+    /// Parse a write call's arguments, keeping the optional third options document.
+    private static func parseWriteArgs(
+        _ args: String,
+        method: String
+    ) throws -> (filter: String, document: String, options: MongoWriteOptions?) {
+        let parts = try splitTopLevelArgs(args)
+        guard parts.count >= 2 else {
+            throw MongoShellParseError.missingArgument("\(method) requires 2 arguments")
+        }
+        guard parts.count > 2 else {
+            return (parts[0], parts[1], nil)
+        }
+        return (parts[0], parts[1], try parseWriteOptions(parts[2], method: method))
+    }
+
+    private static func parseWriteOptions(_ raw: String, method: String) throws -> MongoWriteOptions {
+        let trimmed = raw.trimmingCharacters(in: .whitespacesAndNewlines)
+        guard trimmed.hasPrefix("{"), trimmed.hasSuffix("}") else {
+            throw MongoShellParseError.invalidSyntax("\(method) options must be a document")
+        }
+
+        var options = MongoWriteOptions(raw: trimmed)
+        let inner = String(trimmed.dropFirst().dropLast())
+
+        for entry in try splitTopLevelArgs(inner) {
+            let pair = entry.split(separator: ":", maxSplits: 1).map {
+                $0.trimmingCharacters(in: CharacterSet(charactersIn: " \t\n\"'`"))
+            }
+            guard pair.count == 2 else { continue }
+
+            switch pair[0] {
+            case "upsert":
+                options.upsert = pair[1] == "true"
+            case "arrayFilters":
+                options.arrayFilters = pair[1]
+            case "hint":
+                options.hint = pair[1]
+            default:
+                continue
+            }
+        }
+
+        return options
+    }
+
+    private static func makeWrite(
+        _ kind: MongoWriteKind,
+        collection: String,
+        args: (filter: String, document: String, options: MongoWriteOptions?)
+    ) -> MongoOperation {
+        guard let options = args.options else {
+            switch kind {
+            case .updateOne:
+                return .updateOne(collection: collection, filter: args.filter, update: args.document)
+            case .updateMany:
+                return .updateMany(collection: collection, filter: args.filter, update: args.document)
+            case .replaceOne:
+                return .replaceOne(collection: collection, filter: args.filter, replacement: args.document)
+            case .findOneAndUpdate:
+                return .findOneAndUpdate(collection: collection, filter: args.filter, update: args.document)
+            }
+        }
+
+        return .write(
+            kind: kind,
+            collection: collection,
+            filter: args.filter,
+            document: args.document,
+            options: options
+        )
+    }
+
     /// Parse two arguments where the second is optional
     private static func parseTwoArgsOptional(_ args: String) throws -> (String, String?) {
         let parts = try splitTopLevelArgs(args)

Plugins/TableProPluginKit/PluginDatabaseDriver.swift

@@ -115,6 +115,7 @@ public protocol PluginDatabaseDriver: AnyObject, Sendable {
 
     func fetchApproximateRowCount(table: String, schema: String?) async throws -> Int?
     func fetchAllColumns(schema: String?) async throws -> [String: [PluginColumnInfo]]
+    func sampleFieldPaths(table: String, schema: String?, limit: Int) async throws -> [PluginFieldPath]
     func fetchAllForeignKeys(schema: String?) async throws -> [String: [PluginForeignKeyInfo]]
     func fetchAllDatabaseMetadata() async throws -> [PluginDatabaseMetadata]
     func fetchDependentTypes(table: String, schema: String?) async throws -> [(name: String, labels: [String])]
@@ -286,6 +287,12 @@ public extension PluginDatabaseDriver {
         return result
     }
 
+    /// Default: no nested field paths. Document stores override this to sample documents and
+    /// report the dotted paths their nested structure exposes, which a flat column list cannot.
+    func sampleFieldPaths(table: String, schema: String?, limit: Int) async throws -> [PluginFieldPath] {
+        []
+    }
+
     /// Default: fetches foreign keys per-table sequentially (N+1 round-trips).
     /// SQL drivers should override with a single bulk query (e.g. INFORMATION_SCHEMA.KEY_COLUMN_USAGE).
     func fetchAllForeignKeys(schema: String?) async throws -> [String: [PluginForeignKeyInfo]] {

Plugins/TableProPluginKit/PluginFieldPath.swift

@@ -0,0 +1,16 @@
+import Foundation
+
+/// A field a document store exposes at a nested path, reported for query authoring.
+/// A flat column list cannot express these: a document store types a nested object as one
+/// opaque column, so `address.city` never appears among its columns.
+public struct PluginFieldPath: Codable, Sendable, Hashable {
+    public let path: String
+    public let typeName: String
+    public let depth: Int
+
+    public init(path: String, typeName: String, depth: Int) {
+        self.path = path
+        self.typeName = typeName
+        self.depth = depth
+    }
+}

TablePro/Core/Autocomplete/Mongo/MongoCompletionService.swift

@@ -139,7 +139,22 @@ final class MongoCompletionService: QueryCompletionService {
 
     private func fieldItems(for collection: String?) async -> [SQLCompletionItem] {
         guard let schemaProvider, let collection else { return [] }
-        return await schemaProvider.columnCompletionItems(for: collection)
+
+        let paths = await schemaProvider.fieldPaths(for: collection)
+        guard !paths.isEmpty else {
+            return await schemaProvider.columnCompletionItems(for: collection)
+        }
+
+        return paths.map { path in
+            var item = SQLCompletionItem(
+                label: path.path,
+                kind: .column,
+                insertText: path.path,
+                detail: path.typeName
+            )
+            item.sortPriority = SQLCompletionKind.column.basePriority + path.depth
+            return item
+        }
     }
 
     private func fieldPathItems(for collection: String?) async -> [SQLCompletionItem] {