feat(plugin-mongodb): edit MongoDB documents as Extended JSON, found by their stored _id and saved under an exact guard

CHANGELOG.md

@@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 
 - **Insert Document…** for MongoDB collections, written as Extended JSON. (#3132)
+- **Edit Document…** for MongoDB documents, written as Extended JSON. (#3132)
 - **Remove Field** for MongoDB cells, and **No Field** for a field a document does not have. (#3132)
 - **Agent** mode: one session with the whole connection window, sessions to start and delete, and what each one ran.
 - Row previews and the query editor sized to the display on iPad and on iPhone Duo's inner display.

Plugins/MongoDBDriverPlugin/MongoDBCapabilities.swift

@@ -19,6 +19,12 @@ struct MongoDBCapabilities: Sendable, Equatable {
         major >= 4
     }
 
+    /// The guard an edit is saved under uses `$convert`, which MongoDB 4.0 added. A server whose
+    /// version is not known is let through, so its own error reaches the user.
+    var supportsDocumentReplaceGuard: Bool {
+        self == .unknown || major >= 4
+    }
+
     /// `$setField` and `$unsetField` arrived in 5.0. Nil when the version is not known, so the
     /// server answers for itself.
     var supportsFieldExpressions: Bool? {

Plugins/MongoDBDriverPlugin/MongoDBConnection+Documents.swift

@@ -9,7 +9,7 @@ import CLibMongoc
 import Foundation
 import TableProPluginKit
 
-/// Whole-document writes for Insert Document.
+/// Whole-document reads and writes for Insert Document and Edit Document.
 ///
 /// These go straight to libmongoc rather than through the shell. The shell turns a document into a
 /// JavaScript object on its way to the server, and JavaScript puts integer-like field names first,
@@ -54,22 +54,202 @@ extension MongoDBConnection {
         throw MongoDBError.libmongocUnavailable
         #endif
     }
+
+    /// The documents whose `_id` is byte for byte the one `identity` names, read from the primary so
+    /// an edit starts from the latest write, once the server and the namespace allow the edit.
+    ///
+    /// Every step, the version probe included, runs in one call on the connection's own queue, so
+    /// none of it blocks the caller's thread and a cancel that lands between steps stops the next.
+    func readStoredDocuments(
+        database: String,
+        collection: String,
+        identity: MongoDocumentIdentity
+    ) async throws -> [MongoStoredDocument] {
+        #if canImport(CLibMongoc)
+        let options = MongoEditableDocument.readOptions(maxTimeMS: effectiveMaxTimeMS(background: false))
+        let namespaceCommand = MongoEditableDocument.namespaceTypeCommand(for: collection)
+        return try await readOnClient { [self] client in
+            try MongoEditableDocument.readStored(
+                serverVersion: serverVersion,
+                listCollectionsReply: {
+                    try runCommandSync(client: client, command: namespaceCommand, database: database).first
+                },
+                storedDocuments: {
+                    try exactMatches(
+                        client: client,
+                        database: database,
+                        collection: collection,
+                        identity: identity,
+                        options: options
+                    )
+                },
+                checkCancelled: checkCancelled
+            )
+        }
+        #else
+        throw MongoDBError.libmongocUnavailable
+        #endif
+    }
+
+    /// Replaces the document `filter` matches and answers how many it matched, which is 0 when the
+    /// stored document is no longer the one the edit started from.
+    ///
+    /// libmongoc refuses an empty field name before sending unless `validate` is off, and the server
+    /// still refuses a top-level `$` field. The guard's answer has to be known, so a collection
+    /// whose write concern asks for no acknowledgement is asked for one here.
+    func replaceDocument(database: String, collection: String, filter: String, replacement: String) async throws -> Int64 {
+        #if canImport(CLibMongoc)
+        try await onClient { [self] client in
+            let filterBson = try parsedBson(filter)
+            defer { bson_destroy(filterBson) }
+            let replacementBson = try parsedBson(replacement)
+            defer { bson_destroy(replacementBson) }
+            guard MongoLibbsonCodec.size(of: filterBson) + MongoLibbsonCodec.size(of: replacementBson)
+                <= MongoEditableDocument.commandSizeLimit else {
+                throw MongoDBDocumentEditingError.tooLarge
+            }
+            let handle = try getCollection(client, database: database, collection: collection)
+            defer { mongoc_collection_destroy(handle) }
+            let isAcknowledged = mongoc_write_concern_is_acknowledged(mongoc_collection_get_write_concern(handle))
+            let optionsBson = try parsedBson(Self.replaceOptions(acknowledged: isAcknowledged))
+            defer { bson_destroy(optionsBson) }
+            guard let reply = bson_new() else { throw MongoDBError.connectionFailed }
+            defer { bson_destroy(reply) }
+            var error = bson_error_t()
+            try checkCancelled()
+            guard mongoc_collection_replace_one(handle, filterBson, replacementBson, optionsBson, reply, &error) else {
+                if let failure = (try? canonicalText(of: reply)).flatMap(MongoWriteFailure.read(fromReply:)) {
+                    throw MongoDBError(code: failure.code, message: failure.message)
+                }
+                throw makeError(error)
+            }
+            return try matchedCount(in: reply)
+        }
+        #else
+        throw MongoDBError.libmongocUnavailable
+        #endif
+    }
+
+    private static func replaceOptions(acknowledged: Bool) -> String {
+        let base = #""collation":{"locale":"simple"},"validate":false"#
+        return acknowledged ? "{\(base)}" : #"{\#(base),"writeConcern":{"w":1}}"#
+    }
 }
 
 #if canImport(CLibMongoc)
+/// libbson's own reading of Extended JSON, which is the only authority on what a text stores.
+struct MongoLibbsonCodec: MongoDocumentCodec {
+    func isSameDocument(_ text: String, asCanonical canonical: String) -> Bool {
+        guard let stored = Self.parse(canonical) else { return false }
+        defer { bson_destroy(stored) }
+        return Self.reads(text, as: stored)
+    }
+
+    func bsonSize(of json: String) -> Int? {
+        guard let bson = Self.parse(json) else { return nil }
+        defer { bson_destroy(bson) }
+        return Self.size(of: bson)
+    }
+
+    /// Whether `text` reads back as exactly `stored`. Text libbson cannot read is not the same.
+    static func reads(_ text: String, as stored: OpaquePointer) -> Bool {
+        guard let reread = parse(text) else { return false }
+        defer { bson_destroy(reread) }
+        return bson_equal(reread, stored)
+    }
+
+    /// Every BSON document opens with its own length as a little-endian int32.
+    static func size(of bson: OpaquePointer) -> Int {
+        guard let data = bson_get_data(bson) else { return 0 }
+        let length = UnsafeRawPointer(data).loadUnaligned(as: Int32.self)
+        return Int(Int32(littleEndian: length))
+    }
+
+    private static func parse(_ json: String) -> OpaquePointer? {
+        var error = bson_error_t()
+        return json.withCString { bson_new_from_json($0, -1, &error) }
+    }
+}
+
 fileprivate extension MongoDBConnection {
-    /// On a thread of its own rather than the cooperative pool, since `withClientSync` blocks until
-    /// the connection's queue runs the call.
+    func matchedCount(in reply: OpaquePointer) throws -> Int64 {
+        let text = try canonicalText(of: reply)
+        guard let data = text.data(using: .utf8),
+              let object = try? JSONSerialization.jsonObject(with: data) as? [String: Any],
+              let matched = MongoScriptJson.numeric(object["matchedCount"]) else {
+            throw MongoDBError(code: 0, message: MongoScriptText.writeRefused(code: 0))
+        }
+        return matched
+    }
+
     func onClient<T: Sendable>(_ body: @escaping @Sendable (OpaquePointer) throws -> T) async throws -> T {
         beginScriptRun()
-        return try await pluginDispatchAsync(on: .global(qos: .userInitiated)) { [self] in
+        return try await onClientQueue(body)
+    }
+
+    /// A read the app may cancel. Clearing the latch as the read starts also clears a cancel the app
+    /// sent a moment earlier, so the task is asked after the clear: a task is marked cancelled before
+    /// the app's cancel reaches the latch.
+    func readOnClient<T: Sendable>(_ body: @escaping @Sendable (OpaquePointer) throws -> T) async throws -> T {
+        beginScriptRun()
+        try Task.checkCancellation()
+        return try await onClientQueue(body)
+    }
+
+    /// On a thread of its own rather than the cooperative pool, since `withClientSync` blocks until
+    /// the connection's queue runs the call.
+    private func onClientQueue<T: Sendable>(
+        _ body: @escaping @Sendable (OpaquePointer) throws -> T
+    ) async throws -> T {
+        try await pluginDispatchAsync(on: .global(qos: .userInitiated)) { [self] in
             try withClientSync { client in
                 try checkCancelled()
                 return try body(client)
             }
         }
     }
 
+    /// Read under the collection's own collation, which is what lets a string `_id` use its index.
+    /// A lenient match is read past rather than kept, and the read stops at the second exact one.
+    func exactMatches(
+        client: OpaquePointer,
+        database: String,
+        collection: String,
+        identity: MongoDocumentIdentity,
+        options: String
+    ) throws -> [MongoStoredDocument] {
+        let filterBson = try parsedBson(identity.filter)
+        defer { bson_destroy(filterBson) }
+        let optionsBson = try parsedBson(options)
+        defer { bson_destroy(optionsBson) }
+        let session = attachCancellableSession(client: client, opts: optionsBson)
+        defer {
+            if let session {
+                releaseSessionLsid()
+                mongoc_client_session_destroy(session)
+            }
+        }
+        let handle = try getCollection(client, database: database, collection: collection)
+        defer { mongoc_collection_destroy(handle) }
+        let primary = mongoc_read_prefs_new(MONGOC_READ_PRIMARY)
+        defer { mongoc_read_prefs_destroy(primary) }
+        guard let cursor = mongoc_collection_find_with_opts(handle, filterBson, optionsBson, primary) else {
+            throw MongoDBError(code: 0, message: MongoScriptText.cursorFailed)
+        }
+        defer { mongoc_cursor_destroy(cursor) }
+        var matches = MongoExactMatches(identity: identity)
+        var pointer: OpaquePointer?
+        while !matches.isDecided, mongoc_cursor_next(cursor, &pointer) {
+            try checkCancelled()
+            guard let stored = pointer else { continue }
+            let canonical = try canonicalText(of: stored)
+            matches.consider(canonical) { MongoLibbsonCodec.reads(canonical, as: stored) }
+        }
+        var error = bson_error_t()
+        if mongoc_cursor_error(cursor, &error) { throw makeError(error) }
+        return matches.documents
+    }
+
     func parsedBson(_ json: String) throws -> OpaquePointer {
         var error = bson_error_t()
         guard let bson = json.withCString({ bson_new_from_json($0, -1, &error) }) else {

Plugins/MongoDBDriverPlugin/MongoDBConnection+ScriptHelpers.swift

@@ -12,6 +12,10 @@ struct MongoScriptDocumentBatch: Sendable {
     var json: [String]
     var isTruncated: Bool
 
+    /// Whether these are whole documents as stored, which is what makes each one's `_id` a locator
+    /// the grid can hand back to edit it. A projection or a pipeline builds documents of its own.
+    var holdsStoredDocuments = false
+
     static let empty = MongoScriptDocumentBatch(json: [], isTruncated: false)
 
     var jsonArray: String { "[\(json.joined(separator: ","))]" }

Plugins/MongoDBDriverPlugin/MongoDBDocumentEditingError.swift

@@ -0,0 +1,67 @@
+import Foundation
+
+/// Why a document cannot be opened for editing or saved, in words for the user.
+enum MongoDBDocumentEditingError: Error, Equatable, LocalizedError {
+    case unsupportedOperation
+    case unknownDocument
+    case documentChanged
+    case ambiguousIdentity
+    case identityChanged
+    case missingIdentity
+    case inexactAsText
+    case emptyTimestamp(String)
+    case tooLarge
+    case serverTooOld
+    case view
+    case timeSeries
+    case notACollection(String)
+
+    var errorDescription: String? {
+        switch self {
+        case .unsupportedOperation:
+            return String(localized: "MongoDB cannot make this change to a document.")
+        case .unknownDocument:
+            return String(localized: "This row does not name a stored document.")
+        case .documentChanged:
+            return String(
+                localized: "The document changed on the server after it was opened, so nothing was saved. Copy your text, then open it again."
+            )
+        case .ambiguousIdentity:
+            return String(localized: "More than one document has this _id, so it cannot be edited here.")
+        case .identityChanged:
+            return String(
+                localized: "The _id cannot change. Put the original _id back, or insert a new document instead."
+            )
+        case .missingIdentity:
+            return String(localized: "This document has no _id, so it cannot be edited here.")
+        case .inexactAsText:
+            return String(
+                localized: "This document holds a value text cannot write back exactly, such as a subdocument like {\"$numberInt\": \"5\"}."
+            )
+        case .emptyTimestamp(let field):
+            return String(
+                format: String(
+                    localized: "The top-level field \u{201C}%@\u{201D} holds Timestamp(0, 0), which MongoDB replaces with the current time on every save."
+                ),
+                field
+            )
+        case .tooLarge:
+            return String(
+                localized: "This document is too large to edit here. Saving it would send more than 16 MB to the server."
+            )
+        case .serverTooOld:
+            return String(localized: "Editing a document needs MongoDB 4.0 or later.")
+        case .view:
+            return String(localized: "This is a view, so its documents cannot be edited. Edit them in the collection the view reads.")
+        case .timeSeries:
+            return String(
+                localized: "This is a time-series collection, which cannot replace a single document, so its documents cannot be edited here."
+            )
+        case .notACollection(let type):
+            return String(
+                format: String(localized: "MongoDB lists this as a \u{201C}%@\u{201D} rather than a collection, so its documents cannot be edited here."),
+                type
+            )
+        }
+    }
+}

Plugins/MongoDBDriverPlugin/MongoDBPluginDriver+Documents.swift

@@ -8,16 +8,39 @@ import TableProPluginKit
 
 extension MongoDBPluginDriver {
     func documentWriteStatement(_ write: PluginDocumentWrite) throws -> String? {
-        try documentWritePlan(write).statement
+        try documentWritePlan(write)?.statement
     }
 
+    /// The plan is built on a queue of its own, since an edit reads and compares the whole document
+    /// and the caller may be the main actor.
     func executeDocumentWrite(_ write: PluginDocumentWrite) async throws {
         guard let conn = mongoConnection else { throw MongoDBPluginError.notConnected }
-        let plan = try documentWritePlan(write)
-        try await conn.insertDocument(database: currentDb, collection: write.table, document: plan.document)
+        let plan = try await pluginDispatchAsync(on: .global(qos: .userInitiated)) { [self] in
+            try documentWritePlan(write)
+        }
+        guard let plan else { return }
+        switch plan.write {
+        case .insert(let document):
+            try await conn.insertDocument(database: currentDb, collection: write.table, document: document)
+        case .replace(let filter, let replacement):
+            let matched = try await conn.replaceDocument(
+                database: currentDb, collection: write.table, filter: filter, replacement: replacement
+            )
+            guard matched > 0 else { throw MongoDBDocumentEditingError.documentChanged }
+        }
     }
 
-    private func documentWritePlan(_ write: PluginDocumentWrite) throws -> MongoDocumentWritePlan {
+    func fetchDocument(table: String, schema: String?, locator: String) async throws -> String? {
+        guard let conn = mongoConnection else { throw MongoDBPluginError.notConnected }
+        let identity = try MongoDocumentIdentity(locator: locator)
+        let stored = try await conn.readStoredDocuments(database: currentDb, collection: table, identity: identity)
+        try Task.checkCancellation()
+        return try await pluginDispatchAsync(on: .global(qos: .userInitiated)) {
+            try MongoEditableDocument.text(for: identity, among: stored, codec: MongoLibbsonCodec())
+        }
+    }
+
+    private func documentWritePlan(_ write: PluginDocumentWrite) throws -> MongoDocumentWritePlan? {
         guard let conn = mongoConnection else { throw MongoDBPluginError.notConnected }
         return try MongoDocumentWritePlan.make(
             collection: write.table,

Plugins/MongoDBDriverPlugin/MongoDBPluginDriver.swift

@@ -266,17 +266,7 @@ final class MongoDBPluginDriver: PluginDatabaseDriver, @unchecked Sendable {
         guard let rowCap, MongoDBFindLimitPolicy.isTruncated(rowCount: result.rows.count, rowCap: rowCap) else {
             return result
         }
-        var capped = PluginQueryResult(
-            columns: result.columns,
-            columnTypeNames: result.columnTypeNames,
-            rows: Array(result.rows.prefix(rowCap)),
-            rowsAffected: result.rowsAffected,
-            executionTime: result.executionTime,
-            isTruncated: true,
-            statusMessage: result.statusMessage
-        )
-        capped.absentCells = result.absentCells?.filter { $0.key < rowCap }
-        return capped
+        return result.capped(to: rowCap)
     }
 
     private func mapExecutionError(_ error: Error) -> Error {