From 9b24d962bb83dd00f64e394d4cfbb3b8347e3f45 Mon Sep 17 00:00:00 2001 From: Guoyu Su Date: Wed, 12 Aug 2026 21:05:13 +0800 Subject: [PATCH] feat(mongodb): support generic runCommand execution --- .../java/com/dbx/agent/AgentProtocol.java | 10 ++- .../src/main/resources/agent-protocol-v1.json | 4 +- .../src/main/resources/agent-protocol-v2.json | 4 +- .../com/dbx/agent/mongodb/MongoAgent.java | 16 ++++ .../com/dbx/agent/mongodb/MongoAgentTest.java | 63 ++++++++++++++ apps/desktop/src/lib/backend/api.ts | 1 + apps/desktop/src/lib/backend/http.ts | 9 ++ apps/desktop/src/lib/backend/tauri.ts | 9 ++ .../src/lib/mongo/mongoShellCommand.ts | 26 +++++- apps/desktop/src/stores/queryStore.ts | 20 +++++ crates/dbx-core/assets/agent-protocol-v1.json | 4 +- crates/dbx-core/assets/agent-protocol-v2.json | 4 +- crates/dbx-core/src/agent_tools.rs | 1 + crates/dbx-core/src/db/agent_driver.rs | 29 ++++++- crates/dbx-core/src/db/mongo_driver.rs | 19 +++++ crates/dbx-core/src/mongo_ops.rs | 82 +++++++++++++++++++ crates/dbx-core/src/mongo_shell.rs | 63 +++++++++++++- .../tests/live_mongodb_run_command.rs | 45 ++++++++++ crates/dbx-mcp/src/backend.rs | 3 + crates/dbx-mcp/src/server.rs | 22 +++++ crates/dbx-web/src/main.rs | 1 + crates/dbx-web/src/routes/mongo.rs | 32 ++++++++ packages/app-tests/mongoShellCommand.test.ts | 50 ++++++++++- packages/app-tests/queryStore.test.ts | 50 +++++++++++ packages/mongo-shell/src/aggregate.ts | 2 +- src-tauri/src/commands/mongo_cmd.rs | 28 +++++++ src-tauri/src/lib.rs | 1 + 27 files changed, 580 insertions(+), 18 deletions(-) create mode 100644 crates/dbx-core/tests/live_mongodb_run_command.rs 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 9e647c05fd..178edc4cc0 100644 --- a/agents/common/src/main/java/com/dbx/agent/AgentProtocol.java +++ b/agents/common/src/main/java/com/dbx/agent/AgentProtocol.java @@ -68,6 +68,7 @@ public final class AgentProtocol { public static final String MONGO_METHOD_UPDATE_DOCUMENTS = "update_documents"; 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"; public static final String KV_METHOD_LIST_PREFIX = "kv_list_prefix"; public static final String KV_METHOD_GET = "kv_get"; @@ -120,6 +121,7 @@ public final class AgentProtocol { public static final String CAPABILITY_ETCD_AUTH = "etcd_auth"; public static final String CAPABILITY_MONGO_DROP_DATABASE = "mongo_drop_database"; 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_MULTI_SESSION = "multi_session"; public static final String CAPABILITY_STRUCTURED_ERROR_V1 = "structured_error_v1"; @@ -153,7 +155,8 @@ public final class AgentProtocol { CAPABILITY_ETCD_LEASE, CAPABILITY_ETCD_AUTH, CAPABILITY_MONGO_DROP_DATABASE, - CAPABILITY_MONGO_CLONE_COLLECTION + CAPABILITY_MONGO_CLONE_COLLECTION, + CAPABILITY_MONGO_RUN_COMMAND )); public static final List MULTI_SESSION_CAPABILITIES; @@ -212,11 +215,13 @@ public final class AgentProtocol { List mongoCapabilities = new java.util.ArrayList<>(CAPABILITIES); mongoCapabilities.add(CAPABILITY_MONGO_DROP_DATABASE); mongoCapabilities.add(CAPABILITY_MONGO_CLONE_COLLECTION); + mongoCapabilities.add(CAPABILITY_MONGO_RUN_COMMAND); MONGO_LEGACY_CAPABILITIES = Collections.unmodifiableList(mongoCapabilities); List mongoMultiSessionCapabilities = new java.util.ArrayList<>(MULTI_SESSION_CAPABILITIES); mongoMultiSessionCapabilities.add(CAPABILITY_MONGO_DROP_DATABASE); mongoMultiSessionCapabilities.add(CAPABILITY_MONGO_CLONE_COLLECTION); + mongoMultiSessionCapabilities.add(CAPABILITY_MONGO_RUN_COMMAND); MONGO_LEGACY_MULTI_SESSION_CAPABILITIES = Collections.unmodifiableList(mongoMultiSessionCapabilities); List jdbcCapabilities = new java.util.ArrayList<>(MULTI_SESSION_CAPABILITIES); @@ -258,7 +263,8 @@ public final class AgentProtocol { MONGO_METHOD_UPDATE_DOCUMENT, MONGO_METHOD_UPDATE_DOCUMENTS, MONGO_METHOD_DELETE_DOCUMENT, - MONGO_METHOD_DELETE_DOCUMENTS + MONGO_METHOD_DELETE_DOCUMENTS, + MONGO_METHOD_RUN_COMMAND )); public static final List KV_METHODS = Collections.unmodifiableList(Arrays.asList( diff --git a/agents/common/src/main/resources/agent-protocol-v1.json b/agents/common/src/main/resources/agent-protocol-v1.json index 7b5b10f6ec..ffb0ad0f7d 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"], + "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"], "capabilities": ["connect", "test_connection", "metadata", "query", "paged_query", "transaction", "ddl"], "defaultSqlCapabilities": ["connect", "test_connection", "metadata", "query", "paged_query", "transaction", "ddl"], "commonMethods": [ @@ -39,6 +39,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", "update_document", "update_documents", "delete_document", "delete_documents"], + "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", "update_document", "update_documents", "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 211f63db38..b8363bfb2b 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", "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", "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": [ @@ -43,7 +43,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", "update_document", "update_documents", "delete_document", "delete_documents"], + "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", "update_document", "update_documents", "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 5f5155312b..dde1dcac4f 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 @@ -940,6 +940,21 @@ public final class MongoAgent { return serverVersionFromBuildInfo(buildInfo); } + private static Object runCommand(JsonObject params) { + MongoClient c = requireClient(); + String database = params.get("database").getAsString(); + Document command = documentOrNull(params, "command_json"); + if (command == null || command.isEmpty()) { + throw new IllegalArgumentException("runCommand requires a non-empty command document"); + } + Document response = c.getDatabase(database).runCommand(command); + List> documents = new ArrayList<>(); + documents.add(bsonToJson(response)); + List extendedDocuments = new ArrayList<>(); + extendedDocuments.add(bsonToExtendedJson(response)); + return documentQueryResultWithExtended(documents, extendedDocuments, 1); + } + static String serverVersionFromBuildInfo(Document buildInfo) { String version = buildInfo.getString("version"); if (version == null || version.isBlank()) { @@ -1819,6 +1834,7 @@ public final class MongoAgent { case AgentProtocol.MONGO_METHOD_UPDATE_DOCUMENTS -> updateDocuments(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); case AgentProtocol.METHOD_DISCONNECT, AgentProtocol.METHOD_SHUTDOWN -> { closeLegacyClient(); if (AgentProtocol.METHOD_SHUTDOWN.equals(method)) { 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 17642a640a..5aee2fac68 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 @@ -126,6 +126,7 @@ class MongoAgentTest { assertTrue(containsCapability(result.getAsJsonArray("capabilities"), AgentProtocol.CAPABILITY_METADATA)); assertTrue(containsCapability(result.getAsJsonArray("capabilities"), AgentProtocol.CAPABILITY_MONGO_DROP_DATABASE)); assertTrue(containsCapability(result.getAsJsonArray("capabilities"), AgentProtocol.CAPABILITY_MONGO_CLONE_COLLECTION)); + assertTrue(containsCapability(result.getAsJsonArray("capabilities"), AgentProtocol.CAPABILITY_MONGO_RUN_COMMAND)); } @Test @@ -146,6 +147,7 @@ class MongoAgentTest { assertTrue(containsCapability(result.getAsJsonArray("capabilities"), AgentProtocol.CAPABILITY_MULTI_SESSION)); assertTrue(containsCapability(result.getAsJsonArray("capabilities"), AgentProtocol.CAPABILITY_MONGO_DROP_DATABASE)); assertTrue(containsCapability(result.getAsJsonArray("capabilities"), AgentProtocol.CAPABILITY_MONGO_CLONE_COLLECTION)); + assertTrue(containsCapability(result.getAsJsonArray("capabilities"), AgentProtocol.CAPABILITY_MONGO_RUN_COMMAND)); } @Test @@ -536,6 +538,67 @@ class MongoAgentTest { assertFalse(json.getAsJsonObject("error").get("message").getAsString().contains("Unknown method")); } + @Test + void runCommandExecutesTheCommandDocumentAndPreservesNestedExtendedJson() { + List commands = new ArrayList<>(); + MongoDatabase database = (MongoDatabase) Proxy.newProxyInstance( + MongoDatabase.class.getClassLoader(), + new Class[] {MongoDatabase.class}, + (proxy, method, args) -> { + if ("runCommand".equals(method.getName())) { + commands.add(new Document((Document) args[0])); + return new Document("ok", 1) + .append("cursor", new Document("firstBatch", List.of( + new Document("_id", new ObjectId("507f1f77bcf86cd799439011")) + .append("counter", 9_007_199_254_740_992L) + ))); + } + throw new UnsupportedOperationException(method.getName()); + } + ); + MongoClient client = (MongoClient) Proxy.newProxyInstance( + MongoClient.class.getClassLoader(), + new Class[] {MongoClient.class}, + (proxy, method, args) -> { + if ("getDatabase".equals(method.getName())) { + assertEquals("app", args[0]); + return database; + } + throw new UnsupportedOperationException(method.getName()); + } + ); + JsonObject params = new JsonObject(); + params.addProperty("database", "app"); + params.addProperty("command_json", "{\"find\":\"orders\",\"filter\":{\"active\":true}}"); + JsonObject request = new JsonObject(); + request.addProperty("jsonrpc", "2.0"); + request.addProperty("id", 10); + request.addProperty("method", "run_command"); + request.add("params", params); + + JsonObject json = JsonParser.parseString(MongoAgent.handleRequest(request.toString(), client)).getAsJsonObject(); + + assertFalse(json.has("error"), json.toString()); + assertEquals( + new Document("find", "orders").append("filter", new Document("active", true)), + commands.get(0) + ); + JsonObject result = json.getAsJsonObject("result"); + assertEquals(1, result.get("total").getAsInt()); + assertEquals(1, result.getAsJsonArray("documents").get(0).getAsJsonObject().get("ok").getAsInt()); + JsonObject copiedId = result.getAsJsonArray("extended_documents") + .get(0).getAsJsonObject() + .getAsJsonObject("cursor").getAsJsonArray("firstBatch") + .get(0).getAsJsonObject().getAsJsonObject("_id"); + assertEquals("507f1f77bcf86cd799439011", copiedId.get("$oid").getAsString()); + JsonObject copiedCounter = result.getAsJsonArray("extended_documents") + .get(0).getAsJsonObject() + .getAsJsonObject("cursor").getAsJsonArray("firstBatch") + .get(0).getAsJsonObject().getAsJsonObject("counter"); + assertEquals("9007199254740992", copiedCounter.get("$numberLong").getAsString()); + assertTrue(AgentProtocol.MONGO_LEGACY_METHODS.contains(AgentProtocol.MONGO_METHOD_RUN_COMMAND)); + } + @Test void createIndexMethodIsRecognizedOverJsonRpc() { String response = MongoAgent.handleRequest( diff --git a/apps/desktop/src/lib/backend/api.ts b/apps/desktop/src/lib/backend/api.ts index c2b402dd2d..36b3b6982e 100644 --- a/apps/desktop/src/lib/backend/api.ts +++ b/apps/desktop/src/lib/backend/api.ts @@ -690,6 +690,7 @@ export const mongoDistinct = forward("mongoDistinct"); export const mongoCollectionStats = forward("mongoCollectionStats"); export const mongoCreateIndex = forward("mongoCreateIndex"); export const mongoCreateUser = forward("mongoCreateUser"); +export const mongoRunCommand = forward("mongoRunCommand"); export const mongoDropIndexes = forward("mongoDropIndexes"); export const documentInsertDocument = forward("documentInsertDocument"); export const mongoInsertDocument = forward("mongoInsertDocument"); diff --git a/apps/desktop/src/lib/backend/http.ts b/apps/desktop/src/lib/backend/http.ts index c9e78010a3..15032c55c4 100644 --- a/apps/desktop/src/lib/backend/http.ts +++ b/apps/desktop/src/lib/backend/http.ts @@ -3583,6 +3583,15 @@ export async function mongoCreateUser(connectionId: string, database: string, us }); } +export async function mongoRunCommand(connectionId: string, database: string, commandJson: string, executionId?: string): Promise { + return post("/api/mongo/run-command", { + connectionId, + database, + commandJson, + executionId, + }); +} + export async function mongoDropIndexes(connectionId: string, database: string, collection: string, indexesJson?: string, single = false): Promise { return post("/api/mongo/drop-indexes", { connectionId, diff --git a/apps/desktop/src/lib/backend/tauri.ts b/apps/desktop/src/lib/backend/tauri.ts index ea7e6d1712..4dcaaf4b0b 100644 --- a/apps/desktop/src/lib/backend/tauri.ts +++ b/apps/desktop/src/lib/backend/tauri.ts @@ -3527,6 +3527,15 @@ export async function mongoCreateUser(connectionId: string, database: string, us }); } +export async function mongoRunCommand(connectionId: string, database: string, commandJson: string, executionId?: string): Promise { + return invoke("mongo_run_command", { + connectionId, + database, + commandJson, + executionId, + }); +} + export async function mongoDropIndexes(connectionId: string, database: string, collection: string, indexesJson?: string, single = false): Promise { return invoke("mongo_drop_indexes", { connectionId, diff --git a/apps/desktop/src/lib/mongo/mongoShellCommand.ts b/apps/desktop/src/lib/mongo/mongoShellCommand.ts index 90375653cd..52f648cfb9 100644 --- a/apps/desktop/src/lib/mongo/mongoShellCommand.ts +++ b/apps/desktop/src/lib/mongo/mongoShellCommand.ts @@ -67,6 +67,10 @@ export interface MongoCreateUserCommand { writeConcernJson?: string; } +export interface MongoRunCommand { + commandJson: string; +} + export type MongoCollectionStatsMetric = "stats" | "dataSize" | "storageSize" | "totalIndexSize"; export interface MongoCollectionStatsCommand { @@ -81,7 +85,7 @@ export interface MongoDistinctCommand { filter?: string; } -type MongoWriteKind = "insert" | "update" | "delete" | "createIndex" | "createUser" | "dropIndex" | "dropIndexes" | "dropCollection" | "findOneAndUpdate" | "findOneAndReplace" | "findOneAndDelete"; +type MongoWriteKind = "runCommand" | "insert" | "update" | "delete" | "createIndex" | "createUser" | "dropIndex" | "dropIndexes" | "dropCollection" | "findOneAndUpdate" | "findOneAndReplace" | "findOneAndDelete"; export type MongoCommand = | ({ kind: "find" } & MongoFindCommand) @@ -94,6 +98,7 @@ export type MongoCommand = | ({ kind: "collectionStats" } & MongoCollectionStatsCommand) | ({ kind: "use" } & MongoUseCommand) | ({ kind: "createUser" } & MongoCreateUserCommand) + | ({ kind: "runCommand" } & MongoRunCommand) | { kind: "insert"; collection: string; docsJson: string } | { kind: "update"; collection: string; filter: string; update: string; options?: string; many: boolean } | { kind: "delete"; collection: string; filter: string; many: boolean } @@ -499,6 +504,21 @@ export function parseMongoCreateUserCommand(input: string): MongoCreateUserComma }; } +export function parseMongoRunCommand(input: string): MongoRunCommand | null { + const source = input.trim().replace(/;$/, "").trim(); + const match = /^db\s*\.\s*runCommand\s*\(/i.exec(source); + if (!match) return null; + const openIndex = source.indexOf("(", match.index); + const closeIndex = findMatchingParen(source, openIndex); + if (closeIndex < 0 || source.slice(closeIndex + 1).trim()) return null; + const args = splitTopLevel(source.slice(openIndex + 1, closeIndex)); + if (args.length !== 1) return null; + const commandJson = parseMongoObjectArgument(args[0]); + if (!commandJson) return null; + const command = JSON.parse(commandJson) as Record; + return Object.keys(command).length > 0 ? { commandJson: JSON.stringify(command) } : null; +} + export function parseMongoWriteCommand(input: string): MongoWriteCommand | null { const source = input.trim().replace(/;$/, "").trim(); const insertOne = parseCollectionMethodTarget(source, "insertOne"); @@ -607,6 +627,10 @@ export function parseMongoCommand(input: string): ParsedMongoCommand | null { const createUser = parseMongoCreateUserCommand(source); return createUser ? { kind: "createUser", ...createUser } : null; }, + (source) => { + const runCommand = parseMongoRunCommand(source); + return runCommand ? { kind: "runCommand", ...runCommand } : null; + }, (source) => { // Legacy Mongo shell uses count()/find().count(); keep accepting it // while mapping to DBX's countDocuments-compatible result path. diff --git a/apps/desktop/src/stores/queryStore.ts b/apps/desktop/src/stores/queryStore.ts index 58bf84d8c3..765fbc7317 100644 --- a/apps/desktop/src/stores/queryStore.ts +++ b/apps/desktop/src/stores/queryStore.ts @@ -4204,6 +4204,26 @@ export const useQueryStore = defineStore("query", () => { }); break; } + case "runCommand": { + if (options?.mongoSafety) { + const safety = evaluateMongoWriteSafety(mongoCommand, options.mongoSafety); + if (!safety.allowed) throw new Error(safety.reason); + } + queryExecutionLog("info", "mongo-run-command:start", { + traceId, + database: currentDatabase, + }); + const result = await api.mongoRunCommand(executionConnectionId, currentDatabase, mongoCommand.commandJson, executionId); + allResults.push(markQueryResultRowsRaw(annotateMongoResult(mongoDocumentsToQueryResult(result.documents, performance.now() - commandStartedAt, result.total, result.extended_documents, result.total_is_exact !== false)))); + mongoEditTarget = undefined; + queryExecutionLog("info", "mongo-run-command:done", { + traceId, + database: currentDatabase, + rowCount: result.documents.length, + elapsed: elapsed(), + }); + break; + } case "insert": case "update": case "delete": diff --git a/crates/dbx-core/assets/agent-protocol-v1.json b/crates/dbx-core/assets/agent-protocol-v1.json index 7b5b10f6ec..ffb0ad0f7d 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"], + "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"], "capabilities": ["connect", "test_connection", "metadata", "query", "paged_query", "transaction", "ddl"], "defaultSqlCapabilities": ["connect", "test_connection", "metadata", "query", "paged_query", "transaction", "ddl"], "commonMethods": [ @@ -39,6 +39,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", "update_document", "update_documents", "delete_document", "delete_documents"], + "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", "update_document", "update_documents", "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 211f63db38..b8363bfb2b 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", "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", "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": [ @@ -43,7 +43,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", "update_document", "update_documents", "delete_document", "delete_documents"], + "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", "update_document", "update_documents", "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/agent_tools.rs b/crates/dbx-core/src/agent_tools.rs index a0b68e4f2f..04318362f9 100644 --- a/crates/dbx-core/src/agent_tools.rs +++ b/crates/dbx-core/src/agent_tools.rs @@ -1387,6 +1387,7 @@ for line in sys.stdin: "db.items.updateMany({tenant: 7}, {$set: {active: false}})", "db.items.deleteMany({tenant: 7})", "db.items.createIndex({tenant: 1})", + "db.runCommand({compact: 'items'})", r#"db.items.aggregate([{"$out":"items_backup"}])"#, ] .into_iter() diff --git a/crates/dbx-core/src/db/agent_driver.rs b/crates/dbx-core/src/db/agent_driver.rs index 36a5d687d1..726243d721 100644 --- a/crates/dbx-core/src/db/agent_driver.rs +++ b/crates/dbx-core/src/db/agent_driver.rs @@ -973,6 +973,7 @@ pub enum AgentCapability { EtcdAuth, MongoDropDatabase, MongoCloneCollection, + MongoRunCommand, MultiSession, StructuredErrorV1, } @@ -1055,7 +1056,7 @@ fn parse_agent_rpc_error_header(header: &str) -> (Option, String) { } impl AgentCapability { - pub const ALL: [Self; 22] = [ + pub const ALL: [Self; 23] = [ Self::Connect, Self::TestConnection, Self::Metadata, @@ -1076,6 +1077,7 @@ impl AgentCapability { Self::EtcdAuth, Self::MongoDropDatabase, Self::MongoCloneCollection, + Self::MongoRunCommand, Self::MultiSession, Self::StructuredErrorV1, ]; @@ -1102,6 +1104,7 @@ impl AgentCapability { Self::EtcdAuth => "etcd_auth", Self::MongoDropDatabase => "mongo_drop_database", Self::MongoCloneCollection => "mongo_clone_collection", + Self::MongoRunCommand => "mongo_run_command", Self::MultiSession => "multi_session", Self::StructuredErrorV1 => "structured_error_v1", } @@ -1294,10 +1297,11 @@ pub enum MongoAgentMethod { UpdateDocuments, DeleteDocument, DeleteDocuments, + RunCommand, } impl MongoAgentMethod { - pub const ALL: [Self; 20] = [ + pub const ALL: [Self; 21] = [ Self::ListDatabases, Self::ListCollections, Self::FindDocuments, @@ -1318,6 +1322,7 @@ impl MongoAgentMethod { Self::UpdateDocuments, Self::DeleteDocument, Self::DeleteDocuments, + Self::RunCommand, ]; pub fn as_str(self) -> &'static str { @@ -1342,6 +1347,7 @@ impl MongoAgentMethod { Self::UpdateDocuments => "update_documents", Self::DeleteDocument => "delete_document", Self::DeleteDocuments => "delete_documents", + Self::RunCommand => "run_command", } } } @@ -2539,6 +2545,13 @@ impl AgentDriverClient { self.call_mongo_method(MongoAgentMethod::ServerVersion, mongo_database_params(database)).await } + pub async fn mongo_run_command( + &mut self, + params: Value, + ) -> Result { + self.call_mongo_method(MongoAgentMethod::RunCommand, params).await + } + pub async fn mongo_create_index( &mut self, params: Value, @@ -2755,6 +2768,7 @@ pub fn agent_supports_capability(handshake: Option<&AgentHandshake>, capability: | AgentCapability::EtcdAuth | AgentCapability::MongoDropDatabase | AgentCapability::MongoCloneCollection + | AgentCapability::MongoRunCommand ) { return handshake.map(|value| value.supports(capability)).unwrap_or(false); } @@ -4343,9 +4357,10 @@ for line in sys.stdin: assert_eq!(AgentCapability::EtcdAuth.as_str(), "etcd_auth"); assert_eq!(AgentCapability::MongoDropDatabase.as_str(), "mongo_drop_database"); assert_eq!(AgentCapability::MongoCloneCollection.as_str(), "mongo_clone_collection"); + assert_eq!(AgentCapability::MongoRunCommand.as_str(), "mongo_run_command"); assert_eq!(AgentCapability::MultiSession.as_str(), "multi_session"); assert_eq!(AgentCapability::StructuredErrorV1.as_str(), "structured_error_v1"); - assert_eq!(AgentCapability::ALL.len(), 22); + assert_eq!(AgentCapability::ALL.len(), 23); } #[test] @@ -4702,6 +4717,8 @@ for line in sys.stdin: assert!(!agent_supports_capability(Some(&handshake), AgentCapability::MongoDropDatabase)); assert!(!agent_supports_capability(None, AgentCapability::MongoCloneCollection)); assert!(!agent_supports_capability(Some(&handshake), AgentCapability::MongoCloneCollection)); + assert!(!agent_supports_capability(None, AgentCapability::MongoRunCommand)); + assert!(!agent_supports_capability(Some(&handshake), AgentCapability::MongoRunCommand)); let mongo_handshake = AgentHandshake { capabilities: vec![AgentCapability::MongoDropDatabase.as_str().to_string()], ..handshake }; @@ -4712,6 +4729,12 @@ for line in sys.stdin: ..mongo_handshake }; assert!(agent_supports_capability(Some(&mongo_clone_handshake), AgentCapability::MongoCloneCollection)); + + let mongo_run_command_handshake = AgentHandshake { + capabilities: vec![AgentCapability::MongoRunCommand.as_str().to_string()], + ..mongo_clone_handshake + }; + assert!(agent_supports_capability(Some(&mongo_run_command_handshake), AgentCapability::MongoRunCommand)); } #[test] diff --git a/crates/dbx-core/src/db/mongo_driver.rs b/crates/dbx-core/src/db/mongo_driver.rs index 49f0b7b73a..842989ba4a 100644 --- a/crates/dbx-core/src/db/mongo_driver.rs +++ b/crates/dbx-core/src/db/mongo_driver.rs @@ -173,6 +173,25 @@ pub async fn server_version(client: &Client, database: &str) -> Result Result { + let value: serde_json::Value = + serde_json::from_str(command_json).map_err(|error| format!("Invalid MongoDB command JSON: {error}"))?; + let command = json_object_to_document_extended_json(&value) + .map_err(|error| format!("Invalid MongoDB command document: {error}"))?; + if command.is_empty() { + return Err("MongoDB runCommand requires a non-empty command document".to_string()); + } + let result = client.database(database).run_command(command).await.map_err(|error| error.to_string())?; + let (document, extended_document) = document_json_views(result); + Ok(MongoDocumentResult { + documents: vec![document], + raw_documents: None, + extended_documents: Some(vec![extended_document]), + total: 1, + total_is_exact: true, + }) +} + fn server_version_from_build_info(result: &Document) -> Result { result.get_str("version").map(str::to_string).map_err(|e| format!("MongoDB server version not found: {e}")) } diff --git a/crates/dbx-core/src/mongo_ops.rs b/crates/dbx-core/src/mongo_ops.rs index 739ee31923..6b7cba3394 100644 --- a/crates/dbx-core/src/mongo_ops.rs +++ b/crates/dbx-core/src/mongo_ops.rs @@ -148,6 +148,35 @@ pub async fn mongo_server_version_core( } } +pub async fn mongo_run_command_core( + state: &AppState, + connection_id: &str, + database: &str, + command_json: &str, +) -> Result { + ensure_document_pool(state, connection_id).await?; + let connections = state.connections.read().await; + match connections.get(connection_id).ok_or("Not found")? { + PoolKind::MongoDb(client) => mongo_driver::run_command(client, database, command_json).await, + PoolKind::Agent(client) => { + let mut client = client.lock().await; + if !client.supports_capability(AgentCapability::MongoRunCommand) { + return Err( + "MongoDB Legacy Agent does not support runCommand; upgrade or reinstall the MongoDB Legacy driver" + .to_string(), + ); + } + client + .mongo_run_command(serde_json::json!({ + "database": database, + "command_json": command_json, + })) + .await + } + _ => Err("Not a MongoDB connection".to_string()), + } +} + pub async fn mongo_collection_stats_core( state: &AppState, connection_id: &str, @@ -821,6 +850,10 @@ pub async fn execute_mongo_command_core( .await .map(|version| scalar_query_result("version", Value::String(version))), MongoCommand::Use { database } => Ok(scalar_query_result("database", Value::String(database.clone()))), + MongoCommand::RunCommand { command_json } => { + let result = mongo_run_command_core(state, connection_id, database, command_json).await?; + Ok(mongo_documents_query_result(result.documents)) + } MongoCommand::Find { collection, filter, projection, sort, collation, skip, limit } => { let limit = bounded_mongo_find_limit(*limit, max_rows); let result = mongo_find_documents_without_total_core( @@ -1347,6 +1380,55 @@ for line in sys.stdin: assert!(!error.contains("Unknown method"), "{error}"); } + #[cfg(unix)] + #[tokio::test] + async fn mongo_run_command_routes_legacy_connections_to_the_agent() { + let expected_result = serde_json::json!({ + "documents": [{"ok": 1, "cursor": {"firstBatch": [{"_id": {"$oid": "507f1f77bcf86cd799439011"}}]}}], + "extended_documents": [{"ok": 1, "cursor": {"firstBatch": [{"_id": {"$oid": "507f1f77bcf86cd799439011"}}]}}], + "total": 1, + }); + let (state, _directory) = legacy_mongo_state_with_capabilities( + "run_command", + serde_json::json!({ + "database": "app", + "command_json": "{\"ping\":1}", + }), + expected_result, + &[AgentCapability::MongoRunCommand.as_str()], + ) + .await; + + let result = mongo_run_command_core(&state, "legacy", "app", "{\"ping\":1}").await.unwrap(); + + assert_eq!(result.total, 1); + assert_eq!(result.documents[0]["ok"], 1); + assert_eq!( + result.extended_documents.as_ref().unwrap()[0]["cursor"]["firstBatch"][0]["_id"], + serde_json::json!({"$oid": "507f1f77bcf86cd799439011"}) + ); + } + + #[cfg(unix)] + #[tokio::test] + async fn mongo_run_command_requires_an_explicit_legacy_agent_capability() { + let (state, _directory) = legacy_mongo_state_with_capabilities( + "run_command", + serde_json::json!({ + "database": "app", + "command_json": "{\"ping\":1}", + }), + serde_json::json!({"documents": [{"ok": 1}], "total": 1}), + &[], + ) + .await; + + let error = mongo_run_command_core(&state, "legacy", "app", "{\"ping\":1}").await.unwrap_err(); + + assert!(error.contains("upgrade or reinstall"), "{error}"); + assert!(!error.contains("Unknown method"), "{error}"); + } + #[cfg(unix)] #[tokio::test] async fn mongo_drop_collection_routes_legacy_connections_to_the_agent() { diff --git a/crates/dbx-core/src/mongo_shell.rs b/crates/dbx-core/src/mongo_shell.rs index 8dde2346dc..a52f81dfd5 100644 --- a/crates/dbx-core/src/mongo_shell.rs +++ b/crates/dbx-core/src/mongo_shell.rs @@ -8,6 +8,11 @@ pub enum MongoCommand { Version, #[serde(rename = "use")] Use { database: String }, + #[serde(rename = "runCommand")] + RunCommand { + #[serde(rename = "commandJson")] + command_json: String, + }, #[serde(rename = "createUser")] CreateUser { #[serde(rename = "userJson")] @@ -84,7 +89,8 @@ impl MongoCommand { pub fn is_mutating(&self) -> bool { matches!( self, - Self::CreateUser { .. } + Self::RunCommand { .. } + | Self::CreateUser { .. } | Self::Insert { .. } | Self::Update { .. } | Self::Delete { .. } @@ -98,7 +104,7 @@ impl MongoCommand { } pub fn is_dangerous(&self) -> bool { - matches!(self, Self::CreateUser { .. } | Self::DropCollection { .. }) + matches!(self, Self::RunCommand { .. } | Self::CreateUser { .. } | Self::DropCollection { .. }) || matches!(self, Self::DropIndexes { indexes: None, single: false, .. }) || matches!(self, Self::Aggregate { pipeline, .. } if aggregate_writes(pipeline)) } @@ -356,6 +362,19 @@ pub fn parse(input: &str) -> Result { if let Some(database) = parse_use_database(source) { return Ok(MongoCommand::Use { database }); } + if let Some((args, tail)) = database_method_call(source, "runCommand") { + if !tail.is_empty() || args.len() != 1 { + return Err("MongoDB runCommand() requires exactly one command document.".to_string()); + } + let command_json = normalized_json(&args[0])?; + let command = parse_json_value(&command_json) + .and_then(|value| value.as_object().cloned()) + .ok_or("MongoDB runCommand() requires a command document.")?; + if command.is_empty() { + return Err("MongoDB runCommand() requires a non-empty command document.".to_string()); + } + return Ok(MongoCommand::RunCommand { command_json }); + } if let Some((args, tail)) = database_method_call(source, "createUser") { if !tail.is_empty() || !(1..=2).contains(&args.len()) { return Err("MongoDB createUser() requires a user document and optional write concern.".to_string()); @@ -1022,6 +1041,43 @@ mod tests { assert!(parse(r#"db.createUser({user: "test"}, "majority")"#).is_err()); } + #[test] + fn parses_run_command_as_a_dangerous_write() { + let command = parse( + r#"db.runCommand({ + find: "orders", + filter: {_id: ObjectId("507f1f77bcf86cd799439011")}, + createdAt: ISODate("2025-01-01T00:00:00Z") + })"#, + ) + .unwrap(); + assert_eq!( + command, + MongoCommand::RunCommand { + command_json: r#"{"find":"orders","filter":{"_id":{"$oid":"507f1f77bcf86cd799439011"}},"createdAt":{"$date":"2025-01-01T00:00:00Z"}}"#.to_string(), + } + ); + assert!(command.is_mutating()); + assert!(command.is_dangerous()); + assert_eq!(validate_safety(&command, false, false, false), Err(MongoSafetyError::WritesDisabled)); + assert_eq!(validate_safety(&command, true, false, false), Err(MongoSafetyError::Dangerous)); + assert_eq!(validate_safety(&command, true, true, false), Ok(())); + } + + #[test] + fn rejects_unsupported_run_command_shapes() { + for source in [ + "db.runCommand()", + "db.runCommand({})", + "db.runCommand('ping')", + "db.runCommand({ping: 1}, {readPreference: 'primary'})", + "db.runCommand({ping: 1}).valueOf()", + "db.runCommand([1, 2, 3])", + ] { + assert!(parse(source).unwrap_err().contains("runCommand"), "{source}"); + } + } + #[test] fn treats_effectively_unbounded_write_filters_as_dangerous() { for command in [ @@ -1171,6 +1227,9 @@ mod tests { serde_json::to_value(parse(r#"db.createUser({user: "app", pwd: "secret", roles: []})"#).unwrap()).unwrap(); assert_eq!(create_user["kind"], "createUser"); assert_eq!(create_user["userJson"], r#"{"user":"app","pwd":"secret","roles":[]}"#); + let run_command = serde_json::to_value(parse("db.runCommand({ping: 1})").unwrap()).unwrap(); + assert_eq!(run_command["kind"], "runCommand"); + assert_eq!(run_command["commandJson"], r#"{"ping":1}"#); } #[test] diff --git a/crates/dbx-core/tests/live_mongodb_run_command.rs b/crates/dbx-core/tests/live_mongodb_run_command.rs new file mode 100644 index 0000000000..0796b3ceb3 --- /dev/null +++ b/crates/dbx-core/tests/live_mongodb_run_command.rs @@ -0,0 +1,45 @@ +//! Live MongoDB `runCommand` regression. +//! +//! Run with a reachable MongoDB server: +//! ```text +//! DBX_LIVE_MONGODB_URL='mongodb://127.0.0.1:27017' \ +//! cargo test -p dbx-core --test live_mongodb_run_command -- --ignored --nocapture +//! ``` + +use std::time::Duration; + +use dbx_core::db::mongo_driver; + +#[tokio::test] +#[ignore = "requires DBX_LIVE_MONGODB_URL pointing at a MongoDB database"] +async fn run_command_returns_browser_and_extended_json_documents() { + let url = std::env::var("DBX_LIVE_MONGODB_URL").expect("DBX_LIVE_MONGODB_URL"); + let client = mongo_driver::connect(&url, Duration::from_secs(10), Duration::from_secs(60)).await.unwrap(); + let database = std::env::var("DBX_LIVE_MONGODB_DATABASE").unwrap_or_else(|_| "admin".to_string()); + + let result = mongo_driver::run_command(&client, &database, r#"{"ping":1,"comment":"DBX #3050"}"#).await.unwrap(); + + assert_eq!(result.total, 1); + assert!(result.total_is_exact); + assert_eq!(result.documents.len(), 1); + let extended_documents = result.extended_documents.as_ref().expect("extended JSON documents"); + assert_eq!(extended_documents.len(), 1); + assert_eq!(result.documents[0]["ok"].as_f64(), Some(1.0)); + assert_eq!(extended_documents[0]["ok"]["$numberDouble"].as_str(), Some("1.0")); + + let build_info = mongo_driver::run_command(&client, &database, r#"{"buildInfo":1}"#).await.unwrap(); + assert!(build_info.documents[0]["versionArray"].as_array().is_some()); + assert!(build_info.extended_documents.as_ref().unwrap()[0]["versionArray"].as_array().is_some()); + + let collection = format!("dbx_run_command_{}", std::process::id()); + let create = + mongo_driver::run_command(&client, &database, &serde_json::json!({ "create": &collection }).to_string()) + .await + .unwrap(); + assert_eq!(create.documents[0]["ok"].as_f64(), Some(1.0)); + + let drop = mongo_driver::run_command(&client, &database, &serde_json::json!({ "drop": &collection }).to_string()) + .await + .unwrap(); + assert_eq!(drop.documents[0]["ok"].as_f64(), Some(1.0)); +} diff --git a/crates/dbx-mcp/src/backend.rs b/crates/dbx-mcp/src/backend.rs index 3c937ce147..84e980c12d 100644 --- a/crates/dbx-mcp/src/backend.rs +++ b/crates/dbx-mcp/src/backend.rs @@ -972,6 +972,9 @@ impl DbxBackend for WebBackend { Ok(scalar_query_result("version", Value::String(version))) } MongoCommand::Use { database } => Ok(scalar_query_result("database", Value::String(database.clone()))), + MongoCommand::RunCommand { .. } => { + Err("MongoDB runCommand is not available through the DBX MCP backend".to_string()) + } MongoCommand::Find { collection, filter, projection, sort, collation, skip, limit } => { let result = self .request( diff --git a/crates/dbx-mcp/src/server.rs b/crates/dbx-mcp/src/server.rs index 7a83b18f3d..3fdfe370be 100644 --- a/crates/dbx-mcp/src/server.rs +++ b/crates/dbx-mcp/src/server.rs @@ -1059,6 +1059,12 @@ fn validate_mongo_command( ), ) })?; + if matches!(command, MongoCommand::RunCommand { .. }) { + return Err(tool_error( + "SQL_BLOCKED", + "MongoDB runCommand is not available through MCP; review and execute it manually in DBX.", + )); + } let permissions = mcp_permissions(connection, policy); let production_database = match &command { MongoCommand::Aggregate { pipeline, .. } => { @@ -1498,6 +1504,22 @@ mod tests { ); } + #[test] + fn mongo_run_command_is_never_exposed_through_mcp() { + let mongo = connection("mongo", "mongo", "mongodb", "staging"); + let source = r#"db.runCommand({ping: 1})"#; + + for policy in [ + McpGlobalPolicy { read_only: true, allow_dangerous_sql: false, allowed_connection_ids: None }, + McpGlobalPolicy { read_only: false, allow_dangerous_sql: false, allowed_connection_ids: None }, + McpGlobalPolicy { read_only: false, allow_dangerous_sql: true, allowed_connection_ids: None }, + ] { + let error = validate_mongo_command(&mongo, &policy, "staging", source).unwrap_err(); + assert!(result_text(&error).contains("SQL_BLOCKED")); + assert!(result_text(&error).contains("runCommand")); + } + } + #[test] fn agent_results_preserve_stable_backend_policy_errors() { let result = agent_result(dbx_core::agent_events::ToolResult { diff --git a/crates/dbx-web/src/main.rs b/crates/dbx-web/src/main.rs index 9434b02344..1715ee1e4d 100644 --- a/crates/dbx-web/src/main.rs +++ b/crates/dbx-web/src/main.rs @@ -785,6 +785,7 @@ async fn main() { .route("/mongo/distinct", post(routes::mongo::distinct)) .route("/mongo/create-index", post(routes::mongo::create_index)) .route("/mongo/create-user", post(routes::mongo::create_user)) + .route("/mongo/run-command", post(routes::mongo::run_command)) .route("/mongo/drop-indexes", post(routes::mongo::drop_indexes)) .route("/mongo/insert-document", post(routes::mongo::insert_document)) .route("/mongo/insert-documents", post(routes::mongo::insert_documents)) diff --git a/crates/dbx-web/src/routes/mongo.rs b/crates/dbx-web/src/routes/mongo.rs index 9065964e4d..0efc6d1152 100644 --- a/crates/dbx-web/src/routes/mongo.rs +++ b/crates/dbx-web/src/routes/mongo.rs @@ -209,6 +209,15 @@ pub struct MongoCreateUserRequest { pub write_concern_json: Option, } +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct MongoRunCommandRequest { + pub connection_id: String, + pub database: String, + pub command_json: String, + pub execution_id: Option, +} + #[derive(Deserialize)] #[serde(rename_all = "camelCase")] pub struct MongoDropIndexesRequest { @@ -674,6 +683,29 @@ pub async fn create_user( Ok(Json(serde_json::json!({ "affected_rows": affected_rows }))) } +pub async fn run_command( + State(state): State>, + headers: HeaderMap, + Json(req): Json, +) -> Result, AppError> { + super::mcp_policy::ensure_dangerous_write( + &state, + &headers, + &req.connection_id, + &req.database, + "Run MongoDB command", + ) + .await?; + ensure_writable(&state.app, &req.connection_id, "Run MongoDB command").await?; + let result = run_cancellable( + &state, + req.execution_id.clone(), + dbx_core::mongo_ops::mongo_run_command_core(&state.app, &req.connection_id, &req.database, &req.command_json), + ) + .await?; + Ok(Json(serde_json::to_value(result).map_err(|e| AppError::from(e.to_string()))?)) +} + pub async fn drop_indexes( State(state): State>, headers: HeaderMap, diff --git a/packages/app-tests/mongoShellCommand.test.ts b/packages/app-tests/mongoShellCommand.test.ts index 2aab1cd90a..0be8a33ce4 100644 --- a/packages/app-tests/mongoShellCommand.test.ts +++ b/packages/app-tests/mongoShellCommand.test.ts @@ -17,6 +17,7 @@ import { parseMongoCollectionStatsCommand, parseMongoCommand, parseMongoCreateUserCommand, + parseMongoRunCommand, parseMongoCountDocumentsCommand, parseMongoDistinctCommand, parseMongoFindCommand, @@ -75,7 +76,10 @@ db.createUser({ roles: [{ role: "readWrite", db: "db1" }] })`); - assert.deepEqual(commands.map(({ command }) => command.kind), ["use", "createUser"]); + assert.deepEqual( + commands.map(({ command }) => command.kind), + ["use", "createUser"], + ); const createUser = commands[1]?.command; assert.equal(createUser?.kind, "createUser"); if (createUser?.kind === "createUser") { @@ -85,6 +89,50 @@ db.createUser({ } }); +test("parseMongoRunCommand normalizes one non-empty command document", () => { + assert.deepEqual( + parseMongoRunCommand(`db.runCommand({ + find: "orders", + filter: {_id: ObjectId("507f1f77bcf86cd799439011")}, + createdAt: ISODate("2025-01-01T00:00:00Z") + })`), + { + commandJson: '{"find":"orders","filter":{"_id":{"$oid":"507f1f77bcf86cd799439011"}},"createdAt":{"$date":"2025-01-01T00:00:00Z"}}', + }, + ); + + for (const source of ["db.runCommand()", "db.runCommand({})", "db.runCommand('ping')", "db.runCommand({ping: 1}, {readPreference: 'primary'})", "db.runCommand({ping: 1}).valueOf()", "db.runCommand([1, 2, 3])"]) { + assert.equal(parseMongoRunCommand(source), null, source); + } +}); + +test("runCommand support does not accept arbitrary Mongo shell JavaScript", () => { + assert.deepEqual( + splitMongoCommands(`for (let i = 0; i < 2; i += 1) { + db.items.insertOne({ index: i }); +}`), + [], + ); +}); + +test("splitMongoCommands keeps runCommand after a database switch", () => { + const commands = splitMongoCommands(`use admin + +db.runCommand({ ping: 1 })`); + + assert.deepEqual( + commands.map(({ command }) => command.kind), + ["use", "runCommand"], + ); + const runCommand = commands[1]?.command; + assert.equal(runCommand?.kind, "runCommand"); + if (runCommand?.kind === "runCommand") { + assert.deepEqual(JSON.parse(runCommand.commandJson), { ping: 1 }); + assert.equal(evaluateMongoWriteSafety(runCommand, { allowWrites: true }).allowed, false); + assert.equal(evaluateMongoWriteSafety(runCommand, { allowWrites: true, allowDangerous: true }).allowed, true); + } +}); + test("parseMongoFindCommand parses getCollection find with chained sort skip and limit", () => { assert.deepEqual(parseMongoFindCommand('db.getCollection("audit.logs").find({"level":"warn"}).sort({"createdAt":-1}).skip(20).limit(10)'), { collection: "audit.logs", diff --git a/packages/app-tests/queryStore.test.ts b/packages/app-tests/queryStore.test.ts index 1079ead1aa..3970029c7b 100644 --- a/packages/app-tests/queryStore.test.ts +++ b/packages/app-tests/queryStore.test.ts @@ -4838,6 +4838,56 @@ db.createUser({ } }); +test("mongo runCommand execution follows use and preserves document results", async () => { + const restoreStorage = installMemoryStorage(); + setActivePinia(createPinia()); + const connectionStore = useConnectionStore(); + const store = useQueryStore(); + const originalFetch = globalThis.fetch; + let runCommandBody: any; + + connectionStore.addEphemeralConnection({ + ...conn("mongo-1"), + db_type: "mongodb", + port: 27017, + }); + + globalThis.fetch = withConnectionHealthMock(async (input, init) => { + if (String(input) === "/api/mongo/run-command") { + runCommandBody = JSON.parse(String(init?.body ?? "{}")); + return Response.json({ + documents: [{ ok: 1, mode: "primary" }], + extended_documents: [{ ok: { $numberInt: "1" }, mode: "primary" }], + total: 1, + total_is_exact: true, + }); + } + return new Response("unexpected request", { status: 500 }); + }); + + try { + const tabId = store.createTab("mongo-1", "accounting", "Query", "query", ""); + await store.executeTabSql(tabId, 'use admin\n\ndb.runCommand({ hello: 1, comment: "DBX #3050" })'); + const tab = store.tabs.find((item) => item.id === tabId); + + assert.deepEqual(runCommandBody, { + connectionId: "mongo-1", + database: "admin", + commandJson: '{"hello":1,"comment":"DBX #3050"}', + executionId: runCommandBody.executionId, + }); + assert.equal(typeof runCommandBody.executionId, "string"); + assert.equal(tab?.database, "admin"); + assert.equal(tab?.results?.length, 2); + assert.deepEqual(tab?.results?.[1]?.columns, ["ok", "mode"]); + assert.deepEqual(tab?.results?.[1]?.rows, [[1, "primary"]]); + assert.deepEqual(tab?.results?.[1]?.mongo_copy_documents, [{ ok: { $numberInt: "1" }, mode: "primary" }]); + } finally { + globalThis.fetch = originalFetch; + restoreStorage(); + } +}); + test("mongo dropIndex execution uses the dedicated drop-indexes endpoint", async () => { const restoreStorage = installMemoryStorage(); setActivePinia(createPinia()); diff --git a/packages/mongo-shell/src/aggregate.ts b/packages/mongo-shell/src/aggregate.ts index 69a8f03c67..9386e3eda7 100644 --- a/packages/mongo-shell/src/aggregate.ts +++ b/packages/mongo-shell/src/aggregate.ts @@ -25,7 +25,7 @@ export const MONGO_SHELL_COMMAND_HINT = "Use MongoDB shell-style commands, for example: db.collection.find({}).limit(100), " + "db.collection.aggregate([]), db.collection.aggregate([], { explain: true }), " + 'db.version(), db.collection.countDocuments({}), db.collection.distinct("field"), ' + - "db.collection.getIndexes(), db.collection.createIndex({...}), db.createUser({...}), " + + "db.collection.getIndexes(), db.collection.createIndex({...}), db.createUser({...}), db.runCommand({...}), " + "or db.collection.insertOne({...})."; const PIPELINE_MUST_BE_ARRAY = diff --git a/src-tauri/src/commands/mongo_cmd.rs b/src-tauri/src/commands/mongo_cmd.rs index 94f82636f2..d46616f8fa 100644 --- a/src-tauri/src/commands/mongo_cmd.rs +++ b/src-tauri/src/commands/mongo_cmd.rs @@ -383,6 +383,34 @@ pub async fn mongo_create_user( Ok(serde_json::json!({ "affected_rows": affected_rows })) } +#[tauri::command] +pub async fn mongo_run_command( + state: State<'_, Arc>, + connection_id: String, + database: String, + command_json: String, + execution_id: Option, + mcp_request: Option, +) -> Result { + if mcp_request == Some(true) { + crate::commands::mcp_bridge::ensure_mcp_dangerous_write_allowed_by_id( + state.inner(), + &connection_id, + &database, + "Run MongoDB command", + ) + .await?; + } + ensure_connection_writable(&state, &connection_id, "Run MongoDB command").await?; + let app = state.inner().clone(); + run_cancellable( + &app, + execution_id, + dbx_core::mongo_ops::mongo_run_command_core(&app, &connection_id, &database, &command_json), + ) + .await +} + #[tauri::command] pub async fn mongo_drop_indexes( state: State<'_, Arc>, diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 81dc1c201c..f04c1997ce 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -2001,6 +2001,7 @@ pub fn run() { commands::mongo_cmd::mongo_distinct, commands::mongo_cmd::mongo_create_index, commands::mongo_cmd::mongo_create_user, + commands::mongo_cmd::mongo_run_command, commands::mongo_cmd::mongo_drop_indexes, commands::document_cmd::document_insert_document, commands::mongo_cmd::mongo_insert_document,