diff --git a/agents/common/src/main/java/com/dbx/agent/AgentProtocol.java b/agents/common/src/main/java/com/dbx/agent/AgentProtocol.java index ff48c016ff..d3fceaf6da 100644 --- a/agents/common/src/main/java/com/dbx/agent/AgentProtocol.java +++ b/agents/common/src/main/java/com/dbx/agent/AgentProtocol.java @@ -70,6 +70,7 @@ public final class AgentProtocol { public static final String MONGO_METHOD_INSERT_DOCUMENTS = "insert_documents"; public static final String MONGO_METHOD_UPDATE_DOCUMENT = "update_document"; public static final String MONGO_METHOD_UPDATE_DOCUMENTS = "update_documents"; + public static final String MONGO_METHOD_REPLACE_DOCUMENT = "replace_document"; public static final String MONGO_METHOD_DELETE_DOCUMENT = "delete_document"; public static final String MONGO_METHOD_DELETE_DOCUMENTS = "delete_documents"; public static final String MONGO_METHOD_RUN_COMMAND = "run_command"; @@ -127,6 +128,7 @@ public final class AgentProtocol { public static final String CAPABILITY_MONGO_CLONE_COLLECTION = "mongo_clone_collection"; public static final String CAPABILITY_MONGO_RUN_COMMAND = "mongo_run_command"; public static final String CAPABILITY_MONGO_INSERT_DOCUMENTS = "mongo_insert_documents"; + public static final String CAPABILITY_MONGO_REPLACE_DOCUMENT = "mongo_replace_document"; public static final String CAPABILITY_MULTI_SESSION = "multi_session"; public static final String CAPABILITY_STRUCTURED_ERROR_V1 = "structured_error_v1"; @@ -162,7 +164,8 @@ public final class AgentProtocol { CAPABILITY_MONGO_DROP_DATABASE, CAPABILITY_MONGO_CLONE_COLLECTION, CAPABILITY_MONGO_RUN_COMMAND, - CAPABILITY_MONGO_INSERT_DOCUMENTS + CAPABILITY_MONGO_INSERT_DOCUMENTS, + CAPABILITY_MONGO_REPLACE_DOCUMENT )); public static final List MULTI_SESSION_CAPABILITIES; @@ -226,6 +229,7 @@ public final class AgentProtocol { mongoCapabilities.add(CAPABILITY_MONGO_CLONE_COLLECTION); mongoCapabilities.add(CAPABILITY_MONGO_RUN_COMMAND); mongoCapabilities.add(CAPABILITY_MONGO_INSERT_DOCUMENTS); + mongoCapabilities.add(CAPABILITY_MONGO_REPLACE_DOCUMENT); MONGO_LEGACY_CAPABILITIES = Collections.unmodifiableList(mongoCapabilities); List mongoMultiSessionCapabilities = new java.util.ArrayList<>(MULTI_SESSION_CAPABILITIES); @@ -233,6 +237,7 @@ public final class AgentProtocol { mongoMultiSessionCapabilities.add(CAPABILITY_MONGO_CLONE_COLLECTION); mongoMultiSessionCapabilities.add(CAPABILITY_MONGO_RUN_COMMAND); mongoMultiSessionCapabilities.add(CAPABILITY_MONGO_INSERT_DOCUMENTS); + mongoMultiSessionCapabilities.add(CAPABILITY_MONGO_REPLACE_DOCUMENT); MONGO_LEGACY_MULTI_SESSION_CAPABILITIES = Collections.unmodifiableList(mongoMultiSessionCapabilities); List jdbcCapabilities = new java.util.ArrayList<>(MULTI_SESSION_CAPABILITIES); @@ -274,6 +279,7 @@ public final class AgentProtocol { MONGO_METHOD_INSERT_DOCUMENTS, MONGO_METHOD_UPDATE_DOCUMENT, MONGO_METHOD_UPDATE_DOCUMENTS, + MONGO_METHOD_REPLACE_DOCUMENT, MONGO_METHOD_DELETE_DOCUMENT, MONGO_METHOD_DELETE_DOCUMENTS, MONGO_METHOD_RUN_COMMAND diff --git a/agents/common/src/main/resources/agent-protocol-v1.json b/agents/common/src/main/resources/agent-protocol-v1.json index 250bed89cf..716e9eb035 100644 --- a/agents/common/src/main/resources/agent-protocol-v1.json +++ b/agents/common/src/main/resources/agent-protocol-v1.json @@ -2,7 +2,7 @@ "protocolVersion": 1, "handshakeMethod": "handshake", "handshakeResponseFields": ["protocolVersion", "agentProtocolVersion", "capabilities"], - "allCapabilities": ["connect", "test_connection", "metadata", "query", "paged_query", "transaction", "ddl", "kv", "kv_ttl", "kv_cas", "kv_list_values", "kv_status", "kv_history", "etcd_compaction", "etcd_defrag", "etcd_watch", "etcd_lease", "etcd_auth", "mongo_drop_database", "mongo_clone_collection", "mongo_run_command", "mongo_insert_documents"], + "allCapabilities": ["connect", "test_connection", "metadata", "query", "paged_query", "transaction", "ddl", "kv", "kv_ttl", "kv_cas", "kv_list_values", "kv_status", "kv_history", "etcd_compaction", "etcd_defrag", "etcd_watch", "etcd_lease", "etcd_auth", "mongo_drop_database", "mongo_clone_collection", "mongo_run_command", "mongo_insert_documents", "mongo_replace_document"], "capabilities": ["connect", "test_connection", "metadata", "query", "paged_query", "transaction", "ddl"], "defaultSqlCapabilities": ["connect", "test_connection", "metadata", "query", "paged_query", "transaction", "ddl"], "commonMethods": [ @@ -42,6 +42,6 @@ "disconnect", "shutdown" ], - "mongoLegacyMethods": ["list_databases", "list_collections", "find_documents", "find_one", "explain_find", "aggregate_documents", "find_documents_extended_json", "count_documents", "server_version", "create_index", "create_user", "drop_indexes", "drop_collection", "clone_collection", "drop_database", "insert_document", "insert_documents", "update_document", "update_documents", "delete_document", "delete_documents", "run_command"], + "mongoLegacyMethods": ["list_databases", "list_collections", "find_documents", "find_one", "explain_find", "aggregate_documents", "find_documents_extended_json", "count_documents", "server_version", "create_index", "create_user", "drop_indexes", "drop_collection", "clone_collection", "drop_database", "insert_document", "insert_documents", "update_document", "update_documents", "replace_document", "delete_document", "delete_documents", "run_command"], "kvMethods": ["kv_list_prefix", "kv_get", "kv_put", "kv_delete", "kv_rename", "kv_history", "kv_status", "etcd_compact", "etcd_defrag", "etcd_watch_start", "etcd_watch_poll", "etcd_watch_stop", "etcd_lease_list", "etcd_lease_get", "etcd_lease_grant", "etcd_lease_keepalive_once", "etcd_lease_revoke", "etcd_auth_user_list", "etcd_auth_user_get", "etcd_auth_user_add", "etcd_auth_user_delete", "etcd_auth_user_change_password", "etcd_auth_user_grant_role", "etcd_auth_user_revoke_role", "etcd_auth_role_list", "etcd_auth_role_get", "etcd_auth_role_add", "etcd_auth_role_delete", "etcd_auth_role_grant_permission", "etcd_auth_role_revoke_permission"] } diff --git a/agents/common/src/main/resources/agent-protocol-v2.json b/agents/common/src/main/resources/agent-protocol-v2.json index 55b86252ee..1b1a2017a9 100644 --- a/agents/common/src/main/resources/agent-protocol-v2.json +++ b/agents/common/src/main/resources/agent-protocol-v2.json @@ -2,7 +2,7 @@ "protocolVersion": 2, "handshakeMethod": "handshake", "handshakeResponseFields": ["protocolVersion", "agentProtocolVersion", "capabilities"], - "allCapabilities": ["connect", "test_connection", "metadata", "query", "paged_query", "transaction", "ddl", "kv", "kv_ttl", "kv_cas", "kv_list_values", "kv_status", "kv_history", "etcd_compaction", "etcd_defrag", "etcd_watch", "etcd_lease", "etcd_auth", "mongo_drop_database", "mongo_clone_collection", "mongo_run_command", "mongo_insert_documents", "multi_session", "structured_error_v1"], + "allCapabilities": ["connect", "test_connection", "metadata", "query", "paged_query", "transaction", "ddl", "kv", "kv_ttl", "kv_cas", "kv_list_values", "kv_status", "kv_history", "etcd_compaction", "etcd_defrag", "etcd_watch", "etcd_lease", "etcd_auth", "mongo_drop_database", "mongo_clone_collection", "mongo_run_command", "mongo_insert_documents", "mongo_replace_document", "multi_session", "structured_error_v1"], "capabilities": ["connect", "test_connection", "metadata", "query", "paged_query", "transaction", "ddl", "multi_session"], "defaultSqlCapabilities": ["connect", "test_connection", "metadata", "query", "paged_query", "transaction", "ddl", "multi_session", "structured_error_v1"], "commonMethods": [ @@ -46,7 +46,7 @@ "disconnect", "shutdown" ], - "mongoLegacyMethods": ["list_databases", "list_collections", "find_documents", "find_one", "explain_find", "aggregate_documents", "find_documents_extended_json", "count_documents", "server_version", "create_index", "create_user", "drop_indexes", "drop_collection", "clone_collection", "drop_database", "insert_document", "insert_documents", "update_document", "update_documents", "delete_document", "delete_documents", "run_command"], + "mongoLegacyMethods": ["list_databases", "list_collections", "find_documents", "find_one", "explain_find", "aggregate_documents", "find_documents_extended_json", "count_documents", "server_version", "create_index", "create_user", "drop_indexes", "drop_collection", "clone_collection", "drop_database", "insert_document", "insert_documents", "update_document", "update_documents", "replace_document", "delete_document", "delete_documents", "run_command"], "kvMethods": ["kv_list_prefix", "kv_get", "kv_put", "kv_delete", "kv_rename", "kv_history", "kv_status", "etcd_compact", "etcd_defrag", "etcd_watch_start", "etcd_watch_poll", "etcd_watch_stop", "etcd_lease_list", "etcd_lease_get", "etcd_lease_grant", "etcd_lease_keepalive_once", "etcd_lease_revoke", "etcd_auth_user_list", "etcd_auth_user_get", "etcd_auth_user_add", "etcd_auth_user_delete", "etcd_auth_user_change_password", "etcd_auth_user_grant_role", "etcd_auth_user_revoke_role", "etcd_auth_role_list", "etcd_auth_role_get", "etcd_auth_role_add", "etcd_auth_role_delete", "etcd_auth_role_grant_permission", "etcd_auth_role_revoke_permission"], "sessionField": "agentSessionId", "cursorSessionField": "sessionId" diff --git a/agents/drivers/mongodb/src/main/java/com/dbx/agent/mongodb/MongoAgent.java b/agents/drivers/mongodb/src/main/java/com/dbx/agent/mongodb/MongoAgent.java index 59065265c3..4f380c33ad 100644 --- a/agents/drivers/mongodb/src/main/java/com/dbx/agent/mongodb/MongoAgent.java +++ b/agents/drivers/mongodb/src/main/java/com/dbx/agent/mongodb/MongoAgent.java @@ -26,6 +26,7 @@ import com.mongodb.client.model.CollationMaxVariable; import com.mongodb.client.model.CollationStrength; import com.mongodb.client.model.CountOptions; import com.mongodb.client.model.InsertManyOptions; +import com.mongodb.client.model.ReplaceOptions; import com.mongodb.client.model.UpdateOptions; import com.mongodb.client.result.UpdateResult; import java.io.BufferedReader; @@ -1682,6 +1683,56 @@ public final class MongoAgent { return true; } + private static Object replaceDocument(JsonObject params) { + MongoClient c = requireClient(); + String database = params.get("database").getAsString(); + String collection = params.get("collection").getAsString(); + String filterJson = params.get("filter_json").getAsString(); + String replacementJson = params.get("replacement_json").getAsString(); + String optionsJson = params.has("options_json") && !params.get("options_json").isJsonNull() + ? params.get("options_json").getAsString() + : null; + + var col = c.getDatabase(database).getCollection(collection); + Document filter = documentForWrite(filterJson); + Document replacement = documentForWrite(replacementJson); + requireReplacementDocument(replacement); + ReplaceOptions options = replaceOptionsForWrite(optionsJson); + var result = col.replaceOne(filter, replacement, options); + return Collections.singletonMap("modified_count", result.getModifiedCount()); + } + + /** replaceOne swaps the whole document; update operators here mean updateOne was intended. */ + static void requireReplacementDocument(Document doc) { + for (String key : doc.keySet()) { + if (key.startsWith("$")) { + throw new IllegalArgumentException( + "Replacement document must not contain update operators such as " + key); + } + } + } + + static ReplaceOptions replaceOptionsForWrite(String optionsJson) { + ReplaceOptions result = new ReplaceOptions(); + if (optionsJson == null || optionsJson.trim().isEmpty()) { + return result; + } + Document options = Document.parse(optionsJson); + for (String key : options.keySet()) { + if (!"upsert".equals(key)) { + throw new IllegalArgumentException("Unsupported replace option: " + key); + } + } + Object rawUpsert = options.get("upsert"); + if (rawUpsert != null) { + if (!(rawUpsert instanceof Boolean)) { + throw new IllegalArgumentException("upsert must be a boolean"); + } + result.upsert((Boolean) rawUpsert); + } + return result; + } + static void requireBulkUpdateOperatorDocument(Document doc) { if (!isUpdateOperatorDocument(doc)) { // updateOne/updateMany are shell-style bulk updates here; replacements stay on the @@ -1865,6 +1916,7 @@ public final class MongoAgent { case AgentProtocol.MONGO_METHOD_INSERT_DOCUMENTS -> insertDocuments(params); case AgentProtocol.MONGO_METHOD_UPDATE_DOCUMENT -> updateDocument(params); case AgentProtocol.MONGO_METHOD_UPDATE_DOCUMENTS -> updateDocuments(params); + case AgentProtocol.MONGO_METHOD_REPLACE_DOCUMENT -> replaceDocument(params); case AgentProtocol.MONGO_METHOD_DELETE_DOCUMENT -> deleteDocument(params); case AgentProtocol.MONGO_METHOD_DELETE_DOCUMENTS -> deleteDocuments(params); case AgentProtocol.MONGO_METHOD_RUN_COMMAND -> runCommand(params); diff --git a/agents/drivers/mongodb/src/test/java/com/dbx/agent/mongodb/MongoAgentTest.java b/agents/drivers/mongodb/src/test/java/com/dbx/agent/mongodb/MongoAgentTest.java index db5b9a4789..8b287bd2e5 100644 --- a/agents/drivers/mongodb/src/test/java/com/dbx/agent/mongodb/MongoAgentTest.java +++ b/agents/drivers/mongodb/src/test/java/com/dbx/agent/mongodb/MongoAgentTest.java @@ -1007,6 +1007,50 @@ class MongoAgentTest { assertEquals("upsert must be a boolean", error.getMessage()); } + @Test + void replaceDocumentMethodIsRecognizedOverJsonRpc() { + String response = MongoAgent.handleRequest( + "{\"jsonrpc\":\"2.0\",\"id\":11,\"method\":\"replace_document\"," + + "\"params\":{\"database\":\"app\",\"collection\":\"orders\",\"filter_json\":\"{\\\"_id\\\":1}\"," + + "\"replacement_json\":\"{\\\"name\\\":\\\"new\\\"}\"}}"); + + JsonObject json = JsonParser.parseString(response).getAsJsonObject(); + assertEquals(11, json.get("id").getAsInt()); + assertEquals("Not connected", json.getAsJsonObject("error").get("message").getAsString()); + assertTrue(AgentProtocol.MONGO_LEGACY_METHODS.contains(AgentProtocol.MONGO_METHOD_REPLACE_DOCUMENT)); + assertTrue(AgentProtocol.MONGO_LEGACY_CAPABILITIES.contains(AgentProtocol.CAPABILITY_MONGO_REPLACE_DOCUMENT)); + } + + @Test + void parsesReplaceOptions() { + assertTrue(MongoAgent.replaceOptionsForWrite("{\"upsert\":true}").isUpsert()); + assertFalse(MongoAgent.replaceOptionsForWrite("{\"upsert\":false}").isUpsert()); + assertFalse(MongoAgent.replaceOptionsForWrite(null).isUpsert()); + assertFalse(MongoAgent.replaceOptionsForWrite("{}").isUpsert()); + + IllegalArgumentException unsupported = assertThrows( + IllegalArgumentException.class, + () -> MongoAgent.replaceOptionsForWrite("{\"arrayFilters\":[]}") + ); + assertEquals("Unsupported replace option: arrayFilters", unsupported.getMessage()); + IllegalArgumentException notBoolean = assertThrows( + IllegalArgumentException.class, + () -> MongoAgent.replaceOptionsForWrite("{\"upsert\":\"yes\"}") + ); + assertEquals("upsert must be a boolean", notBoolean.getMessage()); + } + + @Test + void rejectsReplacementDocumentsWithUpdateOperators() { + IllegalArgumentException error = assertThrows( + IllegalArgumentException.class, + () -> MongoAgent.requireReplacementDocument(Document.parse("{\"$set\":{\"a\":1}}")) + ); + assertEquals("Replacement document must not contain update operators such as $set", error.getMessage()); + // A plain document, including one with nested `$`-keys inside values, is fine. + MongoAgent.requireReplacementDocument(Document.parse("{\"name\":\"new\",\"meta\":{\"$ref\":\"x\"}}")); + } + @Test void rejectsUnsupportedUpdateOptions() { IllegalArgumentException error = assertThrows( diff --git a/apps/desktop/src/lib/backend/api.ts b/apps/desktop/src/lib/backend/api.ts index dad796f5eb..c41137aaa2 100644 --- a/apps/desktop/src/lib/backend/api.ts +++ b/apps/desktop/src/lib/backend/api.ts @@ -815,6 +815,7 @@ export const mongoInsertDocuments = forward("mongoInsertDocuments"); export const documentUpdateDocument = forward("documentUpdateDocument"); export const mongoUpdateDocument = forward("mongoUpdateDocument"); export const mongoUpdateDocuments = forward("mongoUpdateDocuments"); +export const mongoReplaceDocument = forward("mongoReplaceDocument"); export const documentDeleteDocument = forward("documentDeleteDocument"); export const documentSaveMeilisearchBatch = forward("documentSaveMeilisearchBatch"); export const meilisearchSearchDocuments = forward("meilisearchSearchDocuments"); diff --git a/apps/desktop/src/lib/backend/http.ts b/apps/desktop/src/lib/backend/http.ts index 337a2d208e..18f73b4018 100644 --- a/apps/desktop/src/lib/backend/http.ts +++ b/apps/desktop/src/lib/backend/http.ts @@ -4530,6 +4530,17 @@ export async function documentUpdateDocument(connectionId: string, database: str }); } +export async function mongoReplaceDocument(connectionId: string, database: string, collection: string, filterJson: string, replacementJson: string, optionsJson?: string): Promise<{ affected_rows: number }> { + return post("/api/mongo/replace-document", { + connectionId, + database, + collection, + filterJson, + replacementJson, + optionsJson, + }); +} + export async function mongoUpdateDocuments(connectionId: string, database: string, collection: string, filterJson: string, updateJson: string, many: boolean, optionsJson?: string): Promise<{ affected_rows: number }> { return post("/api/mongo/update-documents", { connectionId, diff --git a/apps/desktop/src/lib/backend/tauri.ts b/apps/desktop/src/lib/backend/tauri.ts index ac427c2340..7f0cf7c61b 100644 --- a/apps/desktop/src/lib/backend/tauri.ts +++ b/apps/desktop/src/lib/backend/tauri.ts @@ -4440,6 +4440,18 @@ export async function documentUpdateDocument(connectionId: string, database: str }); } +export async function mongoReplaceDocument(connectionId: string, database: string, collection: string, filterJson: string, replacementJson: string, optionsJson?: string): Promise<{ affected_rows: number }> { + const affectedRows = await invoke("mongo_replace_document", { + connectionId, + database, + collection, + filterJson, + replacementJson, + optionsJson, + }); + return { affected_rows: affectedRows }; +} + export async function mongoUpdateDocuments(connectionId: string, database: string, collection: string, filterJson: string, updateJson: string, many: boolean, optionsJson?: string): Promise<{ affected_rows: number }> { const affectedRows = await invoke("mongo_update_documents", { connectionId, diff --git a/apps/desktop/src/lib/mongo/mongoCompletion.ts b/apps/desktop/src/lib/mongo/mongoCompletion.ts index 5eed4c3861..bb2d77de37 100644 --- a/apps/desktop/src/lib/mongo/mongoCompletion.ts +++ b/apps/desktop/src/lib/mongo/mongoCompletion.ts @@ -93,6 +93,7 @@ const COLLECTION_METHODS = [ { label: "insertMany", detail: "Insert multiple documents", apply: "insertMany([{}])" }, { label: "updateOne", detail: "Update one matching document", apply: "updateOne({}, { $set: {} })" }, { label: "updateMany", detail: "Update all matching documents", apply: "updateMany({}, { $set: {} })" }, + { label: "replaceOne", detail: "Replace one matching document", apply: "replaceOne({}, {})" }, { label: "deleteOne", detail: "Delete one matching document", apply: "deleteOne({})" }, { label: "deleteMany", detail: "Delete all matching documents", apply: "deleteMany({})" }, { label: "findOneAndUpdate", detail: "Atomically update and return a document", apply: "findOneAndUpdate({}, { $set: {} })" }, @@ -120,6 +121,7 @@ const COLLECTION_METHOD_BOOST: Record<(typeof COLLECTION_METHODS)[number]["label insertMany: 170, updateOne: 160, updateMany: 150, + replaceOne: 145, deleteOne: 140, deleteMany: 130, findOneAndUpdate: 120, @@ -190,6 +192,7 @@ const METHOD_ARG_ROLES: Record = { findOneAndDelete: ["filter", "options"], updateOne: ["filter", "update", "options"], updateMany: ["filter", "update", "options"], + replaceOne: ["filter", "replacement", "options"], findOneAndUpdate: ["filter", "update", "options"], findOneAndReplace: ["filter", "replacement", "options"], insertOne: ["document", "options"], diff --git a/apps/desktop/src/lib/mongo/mongoShellCommand.ts b/apps/desktop/src/lib/mongo/mongoShellCommand.ts index 8db704effc..c758151d7f 100644 --- a/apps/desktop/src/lib/mongo/mongoShellCommand.ts +++ b/apps/desktop/src/lib/mongo/mongoShellCommand.ts @@ -48,6 +48,7 @@ const COLLECTION_METHOD_SHAPES: Record = { update: { expects: "a filter, an update, and optional options", roles: ["filter", "update", "options"] }, updateOne: { expects: "a filter, an update, and optional options", roles: ["filter", "update", "options"] }, updateMany: { expects: "a filter, an update, and optional options", roles: ["filter", "update", "options"] }, + replaceOne: { expects: "a filter, a replacement document, and optional options", roles: ["filter", "replacement", "options"] }, deleteOne: { expects: "a filter", roles: ["filter"] }, deleteMany: { expects: "a filter", roles: ["filter"] }, findOneAndUpdate: { expects: "a filter, an update, and optional options", roles: ["filter", "update", "options"] }, @@ -76,6 +77,7 @@ const SUPPORTED_COLLECTION_METHODS = [ "insertMany", "updateOne", "updateMany", + "replaceOne", "deleteOne", "deleteMany", "findOneAndUpdate", @@ -157,6 +159,11 @@ function diagnoseMongoCommand(source: string): string | null { const tail = source.slice(closeIndex + 1).trim(); if (tail) return `Unexpected text after ${method}(...): "${tail.length > 40 ? `${tail.slice(0, 40)}…` : tail}".`; + if (method === "replaceOne" && args[1]) { + const replacement = parseMongoObjectArgument(args[1]); + const operator = replacement ? Object.keys(JSON.parse(replacement) as Record).find((key) => key.startsWith("$")) : undefined; + if (operator) return `replaceOne() replaces the whole document, so it must not contain update operators such as ${operator}; use updateOne() to modify fields.`; + } return `${method}() expects ${shapeSpec.expects}.`; } @@ -256,7 +263,7 @@ export interface MongoDistinctCommand { filter?: string; } -type MongoWriteKind = "runCommand" | "insert" | "update" | "delete" | "createIndex" | "createUser" | "dropIndex" | "dropIndexes" | "dropCollection" | "findOneAndUpdate" | "findOneAndReplace" | "findOneAndDelete"; +type MongoWriteKind = "runCommand" | "insert" | "update" | "replace" | "delete" | "createIndex" | "createUser" | "dropIndex" | "dropIndexes" | "dropCollection" | "findOneAndUpdate" | "findOneAndReplace" | "findOneAndDelete"; export type MongoCommand = | ({ kind: "find" } & MongoFindCommand) @@ -273,6 +280,7 @@ export type MongoCommand = | ({ kind: "runCommand" } & MongoRunCommand) | { kind: "insert"; collection: string; docsJson: string } | { kind: "update"; collection: string; filter: string; update: string; options?: string; many: boolean } + | { kind: "replace"; collection: string; filter: string; replacement: string; options?: string } | { kind: "delete"; collection: string; filter: string; many: boolean } | { kind: "createIndex"; collection: string; keys: string; options?: string } | { kind: "dropIndex"; collection: string; index: string } @@ -744,6 +752,20 @@ export function parseMongoWriteCommand(input: string): MongoWriteCommand | null return value !== null && typeof value === "object" ? { kind: "insert", collection: insert.collection, docsJson: docs } : null; } + const replaceOne = parseCollectionMethodTarget(source, "replaceOne"); + if (replaceOne) { + const args = parseMethodArgs(source, replaceOne.methodCallIndex); + if (!args || args.length < 2 || args.length > 3) return null; + const filter = normalizeJsonArgument(args[0]); + const replacement = parseMongoObjectArgument(args[1]); + if (!filter || !replacement) return null; + // A replacement is a whole document; `{$set: ...}` here means updateOne() was intended. + if (Object.keys(JSON.parse(replacement) as Record).some((key) => key.startsWith("$"))) return null; + const options = args[2]?.trim() ? normalizeJsonArgument(args[2]) : undefined; + if (args[2]?.trim() && !options) return null; + return { kind: "replace", collection: replaceOne.collection, filter, replacement, ...(options ? { options } : {}) }; + } + for (const method of ["updateOne", "updateMany"] as const) { const target = parseCollectionMethodTarget(source, method); if (!target) continue; @@ -1516,6 +1538,7 @@ function isNonEmptyRecord(value: unknown): value is Record { function mongoWriteFilter(command: MongoWriteCommand): string | null { switch (command.kind) { case "update": + case "replace": case "delete": case "findOneAndUpdate": case "findOneAndReplace": diff --git a/apps/desktop/src/stores/queryStore.ts b/apps/desktop/src/stores/queryStore.ts index 8c03b1c01a..ec0bd18bd2 100644 --- a/apps/desktop/src/stores/queryStore.ts +++ b/apps/desktop/src/stores/queryStore.ts @@ -6642,6 +6642,7 @@ export const useQueryStore = defineStore("query", () => { } case "insert": case "update": + case "replace": case "delete": case "createIndex": case "createUser": @@ -6665,6 +6666,9 @@ export const useQueryStore = defineStore("query", () => { } else if (mongoCommand.kind === "update") { const result = await api.mongoUpdateDocuments(executionConnectionId, currentDatabase, mongoCommand.collection, mongoCommand.filter, mongoCommand.update, mongoCommand.many, mongoCommand.options); allResults.push(markQueryResultRowsRaw(annotateMongoResult(mongoWriteToQueryResult(result.affected_rows, performance.now() - commandStartedAt)))); + } else if (mongoCommand.kind === "replace") { + const result = await api.mongoReplaceDocument(executionConnectionId, currentDatabase, mongoCommand.collection, mongoCommand.filter, mongoCommand.replacement, mongoCommand.options); + allResults.push(markQueryResultRowsRaw(annotateMongoResult(mongoWriteToQueryResult(result.affected_rows, performance.now() - commandStartedAt)))); } else if (mongoCommand.kind === "createIndex") { const result = await api.mongoCreateIndex(executionConnectionId, currentDatabase, mongoCommand.collection, mongoCommand.keys, mongoCommand.options); allResults.push(markQueryResultRowsRaw(annotateMongoResult(mongoCreateIndexToQueryResult(result.name, performance.now() - commandStartedAt)))); diff --git a/crates/dbx-core/assets/agent-protocol-v1.json b/crates/dbx-core/assets/agent-protocol-v1.json index 250bed89cf..716e9eb035 100644 --- a/crates/dbx-core/assets/agent-protocol-v1.json +++ b/crates/dbx-core/assets/agent-protocol-v1.json @@ -2,7 +2,7 @@ "protocolVersion": 1, "handshakeMethod": "handshake", "handshakeResponseFields": ["protocolVersion", "agentProtocolVersion", "capabilities"], - "allCapabilities": ["connect", "test_connection", "metadata", "query", "paged_query", "transaction", "ddl", "kv", "kv_ttl", "kv_cas", "kv_list_values", "kv_status", "kv_history", "etcd_compaction", "etcd_defrag", "etcd_watch", "etcd_lease", "etcd_auth", "mongo_drop_database", "mongo_clone_collection", "mongo_run_command", "mongo_insert_documents"], + "allCapabilities": ["connect", "test_connection", "metadata", "query", "paged_query", "transaction", "ddl", "kv", "kv_ttl", "kv_cas", "kv_list_values", "kv_status", "kv_history", "etcd_compaction", "etcd_defrag", "etcd_watch", "etcd_lease", "etcd_auth", "mongo_drop_database", "mongo_clone_collection", "mongo_run_command", "mongo_insert_documents", "mongo_replace_document"], "capabilities": ["connect", "test_connection", "metadata", "query", "paged_query", "transaction", "ddl"], "defaultSqlCapabilities": ["connect", "test_connection", "metadata", "query", "paged_query", "transaction", "ddl"], "commonMethods": [ @@ -42,6 +42,6 @@ "disconnect", "shutdown" ], - "mongoLegacyMethods": ["list_databases", "list_collections", "find_documents", "find_one", "explain_find", "aggregate_documents", "find_documents_extended_json", "count_documents", "server_version", "create_index", "create_user", "drop_indexes", "drop_collection", "clone_collection", "drop_database", "insert_document", "insert_documents", "update_document", "update_documents", "delete_document", "delete_documents", "run_command"], + "mongoLegacyMethods": ["list_databases", "list_collections", "find_documents", "find_one", "explain_find", "aggregate_documents", "find_documents_extended_json", "count_documents", "server_version", "create_index", "create_user", "drop_indexes", "drop_collection", "clone_collection", "drop_database", "insert_document", "insert_documents", "update_document", "update_documents", "replace_document", "delete_document", "delete_documents", "run_command"], "kvMethods": ["kv_list_prefix", "kv_get", "kv_put", "kv_delete", "kv_rename", "kv_history", "kv_status", "etcd_compact", "etcd_defrag", "etcd_watch_start", "etcd_watch_poll", "etcd_watch_stop", "etcd_lease_list", "etcd_lease_get", "etcd_lease_grant", "etcd_lease_keepalive_once", "etcd_lease_revoke", "etcd_auth_user_list", "etcd_auth_user_get", "etcd_auth_user_add", "etcd_auth_user_delete", "etcd_auth_user_change_password", "etcd_auth_user_grant_role", "etcd_auth_user_revoke_role", "etcd_auth_role_list", "etcd_auth_role_get", "etcd_auth_role_add", "etcd_auth_role_delete", "etcd_auth_role_grant_permission", "etcd_auth_role_revoke_permission"] } diff --git a/crates/dbx-core/assets/agent-protocol-v2.json b/crates/dbx-core/assets/agent-protocol-v2.json index 55b86252ee..1b1a2017a9 100644 --- a/crates/dbx-core/assets/agent-protocol-v2.json +++ b/crates/dbx-core/assets/agent-protocol-v2.json @@ -2,7 +2,7 @@ "protocolVersion": 2, "handshakeMethod": "handshake", "handshakeResponseFields": ["protocolVersion", "agentProtocolVersion", "capabilities"], - "allCapabilities": ["connect", "test_connection", "metadata", "query", "paged_query", "transaction", "ddl", "kv", "kv_ttl", "kv_cas", "kv_list_values", "kv_status", "kv_history", "etcd_compaction", "etcd_defrag", "etcd_watch", "etcd_lease", "etcd_auth", "mongo_drop_database", "mongo_clone_collection", "mongo_run_command", "mongo_insert_documents", "multi_session", "structured_error_v1"], + "allCapabilities": ["connect", "test_connection", "metadata", "query", "paged_query", "transaction", "ddl", "kv", "kv_ttl", "kv_cas", "kv_list_values", "kv_status", "kv_history", "etcd_compaction", "etcd_defrag", "etcd_watch", "etcd_lease", "etcd_auth", "mongo_drop_database", "mongo_clone_collection", "mongo_run_command", "mongo_insert_documents", "mongo_replace_document", "multi_session", "structured_error_v1"], "capabilities": ["connect", "test_connection", "metadata", "query", "paged_query", "transaction", "ddl", "multi_session"], "defaultSqlCapabilities": ["connect", "test_connection", "metadata", "query", "paged_query", "transaction", "ddl", "multi_session", "structured_error_v1"], "commonMethods": [ @@ -46,7 +46,7 @@ "disconnect", "shutdown" ], - "mongoLegacyMethods": ["list_databases", "list_collections", "find_documents", "find_one", "explain_find", "aggregate_documents", "find_documents_extended_json", "count_documents", "server_version", "create_index", "create_user", "drop_indexes", "drop_collection", "clone_collection", "drop_database", "insert_document", "insert_documents", "update_document", "update_documents", "delete_document", "delete_documents", "run_command"], + "mongoLegacyMethods": ["list_databases", "list_collections", "find_documents", "find_one", "explain_find", "aggregate_documents", "find_documents_extended_json", "count_documents", "server_version", "create_index", "create_user", "drop_indexes", "drop_collection", "clone_collection", "drop_database", "insert_document", "insert_documents", "update_document", "update_documents", "replace_document", "delete_document", "delete_documents", "run_command"], "kvMethods": ["kv_list_prefix", "kv_get", "kv_put", "kv_delete", "kv_rename", "kv_history", "kv_status", "etcd_compact", "etcd_defrag", "etcd_watch_start", "etcd_watch_poll", "etcd_watch_stop", "etcd_lease_list", "etcd_lease_get", "etcd_lease_grant", "etcd_lease_keepalive_once", "etcd_lease_revoke", "etcd_auth_user_list", "etcd_auth_user_get", "etcd_auth_user_add", "etcd_auth_user_delete", "etcd_auth_user_change_password", "etcd_auth_user_grant_role", "etcd_auth_user_revoke_role", "etcd_auth_role_list", "etcd_auth_role_get", "etcd_auth_role_add", "etcd_auth_role_delete", "etcd_auth_role_grant_permission", "etcd_auth_role_revoke_permission"], "sessionField": "agentSessionId", "cursorSessionField": "sessionId" diff --git a/crates/dbx-core/src/db/agent_driver.rs b/crates/dbx-core/src/db/agent_driver.rs index 462833fb63..2b6309fbf1 100644 --- a/crates/dbx-core/src/db/agent_driver.rs +++ b/crates/dbx-core/src/db/agent_driver.rs @@ -1041,6 +1041,7 @@ pub enum AgentCapability { MongoCloneCollection, MongoRunCommand, MongoInsertDocuments, + MongoReplaceDocument, MultiSession, StructuredErrorV1, } @@ -1123,7 +1124,7 @@ fn parse_agent_rpc_error_header(header: &str) -> (Option, String) { } impl AgentCapability { - pub const ALL: [Self; 24] = [ + pub const ALL: [Self; 25] = [ Self::Connect, Self::TestConnection, Self::Metadata, @@ -1146,6 +1147,7 @@ impl AgentCapability { Self::MongoCloneCollection, Self::MongoRunCommand, Self::MongoInsertDocuments, + Self::MongoReplaceDocument, Self::MultiSession, Self::StructuredErrorV1, ]; @@ -1174,6 +1176,7 @@ impl AgentCapability { Self::MongoCloneCollection => "mongo_clone_collection", Self::MongoRunCommand => "mongo_run_command", Self::MongoInsertDocuments => "mongo_insert_documents", + Self::MongoReplaceDocument => "mongo_replace_document", Self::MultiSession => "multi_session", Self::StructuredErrorV1 => "structured_error_v1", } @@ -1376,13 +1379,14 @@ pub enum MongoAgentMethod { InsertDocuments, UpdateDocument, UpdateDocuments, + ReplaceDocument, DeleteDocument, DeleteDocuments, RunCommand, } impl MongoAgentMethod { - pub const ALL: [Self; 22] = [ + pub const ALL: [Self; 23] = [ Self::ListDatabases, Self::ListCollections, Self::FindDocuments, @@ -1402,6 +1406,7 @@ impl MongoAgentMethod { Self::InsertDocuments, Self::UpdateDocument, Self::UpdateDocuments, + Self::ReplaceDocument, Self::DeleteDocument, Self::DeleteDocuments, Self::RunCommand, @@ -1428,6 +1433,7 @@ impl MongoAgentMethod { Self::InsertDocuments => "insert_documents", Self::UpdateDocument => "update_document", Self::UpdateDocuments => "update_documents", + Self::ReplaceDocument => "replace_document", Self::DeleteDocument => "delete_document", Self::DeleteDocuments => "delete_documents", Self::RunCommand => "run_command", @@ -2803,6 +2809,13 @@ impl AgentDriverClient { self.call_mongo_method(MongoAgentMethod::UpdateDocuments, params).await } + pub async fn mongo_replace_document( + &mut self, + params: Value, + ) -> Result { + self.call_mongo_method(MongoAgentMethod::ReplaceDocument, params).await + } + pub async fn mongo_delete_document( &mut self, params: Value, @@ -2958,6 +2971,7 @@ pub fn agent_supports_capability(handshake: Option<&AgentHandshake>, capability: | AgentCapability::MongoCloneCollection | AgentCapability::MongoRunCommand | AgentCapability::MongoInsertDocuments + | AgentCapability::MongoReplaceDocument ) { return handshake.map(|value| value.supports(capability)).unwrap_or(false); } @@ -4991,9 +5005,10 @@ for line in sys.stdin: assert_eq!(AgentCapability::MongoCloneCollection.as_str(), "mongo_clone_collection"); assert_eq!(AgentCapability::MongoRunCommand.as_str(), "mongo_run_command"); assert_eq!(AgentCapability::MongoInsertDocuments.as_str(), "mongo_insert_documents"); + assert_eq!(AgentCapability::MongoReplaceDocument.as_str(), "mongo_replace_document"); assert_eq!(AgentCapability::MultiSession.as_str(), "multi_session"); assert_eq!(AgentCapability::StructuredErrorV1.as_str(), "structured_error_v1"); - assert_eq!(AgentCapability::ALL.len(), 24); + assert_eq!(AgentCapability::ALL.len(), 25); } #[test] @@ -5055,6 +5070,7 @@ for line in sys.stdin: assert_eq!(MongoAgentMethod::InsertDocuments.as_str(), "insert_documents"); assert_eq!(MongoAgentMethod::UpdateDocument.as_str(), "update_document"); assert_eq!(MongoAgentMethod::UpdateDocuments.as_str(), "update_documents"); + assert_eq!(MongoAgentMethod::ReplaceDocument.as_str(), "replace_document"); assert_eq!(MongoAgentMethod::DeleteDocument.as_str(), "delete_document"); assert_eq!(MongoAgentMethod::DeleteDocuments.as_str(), "delete_documents"); } diff --git a/crates/dbx-core/src/db/mongo_driver.rs b/crates/dbx-core/src/db/mongo_driver.rs index ef6d2a6c71..bac813b5e1 100644 --- a/crates/dbx-core/src/db/mongo_driver.rs +++ b/crates/dbx-core/src/db/mongo_driver.rs @@ -2120,6 +2120,50 @@ pub async fn update_documents( Ok(result.modified_count) } +/// replaceOne(filter, replacement[, { upsert }]): the whole document is swapped, so +/// unlike update there are no array filters to apply. +pub async fn replace_document( + client: &Client, + database: &str, + collection: &str, + filter_json: &str, + replacement_json: &str, + options_json: Option<&str>, +) -> Result { + let filter_value: serde_json::Value = + serde_json::from_str(filter_json).map_err(|e| format!("Invalid filter JSON: {e}"))?; + let replacement_value: serde_json::Value = + serde_json::from_str(replacement_json).map_err(|e| format!("Invalid replacement JSON: {e}"))?; + let filter = json_filter_to_document(&filter_value).map_err(|e| format!("Invalid filter: {e}"))?; + let replacement = json_object_to_document(&replacement_value).map_err(|e| format!("Invalid replacement: {e}"))?; + if let Some(operator) = replacement.keys().find(|key| key.starts_with('$')) { + return Err(format!("Replacement document must not contain update operators such as {operator}")); + } + let upsert = parse_replace_options(options_json)?; + let col = client.database(database).collection::(collection); + let mut action = col.replace_one(filter, replacement); + if let Some(upsert) = upsert { + action = action.upsert(upsert); + } + let result = action.await.map_err(|e| e.to_string())?; + Ok(result.modified_count) +} + +#[derive(Default, Deserialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +struct MongoReplaceOptions { + upsert: Option, +} + +fn parse_replace_options(options_json: Option<&str>) -> Result, String> { + let Some(raw) = options_json.filter(|value| !value.trim().is_empty()) else { + return Ok(None); + }; + let options: MongoReplaceOptions = + serde_json::from_str(raw).map_err(|e| format!("Invalid replace options: {e}"))?; + Ok(options.upsert) +} + #[derive(Default, Deserialize)] #[serde(rename_all = "camelCase")] struct MongoUpdateOptions { diff --git a/crates/dbx-core/src/mongo_ops.rs b/crates/dbx-core/src/mongo_ops.rs index a266483ee6..0b9c34adc4 100644 --- a/crates/dbx-core/src/mongo_ops.rs +++ b/crates/dbx-core/src/mongo_ops.rs @@ -804,6 +804,45 @@ pub async fn mongo_update_documents_core( } } +pub async fn mongo_replace_document_core( + state: &AppState, + connection_id: &str, + database: &str, + collection: &str, + filter_json: &str, + replacement_json: &str, + options_json: Option<&str>, +) -> Result { + ensure_document_pool(state, connection_id).await?; + let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; + match &pool { + PoolKind::MongoDb(client) => { + mongo_driver::replace_document(client, database, collection, filter_json, replacement_json, options_json) + .await + } + PoolKind::Agent(client) => { + let mut client = client.lock().await; + if !client.supports_capability(AgentCapability::MongoReplaceDocument) { + return Err( + "MongoDB Legacy Agent does not support replaceOne; upgrade or reinstall the MongoDB Legacy driver" + .to_string(), + ); + } + let result: serde_json::Value = client + .mongo_replace_document(serde_json::json!({ + "database": database, + "collection": collection, + "filter_json": filter_json, + "replacement_json": replacement_json, + "options_json": options_json, + })) + .await?; + Ok(result.get("modified_count").and_then(|v| v.as_u64()).unwrap_or(0)) + } + _ => Err("Not a MongoDB connection".to_string()), + } +} + pub async fn mongo_delete_document_core( state: &AppState, connection_id: &str, @@ -1031,6 +1070,19 @@ pub async fn execute_mongo_command_core( let affected = mongo_insert_documents_core(state, connection_id, database, collection, documents).await?; Ok(affected_query_result(affected)) } + MongoCommand::Replace { collection, filter, replacement, options } => { + let affected = mongo_replace_document_core( + state, + connection_id, + database, + collection, + filter, + replacement, + options.as_deref(), + ) + .await?; + Ok(affected_query_result(affected)) + } MongoCommand::Update { collection, filter, update, options, many } => { let affected = mongo_update_documents_core( state, diff --git a/crates/dbx-core/src/mongo_shell.rs b/crates/dbx-core/src/mongo_shell.rs index faacae98e5..511416dd47 100644 --- a/crates/dbx-core/src/mongo_shell.rs +++ b/crates/dbx-core/src/mongo_shell.rs @@ -66,6 +66,8 @@ pub enum MongoCommand { }, #[serde(rename = "update")] Update { collection: String, filter: String, update: String, options: Option, many: bool }, + #[serde(rename = "replace")] + Replace { collection: String, filter: String, replacement: String, options: Option }, #[serde(rename = "delete")] Delete { collection: String, filter: String, many: bool }, #[serde(rename = "createIndex")] @@ -98,6 +100,7 @@ impl MongoCommand { | Self::CreateUser { .. } | Self::Insert { .. } | Self::Update { .. } + | Self::Replace { .. } | Self::Delete { .. } | Self::CreateIndex { .. } | Self::DropIndexes { .. } @@ -117,6 +120,7 @@ impl MongoCommand { pub fn has_empty_filter(&self) -> bool { match self { Self::Update { filter, .. } + | Self::Replace { filter, .. } | Self::Delete { filter, .. } | Self::FindOneAndUpdate { filter, .. } | Self::FindOneAndReplace { filter, .. } @@ -128,6 +132,7 @@ impl MongoCommand { pub fn has_effectively_unbounded_filter(&self) -> bool { match self { Self::Update { filter, .. } + | Self::Replace { filter, .. } | Self::Delete { filter, .. } | Self::FindOneAndUpdate { filter, .. } | Self::FindOneAndReplace { filter, .. } @@ -609,6 +614,22 @@ pub fn parse(input: &str) -> Result { return Ok(MongoCommand::Insert { collection, documents }); } + if let Some((args, tail)) = method_call(source, prefix_end, "replaceOne") { + if !tail.is_empty() || !(2..=3).contains(&args.len()) { + return Err( + "MongoDB replaceOne() requires a filter, a replacement document, and optional options.".to_string() + ); + } + let replacement = normalized_json(&args[1])?; + require_replacement_document(&replacement)?; + return Ok(MongoCommand::Replace { + collection, + filter: normalized_json(&args[0])?, + replacement, + options: args.get(2).filter(|arg| !arg.trim().is_empty()).map(|arg| normalized_json(arg)).transpose()?, + }); + } + for (method, many) in [("updateOne", false), ("updateMany", true)] { if let Some((args, tail)) = method_call(source, prefix_end, method) { if !tail.is_empty() || !(2..=3).contains(&args.len()) { @@ -1348,6 +1369,19 @@ fn parse_use_database(source: &str) -> Option { Some(database.to_string()) } +/// A replacement is a whole document; `{$set: ...}` here almost always means updateOne() was intended. +fn require_replacement_document(replacement: &str) -> Result<(), String> { + let Some(Value::Object(document)) = parse_json_value(replacement) else { + return Err("MongoDB replaceOne() replacement must be a document.".to_string()); + }; + match document.keys().find(|key| key.starts_with('$')) { + Some(operator) => Err(format!( + "MongoDB replaceOne() replacement must not contain update operators such as {operator}; use updateOne() to modify fields." + )), + None => Ok(()), + } +} + fn is_empty_object(value: &str) -> bool { parse_json_value(value).is_some_and(|value| value.as_object().is_some_and(|object| object.is_empty())) } @@ -1865,6 +1899,51 @@ mod tests { assert!(matches!(parse("db.orders.stats()").unwrap(), MongoCommand::CollectionStats { .. })); } + #[test] + fn parses_replace_one_as_a_filtered_write() { + let command = parse( + r#"db.orders.replaceOne({_id: ObjectId("507f1f77bcf86cd799439011")}, {name: "new", tags: []}, {upsert: true})"#, + ) + .unwrap(); + assert_eq!( + command, + MongoCommand::Replace { + collection: "orders".to_string(), + filter: r#"{"_id":{"$oid":"507f1f77bcf86cd799439011"}}"#.to_string(), + replacement: r#"{"name":"new","tags":[]}"#.to_string(), + options: Some(r#"{"upsert":true}"#.to_string()), + } + ); + assert!(command.is_mutating()); + assert!(!command.is_dangerous()); + assert!(!command.has_empty_filter()); + assert_eq!(validate_safety(&command, true, false, false), Ok(())); + assert_eq!(validate_safety(&command, false, false, false), Err(MongoSafetyError::WritesDisabled)); + + let without_options = parse("db.orders.replaceOne({a: 1}, {b: 2})").unwrap(); + assert!(matches!(without_options, MongoCommand::Replace { options: None, .. })); + + // An empty filter replaces an arbitrary document, so it is guarded like an update. + let unbounded = parse("db.orders.replaceOne({}, {b: 2})").unwrap(); + assert!(unbounded.has_empty_filter()); + assert_eq!(validate_safety(&unbounded, true, false, false), Err(MongoSafetyError::EmptyFilter)); + } + + #[test] + fn rejects_replace_one_with_operators_or_the_wrong_shape() { + let error = parse("db.orders.replaceOne({a: 1}, {$set: {b: 2}})").unwrap_err(); + assert!(error.contains("$set") && error.contains("updateOne"), "{error}"); + + for source in [ + "db.orders.replaceOne({a: 1})", + "db.orders.replaceOne({a: 1}, {b: 2}, {upsert: true}, 4)", + "db.orders.replaceOne({a: 1}, [{b: 2}])", + "db.orders.replaceOne({a: 1}, {b: 2}).limit(1)", + ] { + assert!(parse(source).is_err(), "{source}"); + } + } + #[test] fn parses_show_databases_aliases_as_read_only() { for source in [ diff --git a/crates/dbx-mcp/src/backend.rs b/crates/dbx-mcp/src/backend.rs index 438204e2e4..ce3d528894 100644 --- a/crates/dbx-mcp/src/backend.rs +++ b/crates/dbx-mcp/src/backend.rs @@ -1860,6 +1860,26 @@ impl DbxBackend for WebBackend { .map_err(|error| format!("Invalid MongoDB insert response: {error}"))?; Ok(affected_query_result(affected_rows_from_value(&value))) } + MongoCommand::Replace { collection, filter, replacement, options } => { + let value: Value = self + .request( + reqwest::Method::POST, + "/api/mongo/replace-document", + Some(json!({ + "connectionId": connection_id, + "database": database, + "collection": collection, + "filterJson": filter, + "replacementJson": replacement, + "optionsJson": options, + })), + ) + .await? + .json() + .await + .map_err(|error| format!("Invalid MongoDB replace response: {error}"))?; + Ok(affected_query_result(affected_rows_from_value(&value))) + } MongoCommand::Update { collection, filter, update, options, many } => { let value: Value = self .request( diff --git a/crates/dbx-web/src/main.rs b/crates/dbx-web/src/main.rs index 7083a4c014..467f16270e 100644 --- a/crates/dbx-web/src/main.rs +++ b/crates/dbx-web/src/main.rs @@ -967,6 +967,7 @@ async fn main() { .route("/mongo/insert-documents", post(routes::mongo::insert_documents)) .route("/mongo/update-document", post(routes::mongo::update_document)) .route("/mongo/update-documents", post(routes::mongo::update_documents)) + .route("/mongo/replace-document", post(routes::mongo::replace_document)) .route("/mongo/delete-document", post(routes::mongo::delete_document)) .route("/mongo/delete-documents", post(routes::mongo::delete_documents)) .route("/mongo/find-one-and-update", post(routes::mongo::find_one_and_update)) diff --git a/crates/dbx-web/src/routes/mongo.rs b/crates/dbx-web/src/routes/mongo.rs index 70e2b4c951..90e33ce4fc 100644 --- a/crates/dbx-web/src/routes/mongo.rs +++ b/crates/dbx-web/src/routes/mongo.rs @@ -258,6 +258,17 @@ pub struct MongoUpdateDocumentsRequest { pub options_json: Option, } +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct MongoReplaceDocumentRequest { + pub connection_id: String, + pub database: String, + pub collection: String, + pub filter_json: String, + pub replacement_json: String, + pub options_json: Option, +} + #[derive(Deserialize)] #[serde(rename_all = "camelCase")] pub struct MongoDeleteDocumentsRequest { @@ -809,6 +820,32 @@ pub async fn update_documents( Ok(Json(serde_json::json!({ "affected_rows": result }))) } +pub async fn replace_document( + State(state): State>, + headers: HeaderMap, + Json(req): Json, +) -> Result, AppError> { + let database = super::mcp_policy::resolve_database(&state, &headers, &req.connection_id, &req.database).await?; + if super::mcp_policy::mongo_filter_is_effectively_unbounded(&req.filter_json) { + super::mcp_policy::ensure_dangerous_write(&state, &headers, &req.connection_id, &database, "Replace").await?; + } else { + super::mcp_policy::ensure_write(&state, &headers, &req.connection_id, &database, "Replace").await?; + } + ensure_writable(&state.app, &req.connection_id, "Replace").await?; + let result = dbx_core::mongo_ops::mongo_replace_document_core( + &state.app, + &req.connection_id, + &database, + &req.collection, + &req.filter_json, + &req.replacement_json, + req.options_json.as_deref(), + ) + .await + .map_err(AppError::from)?; + Ok(Json(serde_json::json!({ "affected_rows": result }))) +} + pub async fn find_one_and_update( State(state): State>, headers: HeaderMap, diff --git a/packages/app-tests/mongoCompletion.test.ts b/packages/app-tests/mongoCompletion.test.ts index e6b298232d..574ceee81a 100644 --- a/packages/app-tests/mongoCompletion.test.ts +++ b/packages/app-tests/mongoCompletion.test.ts @@ -500,8 +500,9 @@ test("suggests only helpers the shell parser accepts", () => { assert.ok(methodLabels.includes("drop")); assert.ok(methodLabels.includes("distinct")); assert.ok(methodLabels.includes("estimatedDocumentCount")); + assert.ok(methodLabels.includes("replaceOne")); // Suggesting a helper DBX cannot run just hands the user a command that fails. - for (const unsupported of ["bulkWrite", "replaceOne", "renameCollection"]) { + for (const unsupported of ["bulkWrite", "renameCollection", "mapReduce"]) { assert.equal(methodLabels.includes(unsupported), false, `${unsupported} is not executable`); } // Cursor methods are not collection methods. diff --git a/packages/app-tests/mongoShellCommand.test.ts b/packages/app-tests/mongoShellCommand.test.ts index 9dfef73dfc..2d4bee7e98 100644 --- a/packages/app-tests/mongoShellCommand.test.ts +++ b/packages/app-tests/mongoShellCommand.test.ts @@ -672,6 +672,45 @@ test("parseMongoWriteCommand accepts unquoted insert and update commands", () => }); }); +test("parseMongoWriteCommand reads replaceOne as a filtered replace", () => { + assert.deepEqual(parseMongoWriteCommand('db.orders.replaceOne({_id: ObjectId("507f1f77bcf86cd799439011")}, {name: "new", tags: []}, {upsert: true})'), { + kind: "replace", + collection: "orders", + filter: '{"_id": {"$oid":"507f1f77bcf86cd799439011"}}', + replacement: '{"name": "new", "tags": []}', + options: '{"upsert": true}', + }); + assert.deepEqual(parseMongoWriteCommand('db["my-coll"].replaceOne({a: 1}, {b: 2});'), { + kind: "replace", + collection: "my-coll", + filter: '{"a": 1}', + replacement: '{"b": 2}', + }); + + for (const source of ["db.orders.replaceOne({a: 1})", "db.orders.replaceOne({a: 1}, {b: 2}, {upsert: true}, 4)", "db.orders.replaceOne({a: 1}, [{b: 2}])", "db.orders.replaceOne({a: 1}, {$set: {b: 2}})"]) { + assert.equal(parseMongoWriteCommand(source), null, source); + } + assert.match(describeMongoCommandParseFailure("db.orders.replaceOne({a: 1}, {$set: {b: 2}})"), /must not contain update operators such as \$set; use updateOne\(\)/); +}); + +test("evaluateMongoWriteSafety guards an unbounded replaceOne like an update", () => { + const bounded = parseMongoWriteCommand("db.orders.replaceOne({a: 1}, {b: 2})")!; + const unbounded = parseMongoWriteCommand("db.orders.replaceOne({}, {b: 2})")!; + const policy = { allowWrites: true, allowDangerous: false } as Parameters[1]; + assert.equal(evaluateMongoWriteSafety(bounded, policy).allowed, true); + assert.equal(evaluateMongoWriteSafety(unbounded, policy).allowed, false); + assert.equal(evaluateMongoWriteSafety(bounded, { ...policy, allowWrites: false }).allowed, false); +}); + +test("normalizeRustMongoCommand passes a replace command through unchanged", () => { + assert.deepEqual(normalizeRustMongoCommand({ kind: "replace", collection: "orders", filter: '{"a":1}', replacement: '{"b":2}', options: null }), { + kind: "replace", + collection: "orders", + filter: '{"a":1}', + replacement: '{"b":2}', + }); +}); + test("parseMongoWriteCommand unwraps EJSON.deserialize values", () => { assert.deepEqual(parseMongoWriteCommand('db.products.updateOne({_id: ObjectId("507f1f77bcf86cd799439011")}, {$set: {price: EJSON.deserialize({"$numberDecimal":"12.34"}), payload: EJSON.deserialize({"$binary":{"base64":"AQI=","subType":"00"}})}})'), { kind: "update", @@ -1210,8 +1249,8 @@ test("describeMongoCommandParseFailure explains the argument shape a known metho }); test("describeMongoCommandParseFailure names unsupported methods and points at alternatives", () => { - const replaceOne = describeMongoCommandParseFailure("db.c.replaceOne({a: 1}, {b: 2})"); - assert.match(replaceOne, /^Collection method replaceOne\(\) is not supported\. Supported collection methods: find, findOne/); + const bulkWrite = describeMongoCommandParseFailure("db.c.bulkWrite([{insertOne: {document: {a: 1}}}])"); + assert.match(bulkWrite, /^Collection method bulkWrite\(\) is not supported\. Supported collection methods: find, findOne/); // Bracket and getCollection targets are recognised too. assert.match(describeMongoCommandParseFailure('db["my-coll"].renameCollection("x")'), /renameCollection\(\) is not supported/); assert.match(describeMongoCommandParseFailure('db.getCollection("my-coll").watch()'), /watch\(\) is not supported/); diff --git a/src-tauri/src/commands/mongo_cmd.rs b/src-tauri/src/commands/mongo_cmd.rs index df23d3512b..4a3e6f55d4 100644 --- a/src-tauri/src/commands/mongo_cmd.rs +++ b/src-tauri/src/commands/mongo_cmd.rs @@ -541,6 +541,40 @@ pub async fn mongo_update_documents( .await } +#[tauri::command] +pub async fn mongo_replace_document( + state: State<'_, Arc>, + connection_id: String, + database: String, + collection: String, + filter_json: String, + replacement_json: String, + options_json: Option, + mcp_request: Option, +) -> Result { + if mcp_request == Some(true) { + crate::commands::mcp_bridge::ensure_mcp_mongo_filtered_write_allowed_by_id( + state.inner(), + &connection_id, + &database, + "Replace", + &filter_json, + ) + .await?; + } + ensure_connection_writable(&state, &connection_id, "Replace").await?; + dbx_core::mongo_ops::mongo_replace_document_core( + &state, + &connection_id, + &database, + &collection, + &filter_json, + &replacement_json, + options_json.as_deref(), + ) + .await +} + #[tauri::command] pub async fn mongo_delete_document( state: State<'_, Arc>, diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 20b0d9837d..028b532fc7 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -2335,6 +2335,7 @@ pub fn run() { commands::document_cmd::document_update_document, commands::mongo_cmd::mongo_update_document, commands::mongo_cmd::mongo_update_documents, + commands::mongo_cmd::mongo_replace_document, commands::document_cmd::document_delete_document, commands::document_cmd::document_save_meilisearch_batch, commands::document_cmd::meilisearch_search_documents,