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 474fee0935..1b6e849fef 100644 --- a/agents/common/src/main/java/com/dbx/agent/AgentProtocol.java +++ b/agents/common/src/main/java/com/dbx/agent/AgentProtocol.java @@ -74,6 +74,7 @@ public final class AgentProtocol { 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_BULK_WRITE = "bulk_write"; 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"; @@ -132,6 +133,7 @@ public final class AgentProtocol { 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_MONGO_BULK_WRITE = "mongo_bulk_write"; public static final String CAPABILITY_MONGO_FIND_CURSOR = "mongo_find_cursor"; public static final String CAPABILITY_MULTI_SESSION = "multi_session"; public static final String CAPABILITY_STRUCTURED_ERROR_V1 = "structured_error_v1"; @@ -170,6 +172,7 @@ public final class AgentProtocol { CAPABILITY_MONGO_RUN_COMMAND, CAPABILITY_MONGO_INSERT_DOCUMENTS, CAPABILITY_MONGO_REPLACE_DOCUMENT, + CAPABILITY_MONGO_BULK_WRITE, CAPABILITY_MONGO_FIND_CURSOR )); @@ -235,6 +238,7 @@ public final class AgentProtocol { mongoCapabilities.add(CAPABILITY_MONGO_RUN_COMMAND); mongoCapabilities.add(CAPABILITY_MONGO_INSERT_DOCUMENTS); mongoCapabilities.add(CAPABILITY_MONGO_REPLACE_DOCUMENT); + mongoCapabilities.add(CAPABILITY_MONGO_BULK_WRITE); mongoCapabilities.add(CAPABILITY_MONGO_FIND_CURSOR); MONGO_LEGACY_CAPABILITIES = Collections.unmodifiableList(mongoCapabilities); @@ -244,6 +248,7 @@ public final class AgentProtocol { mongoMultiSessionCapabilities.add(CAPABILITY_MONGO_RUN_COMMAND); mongoMultiSessionCapabilities.add(CAPABILITY_MONGO_INSERT_DOCUMENTS); mongoMultiSessionCapabilities.add(CAPABILITY_MONGO_REPLACE_DOCUMENT); + mongoMultiSessionCapabilities.add(CAPABILITY_MONGO_BULK_WRITE); mongoMultiSessionCapabilities.add(CAPABILITY_MONGO_FIND_CURSOR); MONGO_LEGACY_MULTI_SESSION_CAPABILITIES = Collections.unmodifiableList(mongoMultiSessionCapabilities); @@ -287,6 +292,7 @@ public final class AgentProtocol { MONGO_METHOD_UPDATE_DOCUMENT, MONGO_METHOD_UPDATE_DOCUMENTS, MONGO_METHOD_REPLACE_DOCUMENT, + MONGO_METHOD_BULK_WRITE, 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 abb7ebbc41..1df74f9a48 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", "mongo_replace_document", "mongo_find_cursor"], + "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", "mongo_bulk_write", "mongo_find_cursor"], "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", "replace_document", "delete_document", "delete_documents", "run_command", "start_find_cursor", "fetch_find_cursor", "close_find_cursor"], + "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", "bulk_write", "delete_document", "delete_documents", "run_command", "start_find_cursor", "fetch_find_cursor", "close_find_cursor"], "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 c1ef26af18..8ffe635f9e 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", "mongo_replace_document", "mongo_find_cursor", "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", "mongo_bulk_write", "mongo_find_cursor", "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", "replace_document", "delete_document", "delete_documents", "run_command", "start_find_cursor", "fetch_find_cursor", "close_find_cursor"], + "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", "bulk_write", "delete_document", "delete_documents", "run_command", "start_find_cursor", "fetch_find_cursor", "close_find_cursor"], "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 3f6f1b94bc..3be93548c5 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 @@ -28,10 +28,21 @@ import com.mongodb.client.model.CollationAlternate; import com.mongodb.client.model.CollationCaseFirst; import com.mongodb.client.model.CollationMaxVariable; import com.mongodb.client.model.CollationStrength; +import com.mongodb.MongoBulkWriteException; +import com.mongodb.bulk.BulkWriteError; +import com.mongodb.bulk.BulkWriteResult; +import com.mongodb.client.model.BulkWriteOptions; import com.mongodb.client.model.CountOptions; +import com.mongodb.client.model.DeleteManyModel; +import com.mongodb.client.model.DeleteOneModel; +import com.mongodb.client.model.InsertOneModel; import com.mongodb.client.model.InsertManyOptions; +import com.mongodb.client.model.ReplaceOneModel; import com.mongodb.client.model.ReplaceOptions; +import com.mongodb.client.model.UpdateManyModel; +import com.mongodb.client.model.UpdateOneModel; import com.mongodb.client.model.UpdateOptions; +import com.mongodb.client.model.WriteModel; import com.mongodb.client.result.UpdateResult; import java.io.BufferedReader; import java.io.FileInputStream; @@ -1884,6 +1895,168 @@ public final class MongoAgent { return Collections.singletonMap("modified_count", result.getModifiedCount()); } + private static Object bulkWrite(JsonObject params) { + MongoClient c = requireClient(); + String database = params.get("database").getAsString(); + String collection = params.get("collection").getAsString(); + String operationsJson = params.get("operations_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); + List> models = bulkWriteModelsForWrite(operationsJson); + BulkWriteOptions options = bulkWriteOptionsForWrite(optionsJson); + try { + return bulkWriteCounts(col.bulkWrite(models, options)); + } catch (MongoBulkWriteException error) { + // Report what did complete alongside the failures, as the native driver does. + StringBuilder failures = new StringBuilder(); + for (BulkWriteError writeError : error.getWriteErrors()) { + if (failures.length() > 0) { + failures.append("; "); + } + failures.append("operation ").append(writeError.getIndex() + 1).append(": ").append(writeError.getMessage()); + } + Map completed = bulkWriteCounts(error.getWriteResult()); + throw new IllegalArgumentException( + "bulkWrite " + (options.isOrdered() ? "stopped at " : "finished with failed ") + + failures + ". Completed: inserted " + completed.get("inserted_count") + + ", matched " + completed.get("matched_count") + ", modified " + completed.get("modified_count") + + ", deleted " + completed.get("deleted_count") + ", upserted " + completed.get("upserted_count")); + } + } + + private static Map bulkWriteCounts(BulkWriteResult result) { + Map counts = new LinkedHashMap<>(); + counts.put("inserted_count", result.getInsertedCount()); + counts.put("matched_count", result.getMatchedCount()); + counts.put("modified_count", result.getModifiedCount()); + counts.put("deleted_count", result.getDeletedCount()); + counts.put("upserted_count", result.getUpserts().size()); + return counts; + } + + /** The shell's `bulkWrite([...])` entries, each `{ : { ... } }`, as driver write models. */ + static List> bulkWriteModelsForWrite(String operationsJson) { + JsonElement parsed = JsonParser.parseString(operationsJson); + if (!parsed.isJsonArray() || parsed.getAsJsonArray().isEmpty()) { + throw new IllegalArgumentException("bulkWrite requires a non-empty array of operations"); + } + List> models = new ArrayList<>(); + int position = 0; + for (JsonElement entry : parsed.getAsJsonArray()) { + position++; + if (!entry.isJsonObject() || entry.getAsJsonObject().size() != 1) { + throw new IllegalArgumentException("bulkWrite operation " + position + " must have exactly one operation key"); + } + String kind = entry.getAsJsonObject().keySet().iterator().next(); + JsonElement rawSpec = entry.getAsJsonObject().get(kind); + if (!rawSpec.isJsonObject()) { + throw new IllegalArgumentException("bulkWrite operation " + position + " (" + kind + ") must be a document"); + } + Document spec = documentForWrite(rawSpec.toString()); + models.add(bulkWriteModel(position, kind, spec)); + } + return models; + } + + private static WriteModel bulkWriteModel(int position, String kind, Document spec) { + switch (kind) { + case "insertOne": + return new InsertOneModel<>(requireBulkField(position, kind, spec, "document")); + case "updateOne": + case "updateMany": { + Document filter = requireBulkField(position, kind, spec, "filter"); + Object rawUpdate = spec.get("update"); + UpdateOptions options = new UpdateOptions(); + if (spec.containsKey("upsert")) { + options.upsert(requireBulkBoolean(position, kind, spec, "upsert")); + } + if (spec.get("arrayFilters") instanceof List filters) { + List arrayFilters = new ArrayList<>(); + for (Object filterEntry : filters) { + if (!(filterEntry instanceof Document)) { + throw new IllegalArgumentException("bulkWrite operation " + position + " arrayFilters entries must be documents"); + } + arrayFilters.add((Document) filterEntry); + } + options.arrayFilters(arrayFilters); + } + if (rawUpdate instanceof List stages) { + List pipeline = new ArrayList<>(); + for (Object stage : stages) { + if (!(stage instanceof Document)) { + throw new IllegalArgumentException("bulkWrite operation " + position + " (" + kind + ") pipeline stages must be documents"); + } + pipeline.add((Document) stage); + } + if (pipeline.isEmpty()) { + throw new IllegalArgumentException("bulkWrite operation " + position + " (" + kind + ") update pipeline must not be empty"); + } + return kind.equals("updateOne") ? new UpdateOneModel<>(filter, pipeline, options) : new UpdateManyModel<>(filter, pipeline, options); + } + if (!(rawUpdate instanceof Document update)) { + throw new IllegalArgumentException("bulkWrite operation " + position + " (" + kind + ") requires an update document or pipeline"); + } + requireBulkUpdateOperatorDocument(update); + return kind.equals("updateOne") ? new UpdateOneModel<>(filter, update, options) : new UpdateManyModel<>(filter, update, options); + } + case "replaceOne": { + Document replacement = requireBulkField(position, kind, spec, "replacement"); + requireReplacementDocument(replacement); + ReplaceOptions options = new ReplaceOptions(); + if (spec.containsKey("upsert")) { + options.upsert(requireBulkBoolean(position, kind, spec, "upsert")); + } + return new ReplaceOneModel<>(requireBulkField(position, kind, spec, "filter"), replacement, options); + } + case "deleteOne": + return new DeleteOneModel<>(requireBulkField(position, kind, spec, "filter")); + case "deleteMany": + return new DeleteManyModel<>(requireBulkField(position, kind, spec, "filter")); + default: + throw new IllegalArgumentException("bulkWrite operation " + position + " uses unsupported operation " + kind); + } + } + + private static Document requireBulkField(int position, String kind, Document spec, String field) { + Object value = spec.get(field); + if (!(value instanceof Document)) { + throw new IllegalArgumentException("bulkWrite operation " + position + " (" + kind + ") requires a " + field + " document"); + } + return (Document) value; + } + + private static boolean requireBulkBoolean(int position, String kind, Document spec, String field) { + Object value = spec.get(field); + if (!(value instanceof Boolean)) { + throw new IllegalArgumentException("bulkWrite operation " + position + " (" + kind + ") " + field + " must be a boolean"); + } + return (Boolean) value; + } + + static BulkWriteOptions bulkWriteOptionsForWrite(String optionsJson) { + BulkWriteOptions result = new BulkWriteOptions(); + if (optionsJson == null || optionsJson.trim().isEmpty()) { + return result; + } + Document options = Document.parse(optionsJson); + for (String key : options.keySet()) { + if (!"ordered".equals(key)) { + throw new IllegalArgumentException("Unsupported bulkWrite option: " + key); + } + } + Object ordered = options.get("ordered"); + if (ordered != null) { + if (!(ordered instanceof Boolean)) { + throw new IllegalArgumentException("ordered must be a boolean"); + } + result.ordered((Boolean) ordered); + } + return result; + } + /** replaceOne swaps the whole document; update operators here mean updateOne was intended. */ static void requireReplacementDocument(Document doc) { for (String key : doc.keySet()) { @@ -2106,6 +2279,7 @@ public final class MongoAgent { 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_BULK_WRITE -> bulkWrite(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 82c0fead96..a33e6f4a8b 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 @@ -32,6 +32,13 @@ import com.mongodb.client.model.CollationStrength; import com.mongodb.client.model.CountOptions; import com.mongodb.client.model.InsertManyOptions; import com.mongodb.client.model.UpdateOptions; +import com.mongodb.client.model.DeleteManyModel; +import com.mongodb.client.model.DeleteOneModel; +import com.mongodb.client.model.InsertOneModel; +import com.mongodb.client.model.ReplaceOneModel; +import com.mongodb.client.model.UpdateManyModel; +import com.mongodb.client.model.UpdateOneModel; +import com.mongodb.client.model.WriteModel; import com.mongodb.client.result.UpdateResult; import java.io.FileInputStream; import java.lang.reflect.Proxy; @@ -1329,6 +1336,68 @@ class MongoAgentTest { assertTrue(AgentProtocol.MONGO_LEGACY_CAPABILITIES.contains(AgentProtocol.CAPABILITY_MONGO_REPLACE_DOCUMENT)); } + @Test + void bulkWriteMethodIsRecognizedOverJsonRpc() { + String response = MongoAgent.handleRequest( + "{\"jsonrpc\":\"2.0\",\"id\":12,\"method\":\"bulk_write\"," + + "\"params\":{\"database\":\"app\",\"collection\":\"orders\"," + + "\"operations_json\":\"[{\\\"insertOne\\\":{\\\"document\\\":{\\\"a\\\":1}}}]\"}}"); + + JsonObject json = JsonParser.parseString(response).getAsJsonObject(); + assertEquals(12, json.get("id").getAsInt()); + assertEquals("Not connected", json.getAsJsonObject("error").get("message").getAsString()); + assertTrue(AgentProtocol.MONGO_LEGACY_METHODS.contains(AgentProtocol.MONGO_METHOD_BULK_WRITE)); + assertTrue(AgentProtocol.MONGO_LEGACY_CAPABILITIES.contains(AgentProtocol.CAPABILITY_MONGO_BULK_WRITE)); + } + + @Test + void buildsBulkWriteModelsForEveryOperationKind() { + List> models = MongoAgent.bulkWriteModelsForWrite( + "[{\"insertOne\":{\"document\":{\"sku\":\"A1\"}}}," + + "{\"updateOne\":{\"filter\":{\"sku\":\"A1\"},\"update\":{\"$inc\":{\"stock\":1}},\"upsert\":true}}," + + "{\"updateMany\":{\"filter\":{},\"update\":[{\"$set\":{\"stock\":0}}]}}," + + "{\"replaceOne\":{\"filter\":{\"sku\":\"B2\"},\"replacement\":{\"sku\":\"B2\"}}}," + + "{\"deleteOne\":{\"filter\":{\"sku\":\"C3\"}}}," + + "{\"deleteMany\":{\"filter\":{\"stock\":{\"$lt\":0}}}}]"); + + assertEquals(6, models.size()); + assertTrue(models.get(0) instanceof InsertOneModel); + assertTrue(models.get(1) instanceof UpdateOneModel); + assertTrue(((UpdateOneModel) models.get(1)).getOptions().isUpsert()); + assertTrue(models.get(2) instanceof UpdateManyModel); + assertNotNull(((UpdateManyModel) models.get(2)).getUpdatePipeline()); + assertTrue(models.get(3) instanceof ReplaceOneModel); + assertTrue(models.get(4) instanceof DeleteOneModel); + assertTrue(models.get(5) instanceof DeleteManyModel); + } + + @Test + void rejectsMalformedBulkWriteOperations() { + for (String[] item : new String[][] { + {"[]", "non-empty array"}, + {"[{\"insertOne\":{},\"deleteOne\":{}}]", "exactly one operation key"}, + {"[{\"upsertOne\":{\"document\":{}}}]", "unsupported operation upsertOne"}, + {"[{\"insertOne\":{}}]", "requires a document document"}, + {"[{\"updateOne\":{\"filter\":{},\"update\":{\"a\":1}}}]", "update operators such as $set"}, + {"[{\"updateOne\":{\"filter\":{},\"update\":{\"$set\":{\"a\":1}},\"upsert\":1}}]", "upsert must be a boolean"}, + {"[{\"replaceOne\":{\"filter\":{},\"replacement\":{\"$set\":{\"a\":1}}}}]", "must not contain update operators"}, + }) { + IllegalArgumentException error = assertThrows(IllegalArgumentException.class, () -> MongoAgent.bulkWriteModelsForWrite(item[0]), item[0]); + assertTrue(error.getMessage().contains(item[1]), item[0] + " -> " + error.getMessage()); + } + } + + @Test + void parsesBulkWriteOptions() { + assertTrue(MongoAgent.bulkWriteOptionsForWrite(null).isOrdered()); + assertTrue(MongoAgent.bulkWriteOptionsForWrite("{\"ordered\":true}").isOrdered()); + assertFalse(MongoAgent.bulkWriteOptionsForWrite("{\"ordered\":false}").isOrdered()); + IllegalArgumentException unsupported = assertThrows(IllegalArgumentException.class, () -> MongoAgent.bulkWriteOptionsForWrite("{\"writeConcern\":{}}")); + assertEquals("Unsupported bulkWrite option: writeConcern", unsupported.getMessage()); + IllegalArgumentException notBoolean = assertThrows(IllegalArgumentException.class, () -> MongoAgent.bulkWriteOptionsForWrite("{\"ordered\":\"yes\"}")); + assertEquals("ordered must be a boolean", notBoolean.getMessage()); + } + @Test void parsesReplaceOptions() { assertTrue(MongoAgent.replaceOptionsForWrite("{\"upsert\":true}").isUpsert()); diff --git a/apps/desktop/src/lib/backend/api.ts b/apps/desktop/src/lib/backend/api.ts index 238dad99ab..fa1024a2eb 100644 --- a/apps/desktop/src/lib/backend/api.ts +++ b/apps/desktop/src/lib/backend/api.ts @@ -823,6 +823,7 @@ export const documentUpdateDocument = forward("documentUpdateDocument"); export const mongoUpdateDocument = forward("mongoUpdateDocument"); export const mongoUpdateDocuments = forward("mongoUpdateDocuments"); export const mongoReplaceDocument = forward("mongoReplaceDocument"); +export const mongoBulkWrite = forward("mongoBulkWrite"); 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 c2388c4c9d..e3e3cb92fd 100644 --- a/apps/desktop/src/lib/backend/http.ts +++ b/apps/desktop/src/lib/backend/http.ts @@ -53,6 +53,7 @@ import type { } from "@/types/database"; import type { DetachedTabHandoff } from "@/lib/app/detachedTabHandoff"; import { normalizeRustMongoCommand, type MongoCommand } from "@/lib/mongo/mongoShellCommand"; +import type { MongoBulkWriteResult } from "@/lib/mongo/mongoShellCommand"; import { BackendErrorException, type BackendError } from "@/lib/backend/errorUtils"; import { decodeMeilisearchDocumentPage, decodeMeilisearchSearchResult, type MeilisearchDocumentPage, type MeilisearchDocumentPageWire, type MeilisearchSearchResult, type MeilisearchSearchWireResult } from "@/lib/backend/meilisearchTransport"; import type { CreatedKey, EnqueuedTaskSummary, KeyCreateInput, KeyListItem, KeyPage, KeyUpdateInput, MeilisearchSystemOverview, MeilisearchTask, TaskListInput, TaskPage, TaskSelector } from "@/types/meilisearchManagement"; @@ -4670,6 +4671,16 @@ export async function documentUpdateDocument(connectionId: string, database: str }); } +export async function mongoBulkWrite(connectionId: string, database: string, collection: string, operationsJson: string, optionsJson?: string): Promise { + return post("/api/mongo/bulk-write", { + connectionId, + database, + collection, + operationsJson, + optionsJson, + }); +} + 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, diff --git a/apps/desktop/src/lib/backend/tauri.ts b/apps/desktop/src/lib/backend/tauri.ts index 2f665e77a1..589797c735 100644 --- a/apps/desktop/src/lib/backend/tauri.ts +++ b/apps/desktop/src/lib/backend/tauri.ts @@ -12,6 +12,7 @@ import type { DetachedTabHandoff } from "@/lib/app/detachedTabHandoff"; import { BackendErrorException, type BackendError } from "@/lib/backend/errorUtils"; import { listen, type UnlistenFn } from "@tauri-apps/api/event"; import { normalizeRustMongoCommand, type MongoCommand } from "@/lib/mongo/mongoShellCommand"; +import type { MongoBulkWriteResult } from "@/lib/mongo/mongoShellCommand"; import { ExternalSqlFileTooLargeError } from "@/lib/sql/sqlFileOpen"; import { appendDebugLog, isDebugLoggingEnabled } from "@/lib/backend/debugLog"; import { decodeMeilisearchDocumentPage, decodeMeilisearchSearchResult, type MeilisearchDocumentPage, type MeilisearchDocumentPageWire, type MeilisearchSearchResult, type MeilisearchSearchWireResult } from "@/lib/backend/meilisearchTransport"; @@ -4445,6 +4446,16 @@ export async function documentUpdateDocument(connectionId: string, database: str }); } +export async function mongoBulkWrite(connectionId: string, database: string, collection: string, operationsJson: string, optionsJson?: string): Promise { + return invoke("mongo_bulk_write", { + connectionId, + database, + collection, + operationsJson, + optionsJson, + }); +} + 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, diff --git a/apps/desktop/src/lib/mongo/mongoCompletion.ts b/apps/desktop/src/lib/mongo/mongoCompletion.ts index bb2d77de37..335e880074 100644 --- a/apps/desktop/src/lib/mongo/mongoCompletion.ts +++ b/apps/desktop/src/lib/mongo/mongoCompletion.ts @@ -94,6 +94,7 @@ const COLLECTION_METHODS = [ { 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: "bulkWrite", detail: "Run several writes in one batch", apply: "bulkWrite([\n { insertOne: { document: {} } }\n])" }, { 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: {} })" }, @@ -122,6 +123,7 @@ const COLLECTION_METHOD_BOOST: Record<(typeof COLLECTION_METHODS)[number]["label updateOne: 160, updateMany: 150, replaceOne: 145, + bulkWrite: 135, deleteOne: 140, deleteMany: 130, findOneAndUpdate: 120, @@ -180,7 +182,7 @@ const ROOT_SNIPPET_BOOST: Record<(typeof ROOT_SNIPPETS)[number]["label"], number }; /** Role of each positional argument, by collection helper. Drives cursor classification. */ -type MongoArgRole = "filter" | "update" | "replacement" | "document" | "documents" | "pipeline" | "projection" | "keys" | "sortKeys" | "fieldName" | "options"; +type MongoArgRole = "filter" | "update" | "replacement" | "document" | "documents" | "operations" | "pipeline" | "projection" | "keys" | "sortKeys" | "fieldName" | "options"; const METHOD_ARG_ROLES: Record = { find: ["filter", "projection", "options"], @@ -193,6 +195,7 @@ const METHOD_ARG_ROLES: Record = { updateOne: ["filter", "update", "options"], updateMany: ["filter", "update", "options"], replaceOne: ["filter", "replacement", "options"], + bulkWrite: ["operations", "options"], findOneAndUpdate: ["filter", "update", "options"], findOneAndReplace: ["filter", "replacement", "options"], insertOne: ["document", "options"], @@ -538,6 +541,8 @@ function classifyCursorInCall(method: string, scan: MongoCallScan): MongoCursorC return { mode: scan.stack.length === 0 ? "fieldPath" : "none" }; case "pipeline": return classifyPipeline(scan); + // bulkWrite operations are `{ : { filter, update, … } }` entries; completing inside them is a follow-up. + case "operations": case "options": return { mode: "none" }; default: diff --git a/apps/desktop/src/lib/mongo/mongoShellCommand.ts b/apps/desktop/src/lib/mongo/mongoShellCommand.ts index c758151d7f..85f89dd56d 100644 --- a/apps/desktop/src/lib/mongo/mongoShellCommand.ts +++ b/apps/desktop/src/lib/mongo/mongoShellCommand.ts @@ -49,6 +49,7 @@ const COLLECTION_METHOD_SHAPES: Record = { 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"] }, + bulkWrite: { expects: "an array of operations and optional options", roles: ["operations", "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"] }, @@ -78,6 +79,7 @@ const SUPPORTED_COLLECTION_METHODS = [ "updateOne", "updateMany", "replaceOne", + "bulkWrite", "deleteOne", "deleteMany", "findOneAndUpdate", @@ -159,6 +161,14 @@ 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 === "bulkWrite" && args[0]) { + const operations = normalizeJsonArgument(args[0]); + const problem = operations ? validateBulkWriteOperations(operations) : null; + if (problem) return problem; + const options = args[1]?.trim() ? parseMongoObjectArgument(args[1]) : null; + const optionProblem = options ? validateBulkWriteOptions(options) : null; + if (optionProblem) return optionProblem; + } 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; @@ -263,7 +273,7 @@ export interface MongoDistinctCommand { filter?: string; } -type MongoWriteKind = "runCommand" | "insert" | "update" | "replace" | "delete" | "createIndex" | "createUser" | "dropIndex" | "dropIndexes" | "dropCollection" | "findOneAndUpdate" | "findOneAndReplace" | "findOneAndDelete"; +type MongoWriteKind = "runCommand" | "insert" | "update" | "replace" | "bulkWrite" | "delete" | "createIndex" | "createUser" | "dropIndex" | "dropIndexes" | "dropCollection" | "findOneAndUpdate" | "findOneAndReplace" | "findOneAndDelete"; export type MongoCommand = | ({ kind: "find" } & MongoFindCommand) @@ -281,6 +291,7 @@ export type MongoCommand = | { 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: "bulkWrite"; collection: string; operations: 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 } @@ -752,6 +763,17 @@ export function parseMongoWriteCommand(input: string): MongoWriteCommand | null return value !== null && typeof value === "object" ? { kind: "insert", collection: insert.collection, docsJson: docs } : null; } + const bulkWrite = parseCollectionMethodTarget(source, "bulkWrite"); + if (bulkWrite) { + const args = parseMethodArgs(source, bulkWrite.methodCallIndex); + if (!args || args.length < 1 || args.length > 2) return null; + const operations = normalizeJsonArgument(args[0]); + if (!operations || validateBulkWriteOperations(operations) !== null) return null; + const options = args[1]?.trim() ? parseMongoObjectArgument(args[1]) : undefined; + if (args[1]?.trim() && (!options || validateBulkWriteOptions(options) !== null)) return null; + return { kind: "bulkWrite", collection: bulkWrite.collection, operations, ...(options ? { options } : {}) }; + } + const replaceOne = parseCollectionMethodTarget(source, "replaceOne"); if (replaceOne) { const args = parseMethodArgs(source, replaceOne.methodCallIndex); @@ -937,7 +959,7 @@ export function evaluateMongoWriteSafety(command: MongoWriteCommand, options: Mo }; } const filter = mongoWriteFilter(command); - const highRisk = filter !== null ? mongoFilterIsEffectivelyUnbounded(filter) : command.kind !== "insert"; + const highRisk = command.kind === "bulkWrite" ? bulkWriteFilters(command.operations).some(mongoFilterIsEffectivelyUnbounded) : filter !== null ? mongoFilterIsEffectivelyUnbounded(filter) : command.kind !== "insert"; if (!options.allowDangerous && highRisk) { return { allowed: false, @@ -1058,6 +1080,24 @@ export function mongoWriteToQueryResult(affectedRows: number, executionTimeMs: n }; } +export interface MongoBulkWriteResult { + inserted_count: number; + matched_count: number; + modified_count: number; + deleted_count: number; + upserted_count: number; +} + +/** One row of counts, the way the shell prints a `BulkWriteResult`. */ +export function mongoBulkWriteToQueryResult(result: MongoBulkWriteResult, executionTimeMs: number): QueryResult { + return { + columns: ["insertedCount", "matchedCount", "modifiedCount", "deletedCount", "upsertedCount"], + rows: [[result.inserted_count, result.matched_count, result.modified_count, result.deleted_count, result.upserted_count]], + affected_rows: result.inserted_count + result.modified_count + result.deleted_count + result.upserted_count, + execution_time_ms: Math.max(0, Math.round(executionTimeMs)), + }; +} + export function mongoCreateIndexToQueryResult(name: string, executionTimeMs: number): QueryResult { return { columns: ["name"], @@ -1535,6 +1575,90 @@ function isNonEmptyRecord(value: unknown): value is Record { return isRecord(value) && Object.keys(value).length > 0; } +const BULK_WRITE_FIELDS: Record = { + insertOne: ["document"], + updateOne: ["filter", "update", "upsert", "arrayFilters"], + updateMany: ["filter", "update", "upsert", "arrayFilters"], + replaceOne: ["filter", "replacement", "upsert"], + deleteOne: ["filter"], + deleteMany: ["filter"], +}; + +const isDocument = (value: unknown): value is Record => value !== null && typeof value === "object" && !Array.isArray(value); + +/** + * Validate a `bulkWrite([...])` operations array the way the shell does, mirroring the + * Rust parser: each entry is one `{ : { ... } }` with exactly the fields it takes. + * Returns the problem, or null when every operation is well-formed. + */ +export function validateBulkWriteOperations(operationsJson: string): string | null { + let entries: unknown; + try { + entries = JSON.parse(operationsJson); + } catch { + return "bulkWrite() requires an array of operations."; + } + if (!Array.isArray(entries)) return "bulkWrite() requires an array of operations."; + if (entries.length === 0) return "bulkWrite() requires at least one operation."; + + for (const [index, entry] of entries.entries()) { + const position = index + 1; + if (!isDocument(entry)) return `bulkWrite() operation ${position} must be a document such as { insertOne: { document: { ... } } }.`; + const keys = Object.keys(entry); + if (keys.length !== 1) return `bulkWrite() operation ${position} must have exactly one operation key.`; + const kind = keys[0]!; + const spec = entry[kind]; + const allowed = BULK_WRITE_FIELDS[kind]; + if (!allowed) return `bulkWrite() operation ${position} uses unsupported operation ${kind}; supported: ${Object.keys(BULK_WRITE_FIELDS).join(", ")}.`; + if (!isDocument(spec)) return `bulkWrite() operation ${position} (${kind}) must be a document.`; + const unknown = Object.keys(spec).find((key) => !allowed.includes(key)); + if (unknown) return `bulkWrite() operation ${position} (${kind}) has unsupported field ${unknown}.`; + if ("upsert" in spec && typeof spec.upsert !== "boolean") return `bulkWrite() operation ${position} (${kind}) upsert must be a boolean.`; + + for (const field of allowed.filter((name) => name === "document" || name === "filter" || name === "replacement")) { + if (!(field in spec)) return `bulkWrite() operation ${position} (${kind}) requires a ${field} document.`; + if (!isDocument(spec[field])) return `bulkWrite() operation ${position} (${kind}) field ${field} must be a document.`; + } + if (kind === "updateOne" || kind === "updateMany") { + const update = spec.update; + if (isDocument(update)) { + const operatorKeys = Object.keys(update); + if (operatorKeys.length === 0 || !operatorKeys.every((key) => key.startsWith("$"))) { + return `bulkWrite() operation ${position} (${kind}) update must use operators such as $set; use replaceOne for a whole document.`; + } + } else if (!Array.isArray(update)) { + return update === undefined ? `bulkWrite() operation ${position} (${kind}) requires an update.` : `bulkWrite() operation ${position} (${kind}) update must be a document or pipeline.`; + } + if ("arrayFilters" in spec && !Array.isArray(spec.arrayFilters)) return `bulkWrite() operation ${position} (${kind}) arrayFilters must be an array.`; + } + if (kind === "replaceOne") { + const operator = Object.keys(spec.replacement as Record).find((key) => key.startsWith("$")); + if (operator) return `bulkWrite() operation ${position} (replaceOne) replacement must not contain update operators such as ${operator}; use updateOne to modify fields.`; + } + } + return null; +} + +/** Only `ordered` is honoured, so anything else is rejected rather than dropped. */ +export function validateBulkWriteOptions(optionsJson: string): string | null { + const options = JSON.parse(optionsJson) as Record; + for (const [key, value] of Object.entries(options)) { + if (key !== "ordered") return `Unsupported bulkWrite() option: ${key}.`; + if (typeof value !== "boolean") return "bulkWrite() ordered option must be a boolean."; + } + return null; +} + +/** Filters of every non-insert operation, for the safety checks. */ +function bulkWriteFilters(operationsJson: string): string[] { + try { + const entries = JSON.parse(operationsJson) as Array>; + return entries.flatMap((entry) => Object.values(entry)).flatMap((spec) => (isDocument(spec?.filter) ? [JSON.stringify(spec.filter)] : [])); + } catch { + return []; + } +} + function mongoWriteFilter(command: MongoWriteCommand): string | null { switch (command.kind) { case "update": diff --git a/apps/desktop/src/stores/queryStore.ts b/apps/desktop/src/stores/queryStore.ts index 5e0ff7aa94..5c0e11c4a5 100644 --- a/apps/desktop/src/stores/queryStore.ts +++ b/apps/desktop/src/stores/queryStore.ts @@ -30,6 +30,7 @@ import { planMongoFindPagination, mongoUseToQueryResult, mongoVersionToQueryResult, + mongoBulkWriteToQueryResult, mongoWriteToQueryResult, splitMongoCommandRanges, type MongoAggregateSafetyOptions, @@ -6647,6 +6648,7 @@ export const useQueryStore = defineStore("query", () => { case "insert": case "update": case "replace": + case "bulkWrite": case "delete": case "createIndex": case "createUser": @@ -6670,6 +6672,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 === "bulkWrite") { + const result = await api.mongoBulkWrite(executionConnectionId, currentDatabase, mongoCommand.collection, mongoCommand.operations, mongoCommand.options); + allResults.push(markQueryResultRowsRaw(annotateMongoResult(mongoBulkWriteToQueryResult(result, 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)))); diff --git a/crates/dbx-core/assets/agent-protocol-v1.json b/crates/dbx-core/assets/agent-protocol-v1.json index abb7ebbc41..1df74f9a48 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", "mongo_replace_document", "mongo_find_cursor"], + "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", "mongo_bulk_write", "mongo_find_cursor"], "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", "replace_document", "delete_document", "delete_documents", "run_command", "start_find_cursor", "fetch_find_cursor", "close_find_cursor"], + "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", "bulk_write", "delete_document", "delete_documents", "run_command", "start_find_cursor", "fetch_find_cursor", "close_find_cursor"], "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 c1ef26af18..8ffe635f9e 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", "mongo_replace_document", "mongo_find_cursor", "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", "mongo_bulk_write", "mongo_find_cursor", "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", "replace_document", "delete_document", "delete_documents", "run_command", "start_find_cursor", "fetch_find_cursor", "close_find_cursor"], + "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", "bulk_write", "delete_document", "delete_documents", "run_command", "start_find_cursor", "fetch_find_cursor", "close_find_cursor"], "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 bd4247d231..c82d12e3ee 100644 --- a/crates/dbx-core/src/db/agent_driver.rs +++ b/crates/dbx-core/src/db/agent_driver.rs @@ -1042,6 +1042,7 @@ pub enum AgentCapability { MongoRunCommand, MongoInsertDocuments, MongoReplaceDocument, + MongoBulkWrite, MongoFindCursor, MultiSession, StructuredErrorV1, @@ -1125,7 +1126,7 @@ fn parse_agent_rpc_error_header(header: &str) -> (Option, String) { } impl AgentCapability { - pub const ALL: [Self; 26] = [ + pub const ALL: [Self; 27] = [ Self::Connect, Self::TestConnection, Self::Metadata, @@ -1149,6 +1150,7 @@ impl AgentCapability { Self::MongoRunCommand, Self::MongoInsertDocuments, Self::MongoReplaceDocument, + Self::MongoBulkWrite, Self::MongoFindCursor, Self::MultiSession, Self::StructuredErrorV1, @@ -1179,6 +1181,7 @@ impl AgentCapability { Self::MongoRunCommand => "mongo_run_command", Self::MongoInsertDocuments => "mongo_insert_documents", Self::MongoReplaceDocument => "mongo_replace_document", + Self::MongoBulkWrite => "mongo_bulk_write", Self::MongoFindCursor => "mongo_find_cursor", Self::MultiSession => "multi_session", Self::StructuredErrorV1 => "structured_error_v1", @@ -1383,6 +1386,7 @@ pub enum MongoAgentMethod { UpdateDocument, UpdateDocuments, ReplaceDocument, + BulkWrite, DeleteDocument, DeleteDocuments, RunCommand, @@ -1392,7 +1396,7 @@ pub enum MongoAgentMethod { } impl MongoAgentMethod { - pub const ALL: [Self; 26] = [ + pub const ALL: [Self; 27] = [ Self::ListDatabases, Self::ListCollections, Self::FindDocuments, @@ -1413,6 +1417,7 @@ impl MongoAgentMethod { Self::UpdateDocument, Self::UpdateDocuments, Self::ReplaceDocument, + Self::BulkWrite, Self::DeleteDocument, Self::DeleteDocuments, Self::RunCommand, @@ -1443,6 +1448,7 @@ impl MongoAgentMethod { Self::UpdateDocument => "update_document", Self::UpdateDocuments => "update_documents", Self::ReplaceDocument => "replace_document", + Self::BulkWrite => "bulk_write", Self::DeleteDocument => "delete_document", Self::DeleteDocuments => "delete_documents", Self::RunCommand => "run_command", @@ -2849,6 +2855,10 @@ impl AgentDriverClient { self.call_mongo_method(MongoAgentMethod::ReplaceDocument, params).await } + pub async fn mongo_bulk_write(&mut self, params: Value) -> Result { + self.call_mongo_method(MongoAgentMethod::BulkWrite, params).await + } + pub async fn mongo_delete_document( &mut self, params: Value, @@ -3005,6 +3015,7 @@ pub fn agent_supports_capability(handshake: Option<&AgentHandshake>, capability: | AgentCapability::MongoRunCommand | AgentCapability::MongoInsertDocuments | AgentCapability::MongoReplaceDocument + | AgentCapability::MongoBulkWrite | AgentCapability::MongoFindCursor ) { return handshake.map(|value| value.supports(capability)).unwrap_or(false); @@ -5040,10 +5051,11 @@ for line in sys.stdin: 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::MongoBulkWrite.as_str(), "mongo_bulk_write"); assert_eq!(AgentCapability::MongoFindCursor.as_str(), "mongo_find_cursor"); assert_eq!(AgentCapability::MultiSession.as_str(), "multi_session"); assert_eq!(AgentCapability::StructuredErrorV1.as_str(), "structured_error_v1"); - assert_eq!(AgentCapability::ALL.len(), 26); + assert_eq!(AgentCapability::ALL.len(), 27); } #[test] @@ -5106,6 +5118,7 @@ for line in sys.stdin: 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::BulkWrite.as_str(), "bulk_write"); assert_eq!(MongoAgentMethod::DeleteDocument.as_str(), "delete_document"); assert_eq!(MongoAgentMethod::DeleteDocuments.as_str(), "delete_documents"); assert_eq!(MongoAgentMethod::StartFindCursor.as_str(), "start_find_cursor"); diff --git a/crates/dbx-core/src/db/mongo_driver.rs b/crates/dbx-core/src/db/mongo_driver.rs index 83b9785459..1e5dd94293 100644 --- a/crates/dbx-core/src/db/mongo_driver.rs +++ b/crates/dbx-core/src/db/mongo_driver.rs @@ -2120,6 +2120,153 @@ pub async fn update_documents( Ok(result.modified_count) } +/// Counts reported by `bulkWrite`, mirroring the shell's `BulkWriteResult`. +#[derive(Debug, Default, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct MongoBulkWriteResult { + pub inserted_count: u64, + pub matched_count: u64, + pub modified_count: u64, + pub deleted_count: u64, + pub upserted_count: u64, +} + +#[derive(Default, Deserialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +struct MongoBulkWriteOptions { + ordered: Option, +} + +/// bulkWrite(operations[, { ordered }]). +/// +/// The Rust driver's `bulk_write` is client-level and needs MongoDB 8.0, so the +/// operations run here one at a time in order, which is what the shell's +/// collection-level bulkWrite guarantees anyway (it is not transactional). With +/// `ordered: true` (the default) the first failure stops the batch; with +/// `ordered: false` every operation is attempted and the failures are reported together. +pub async fn bulk_write( + client: &Client, + database: &str, + collection: &str, + operations_json: &str, + options_json: Option<&str>, +) -> Result { + use crate::mongo_shell::BulkWriteOperation; + + let operations = crate::mongo_shell::parse_bulk_write_operations(operations_json)?; + let ordered = match options_json.filter(|value| !value.trim().is_empty()) { + Some(raw) => { + let options: MongoBulkWriteOptions = + serde_json::from_str(raw).map_err(|e| format!("Invalid bulkWrite options: {e}"))?; + options.ordered.unwrap_or(true) + } + None => true, + }; + + let col = client.database(database).collection::(collection); + let mut result = MongoBulkWriteResult::default(); + let mut failures = Vec::new(); + + for (index, operation) in operations.iter().enumerate() { + let outcome: Result<(), String> = async { + match operation { + BulkWriteOperation::InsertOne { document } => { + let document = json_object_to_document(document).map_err(|e| format!("Invalid document: {e}"))?; + col.insert_one(document).await.map_err(|e| e.to_string())?; + result.inserted_count += 1; + } + BulkWriteOperation::UpdateOne { filter, update, upsert, array_filters } + | BulkWriteOperation::UpdateMany { filter, update, upsert, array_filters } => { + let many = matches!(operation, BulkWriteOperation::UpdateMany { .. }); + let filter = json_filter_to_document(filter).map_err(|e| format!("Invalid filter: {e}"))?; + let update = json_update_to_modifications(update).map_err(|e| format!("Invalid update: {e}"))?; + let array_filters = array_filters + .as_ref() + .and_then(serde_json::Value::as_array) + .map(|filters| filters.iter().map(json_filter_to_document).collect::, _>>()) + .transpose() + .map_err(|e| format!("Invalid arrayFilters: {e}"))?; + let update_result = if many { + let mut action = col.update_many(filter, update); + if let Some(upsert) = upsert { + action = action.upsert(*upsert); + } + if let Some(filters) = array_filters { + action = action.array_filters(filters); + } + action.await.map_err(|e| e.to_string())? + } else { + let mut action = col.update_one(filter, update); + if let Some(upsert) = upsert { + action = action.upsert(*upsert); + } + if let Some(filters) = array_filters { + action = action.array_filters(filters); + } + action.await.map_err(|e| e.to_string())? + }; + result.matched_count += update_result.matched_count; + result.modified_count += update_result.modified_count; + result.upserted_count += u64::from(update_result.upserted_id.is_some()); + } + BulkWriteOperation::ReplaceOne { filter, replacement, upsert } => { + let filter = json_filter_to_document(filter).map_err(|e| format!("Invalid filter: {e}"))?; + let replacement = + json_object_to_document(replacement).map_err(|e| format!("Invalid replacement: {e}"))?; + let mut action = col.replace_one(filter, replacement); + if let Some(upsert) = upsert { + action = action.upsert(*upsert); + } + let replace_result = action.await.map_err(|e| e.to_string())?; + result.matched_count += replace_result.matched_count; + result.modified_count += replace_result.modified_count; + result.upserted_count += u64::from(replace_result.upserted_id.is_some()); + } + BulkWriteOperation::DeleteOne { filter } | BulkWriteOperation::DeleteMany { filter } => { + let many = matches!(operation, BulkWriteOperation::DeleteMany { .. }); + let filter = json_filter_to_document(filter).map_err(|e| format!("Invalid filter: {e}"))?; + let delete_result = if many { + col.delete_many(filter).await.map_err(|e| e.to_string())? + } else { + col.delete_one(filter).await.map_err(|e| e.to_string())? + }; + result.deleted_count += delete_result.deleted_count; + } + } + Ok(()) + } + .await; + + if let Err(error) = outcome { + let message = format!("operation {} ({}): {error}", index + 1, operation.kind()); + if ordered { + return Err(format!( + "bulkWrite stopped at {message}. Completed before it: {}", + describe_bulk_counts(&result) + )); + } + failures.push(message); + } + } + + if failures.is_empty() { + Ok(result) + } else { + Err(format!( + "bulkWrite finished with {} failed operation(s): {}. Completed: {}", + failures.len(), + failures.join("; "), + describe_bulk_counts(&result) + )) + } +} + +fn describe_bulk_counts(result: &MongoBulkWriteResult) -> String { + format!( + "inserted {}, matched {}, modified {}, deleted {}, upserted {}", + result.inserted_count, result.matched_count, result.modified_count, result.deleted_count, result.upserted_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( diff --git a/crates/dbx-core/src/mongo_ops.rs b/crates/dbx-core/src/mongo_ops.rs index 473c6529f1..bef7c44aaf 100644 --- a/crates/dbx-core/src/mongo_ops.rs +++ b/crates/dbx-core/src/mongo_ops.rs @@ -858,6 +858,58 @@ pub async fn mongo_replace_document_core( } } +pub async fn mongo_bulk_write_core( + state: &AppState, + connection_id: &str, + database: &str, + collection: &str, + operations_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::bulk_write(client, database, collection, operations_json, options_json).await + } + PoolKind::Agent(client) => { + let mut client = client.lock().await; + if !client.supports_capability(AgentCapability::MongoBulkWrite) { + return Err( + "MongoDB Legacy Agent does not support bulkWrite; upgrade or reinstall the MongoDB Legacy driver" + .to_string(), + ); + } + client + .mongo_bulk_write(serde_json::json!({ + "database": database, + "collection": collection, + "operations_json": operations_json, + "options_json": options_json, + })) + .await + } + _ => Err("Not a MongoDB connection".to_string()), + } +} + +/// One row of counts, the way the shell prints a `BulkWriteResult`. +pub fn mongo_bulk_write_query_result(result: &mongo_driver::MongoBulkWriteResult) -> QueryResult { + let columns = ["insertedCount", "matchedCount", "modifiedCount", "deletedCount", "upsertedCount"]; + let values = [ + result.inserted_count, + result.matched_count, + result.modified_count, + result.deleted_count, + result.upserted_count, + ]; + query_result( + columns.iter().map(ToString::to_string).collect(), + vec![values.iter().map(|value| serde_json::Value::from(*value)).collect()], + result.inserted_count + result.modified_count + result.deleted_count + result.upserted_count, + ) +} + pub async fn mongo_delete_document_core( state: &AppState, connection_id: &str, @@ -1085,6 +1137,12 @@ 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::BulkWrite { collection, operations, options } => { + let result = + mongo_bulk_write_core(state, connection_id, database, collection, operations, options.as_deref()) + .await?; + Ok(mongo_bulk_write_query_result(&result)) + } MongoCommand::Replace { collection, filter, replacement, options } => { let affected = mongo_replace_document_core( state, diff --git a/crates/dbx-core/src/mongo_shell.rs b/crates/dbx-core/src/mongo_shell.rs index 511416dd47..c77054b32e 100644 --- a/crates/dbx-core/src/mongo_shell.rs +++ b/crates/dbx-core/src/mongo_shell.rs @@ -68,6 +68,8 @@ pub enum MongoCommand { 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 = "bulkWrite")] + BulkWrite { collection: String, operations: String, options: Option }, #[serde(rename = "delete")] Delete { collection: String, filter: String, many: bool }, #[serde(rename = "createIndex")] @@ -101,6 +103,7 @@ impl MongoCommand { | Self::Insert { .. } | Self::Update { .. } | Self::Replace { .. } + | Self::BulkWrite { .. } | Self::Delete { .. } | Self::CreateIndex { .. } | Self::DropIndexes { .. } @@ -125,6 +128,9 @@ impl MongoCommand { | Self::FindOneAndUpdate { filter, .. } | Self::FindOneAndReplace { filter, .. } | Self::FindOneAndDelete { filter, .. } => is_empty_object(filter), + Self::BulkWrite { operations, .. } => { + bulk_write_filters(operations).iter().any(|filter| is_empty_object(filter)) + } _ => false, } } @@ -137,11 +143,160 @@ impl MongoCommand { | Self::FindOneAndUpdate { filter, .. } | Self::FindOneAndReplace { filter, .. } | Self::FindOneAndDelete { filter, .. } => mongo_filter_is_effectively_unbounded(filter), + Self::BulkWrite { operations, .. } => { + bulk_write_filters(operations).iter().any(|filter| mongo_filter_is_effectively_unbounded(filter)) + } _ => false, } } } +/// One entry of a `bulkWrite([...])` array, validated but not yet converted to BSON. +#[derive(Debug, Clone, PartialEq)] +pub enum BulkWriteOperation { + InsertOne { document: Value }, + UpdateOne { filter: Value, update: Value, upsert: Option, array_filters: Option }, + UpdateMany { filter: Value, update: Value, upsert: Option, array_filters: Option }, + ReplaceOne { filter: Value, replacement: Value, upsert: Option }, + DeleteOne { filter: Value }, + DeleteMany { filter: Value }, +} + +impl BulkWriteOperation { + pub fn kind(&self) -> &'static str { + match self { + Self::InsertOne { .. } => "insertOne", + Self::UpdateOne { .. } => "updateOne", + Self::UpdateMany { .. } => "updateMany", + Self::ReplaceOne { .. } => "replaceOne", + Self::DeleteOne { .. } => "deleteOne", + Self::DeleteMany { .. } => "deleteMany", + } + } + + fn filter(&self) -> Option<&Value> { + match self { + Self::InsertOne { .. } => None, + Self::UpdateOne { filter, .. } + | Self::UpdateMany { filter, .. } + | Self::ReplaceOne { filter, .. } + | Self::DeleteOne { filter } + | Self::DeleteMany { filter } => Some(filter), + } + } +} + +/// Validate a `bulkWrite` operations array the way the shell does: each entry is one +/// `{ : { ... } }` document with exactly the fields that operation takes. +pub fn parse_bulk_write_operations(operations_json: &str) -> Result, String> { + let Some(Value::Array(entries)) = parse_json_value(operations_json) else { + return Err("MongoDB bulkWrite() requires an array of operations.".to_string()); + }; + if entries.is_empty() { + return Err("MongoDB bulkWrite() requires at least one operation.".to_string()); + } + entries.iter().enumerate().map(|(index, entry)| parse_bulk_write_operation(index, entry)).collect() +} + +fn parse_bulk_write_operation(index: usize, entry: &Value) -> Result { + let position = index + 1; + let Value::Object(wrapper) = entry else { + return Err(format!("MongoDB bulkWrite() operation {position} must be a document such as {{ insertOne: {{ document: {{ ... }} }} }}.")); + }; + let (kind, spec) = match wrapper.iter().collect::>().as_slice() { + [(kind, Value::Object(spec))] => (kind.as_str(), spec), + [(kind, _)] => return Err(format!("MongoDB bulkWrite() operation {position} ({kind}) must be a document.")), + _ => return Err(format!("MongoDB bulkWrite() operation {position} must have exactly one operation key.")), + }; + + let allowed: &[&str] = match kind { + "insertOne" => &["document"], + "updateOne" | "updateMany" => &["filter", "update", "upsert", "arrayFilters"], + "replaceOne" => &["filter", "replacement", "upsert"], + "deleteOne" | "deleteMany" => &["filter"], + other => { + return Err(format!( + "MongoDB bulkWrite() operation {position} uses unsupported operation {other}; supported: insertOne, updateOne, updateMany, replaceOne, deleteOne, deleteMany." + )) + } + }; + if let Some(unknown) = spec.keys().find(|key| !allowed.contains(&key.as_str())) { + return Err(format!("MongoDB bulkWrite() operation {position} ({kind}) has unsupported field {unknown}.")); + } + let document = |field: &str| -> Result { + match spec.get(field) { + Some(value @ Value::Object(_)) => Ok(value.clone()), + Some(_) => { + Err(format!("MongoDB bulkWrite() operation {position} ({kind}) field {field} must be a document.")) + } + None => Err(format!("MongoDB bulkWrite() operation {position} ({kind}) requires a {field} document.")), + } + }; + let upsert = match spec.get("upsert") { + None => None, + Some(Value::Bool(value)) => Some(*value), + Some(_) => return Err(format!("MongoDB bulkWrite() operation {position} ({kind}) upsert must be a boolean.")), + }; + + match kind { + "insertOne" => Ok(BulkWriteOperation::InsertOne { document: document("document")? }), + "updateOne" | "updateMany" => { + let update = match spec.get("update") { + Some(value @ Value::Object(object)) => { + if object.is_empty() || !object.keys().all(|key| key.starts_with('$')) { + return Err(format!( + "MongoDB bulkWrite() operation {position} ({kind}) update must use operators such as $set; use replaceOne for a whole document." + )); + } + value.clone() + } + Some(value @ Value::Array(_)) => value.clone(), + Some(_) => { + return Err(format!( + "MongoDB bulkWrite() operation {position} ({kind}) update must be a document or pipeline." + )) + } + None => return Err(format!("MongoDB bulkWrite() operation {position} ({kind}) requires an update.")), + }; + let array_filters = match spec.get("arrayFilters") { + None => None, + Some(value @ Value::Array(_)) => Some(value.clone()), + Some(_) => { + return Err(format!( + "MongoDB bulkWrite() operation {position} ({kind}) arrayFilters must be an array." + )) + } + }; + let filter = document("filter")?; + Ok(if kind == "updateOne" { + BulkWriteOperation::UpdateOne { filter, update, upsert, array_filters } + } else { + BulkWriteOperation::UpdateMany { filter, update, upsert, array_filters } + }) + } + "replaceOne" => { + let replacement = document("replacement")?; + if let Some(operator) = + replacement.as_object().and_then(|object| object.keys().find(|key| key.starts_with('$'))) + { + return Err(format!( + "MongoDB bulkWrite() operation {position} (replaceOne) replacement must not contain update operators such as {operator}; use updateOne to modify fields." + )); + } + Ok(BulkWriteOperation::ReplaceOne { filter: document("filter")?, replacement, upsert }) + } + "deleteOne" => Ok(BulkWriteOperation::DeleteOne { filter: document("filter")? }), + _ => Ok(BulkWriteOperation::DeleteMany { filter: document("filter")? }), + } +} + +/// Filters of every non-insert operation, as JSON text, for the safety checks. +fn bulk_write_filters(operations_json: &str) -> Vec { + parse_bulk_write_operations(operations_json) + .map(|operations| operations.iter().filter_map(|op| op.filter().map(Value::to_string)).collect()) + .unwrap_or_default() +} + pub fn validate_safety( command: &MongoCommand, allow_writes: bool, @@ -614,6 +769,19 @@ pub fn parse(input: &str) -> Result { return Ok(MongoCommand::Insert { collection, documents }); } + if let Some((args, tail)) = method_call(source, prefix_end, "bulkWrite") { + if !tail.is_empty() || !(1..=2).contains(&args.len()) { + return Err("MongoDB bulkWrite() requires an array of operations and optional options.".to_string()); + } + let operations = normalized_json(&args[0])?; + parse_bulk_write_operations(&operations)?; + let options = args.get(1).filter(|arg| !arg.trim().is_empty()).map(|arg| normalized_json(arg)).transpose()?; + if let Some(options) = &options { + validate_bulk_write_options(options)?; + } + return Ok(MongoCommand::BulkWrite { collection, operations, options }); + } + if let Some((args, tail)) = method_call(source, prefix_end, "replaceOne") { if !tail.is_empty() || !(2..=3).contains(&args.len()) { return Err( @@ -1369,6 +1537,21 @@ fn parse_use_database(source: &str) -> Option { Some(database.to_string()) } +/// `bulkWrite` options: only `ordered` is honoured, so anything else is rejected rather than dropped. +fn validate_bulk_write_options(options: &str) -> Result<(), String> { + let Some(Value::Object(object)) = parse_json_value(options) else { + return Err("MongoDB bulkWrite() options must be a document.".to_string()); + }; + for (key, value) in &object { + match (key.as_str(), value) { + ("ordered", Value::Bool(_)) => {} + ("ordered", _) => return Err("MongoDB bulkWrite() ordered option must be a boolean.".to_string()), + (other, _) => return Err(format!("Unsupported MongoDB bulkWrite() option: {other}.")), + } + } + Ok(()) +} + /// 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 { @@ -1944,6 +2127,98 @@ mod tests { } } + #[test] + fn parses_bulk_write_with_every_operation_kind() { + let command = parse( + r#"db.orders.bulkWrite([ + { insertOne: { document: { sku: "A1", stock: 1 } } }, + { updateOne: { filter: { sku: "A1" }, update: { $inc: { stock: 1 } }, upsert: true } }, + { updateMany: { filter: { archived: true }, update: { $set: { stock: 0 } }, arrayFilters: [] } }, + { replaceOne: { filter: { sku: "B2" }, replacement: { sku: "B2", stock: 9 } } }, + { deleteOne: { filter: { sku: "C3" } } }, + { deleteMany: { filter: { stock: { $lt: 0 } } } } + ], { ordered: false })"#, + ) + .unwrap(); + let MongoCommand::BulkWrite { collection, operations, options } = &command else { + panic!("expected a bulkWrite command"); + }; + assert_eq!(collection, "orders"); + assert_eq!(options.as_deref(), Some(r#"{"ordered":false}"#)); + let kinds: Vec<_> = parse_bulk_write_operations(operations).unwrap().iter().map(|op| op.kind()).collect(); + assert_eq!(kinds, ["insertOne", "updateOne", "updateMany", "replaceOne", "deleteOne", "deleteMany"]); + + assert!(command.is_mutating()); + assert!(!command.is_dangerous()); + assert!(!command.has_effectively_unbounded_filter()); + assert_eq!(validate_safety(&command, true, false, false), Ok(())); + assert_eq!(validate_safety(&command, false, false, false), Err(MongoSafetyError::WritesDisabled)); + + // Only an insert: no filters at all, so never unbounded. + let inserts = parse(r#"db.orders.bulkWrite([{ insertOne: { document: { a: 1 } } }])"#).unwrap(); + assert!(!inserts.has_empty_filter()); + } + + #[test] + fn bulk_write_with_any_unbounded_filter_is_guarded_like_an_update() { + let command = parse( + r#"db.orders.bulkWrite([ + { updateOne: { filter: { sku: "A1" }, update: { $set: { a: 1 } } } }, + { deleteMany: { filter: {} } } + ])"#, + ) + .unwrap(); + assert!(command.has_empty_filter()); + assert!(command.has_effectively_unbounded_filter()); + assert_eq!(validate_safety(&command, true, false, false), Err(MongoSafetyError::EmptyFilter)); + } + + #[test] + fn rejects_malformed_bulk_write_operations() { + for (source, expected) in [ + ("db.orders.bulkWrite()", "array of operations"), + ("db.orders.bulkWrite({})", "array of operations"), + ("db.orders.bulkWrite([])", "at least one"), + ("db.orders.bulkWrite([1])", "operation 1 must be a document"), + ("db.orders.bulkWrite([{ insertOne: {}, deleteOne: {} }])", "exactly one operation key"), + ("db.orders.bulkWrite([{ upsertOne: { document: {} } }])", "unsupported operation upsertOne"), + ("db.orders.bulkWrite([{ insertOne: { doc: {} } }])", "unsupported field doc"), + ("db.orders.bulkWrite([{ insertOne: {} }])", "requires a document document"), + ( + "db.orders.bulkWrite([{ updateOne: { filter: {}, update: { a: 1 } } }])", + "must use operators such as $set", + ), + ( + "db.orders.bulkWrite([{ updateOne: { filter: {}, update: { $set: { a: 1 } }, upsert: 1 } }])", + "upsert must be a boolean", + ), + ( + "db.orders.bulkWrite([{ updateOne: { filter: {}, update: { $set: { a: 1 } }, collation: {} } }])", + "unsupported field collation", + ), + ( + "db.orders.bulkWrite([{ replaceOne: { filter: {}, replacement: { $set: { a: 1 } } } }])", + "must not contain update operators such as $set", + ), + ("db.orders.bulkWrite([{ deleteOne: { filter: [] } }])", "field filter must be a document"), + ( + "db.orders.bulkWrite([{ deleteOne: { filter: {} } }], { ordered: 1 })", + "ordered option must be a boolean", + ), + ( + "db.orders.bulkWrite([{ deleteOne: { filter: {} } }], { writeConcern: {} })", + "Unsupported MongoDB bulkWrite() option: writeConcern", + ), + ( + "db.orders.bulkWrite([{ deleteOne: { filter: {} } }]).limit(1)", + "array of operations and optional options", + ), + ] { + let error = parse(source).unwrap_err(); + assert!(error.contains(expected), "{source}\n expected: {expected}\n got: {error}"); + } + } + #[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 ce3d528894..9f4ec0d1bc 100644 --- a/crates/dbx-mcp/src/backend.rs +++ b/crates/dbx-mcp/src/backend.rs @@ -1860,6 +1860,25 @@ impl DbxBackend for WebBackend { .map_err(|error| format!("Invalid MongoDB insert response: {error}"))?; Ok(affected_query_result(affected_rows_from_value(&value))) } + MongoCommand::BulkWrite { collection, operations, options } => { + let result: dbx_core::db::mongo_driver::MongoBulkWriteResult = self + .request( + reqwest::Method::POST, + "/api/mongo/bulk-write", + Some(json!({ + "connectionId": connection_id, + "database": database, + "collection": collection, + "operationsJson": operations, + "optionsJson": options, + })), + ) + .await? + .json() + .await + .map_err(|error| format!("Invalid MongoDB bulkWrite response: {error}"))?; + Ok(dbx_core::mongo_ops::mongo_bulk_write_query_result(&result)) + } MongoCommand::Replace { collection, filter, replacement, options } => { let value: Value = self .request( diff --git a/crates/dbx-web/src/main.rs b/crates/dbx-web/src/main.rs index 41e0a1daee..574f89d40b 100644 --- a/crates/dbx-web/src/main.rs +++ b/crates/dbx-web/src/main.rs @@ -968,6 +968,7 @@ async fn main() { .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/bulk-write", post(routes::mongo::bulk_write)) .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 90e33ce4fc..151ed71733 100644 --- a/crates/dbx-web/src/routes/mongo.rs +++ b/crates/dbx-web/src/routes/mongo.rs @@ -269,6 +269,16 @@ pub struct MongoReplaceDocumentRequest { pub options_json: Option, } +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct MongoBulkWriteRequest { + pub connection_id: String, + pub database: String, + pub collection: String, + pub operations_json: String, + pub options_json: Option, +} + #[derive(Deserialize)] #[serde(rename_all = "camelCase")] pub struct MongoDeleteDocumentsRequest { @@ -846,6 +856,36 @@ pub async fn replace_document( Ok(Json(serde_json::json!({ "affected_rows": result }))) } +pub async fn bulk_write( + 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?; + let command = dbx_core::mongo_shell::MongoCommand::BulkWrite { + collection: req.collection.clone(), + operations: req.operations_json.clone(), + options: None, + }; + if command.has_effectively_unbounded_filter() { + super::mcp_policy::ensure_dangerous_write(&state, &headers, &req.connection_id, &database, "BulkWrite").await?; + } else { + super::mcp_policy::ensure_write(&state, &headers, &req.connection_id, &database, "BulkWrite").await?; + } + ensure_writable(&state.app, &req.connection_id, "BulkWrite").await?; + let result = dbx_core::mongo_ops::mongo_bulk_write_core( + &state.app, + &req.connection_id, + &database, + &req.collection, + &req.operations_json, + req.options_json.as_deref(), + ) + .await + .map_err(AppError::from)?; + Ok(Json(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 574ceee81a..0e54d7bb64 100644 --- a/packages/app-tests/mongoCompletion.test.ts +++ b/packages/app-tests/mongoCompletion.test.ts @@ -501,8 +501,9 @@ test("suggests only helpers the shell parser accepts", () => { assert.ok(methodLabels.includes("distinct")); assert.ok(methodLabels.includes("estimatedDocumentCount")); assert.ok(methodLabels.includes("replaceOne")); + assert.ok(methodLabels.includes("bulkWrite")); // Suggesting a helper DBX cannot run just hands the user a command that fails. - for (const unsupported of ["bulkWrite", "renameCollection", "mapReduce"]) { + for (const unsupported of ["renameCollection", "mapReduce", "watch"]) { 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 2d4bee7e98..d3fec0f49e 100644 --- a/packages/app-tests/mongoShellCommand.test.ts +++ b/packages/app-tests/mongoShellCommand.test.ts @@ -5,6 +5,7 @@ import { evaluateMongoAggregateSafety, evaluateMongoWriteSafety, mongoAggregateWriteStage, + mongoBulkWriteToQueryResult, mongoCollectionStatsToQueryResult, mongoCountToQueryResult, mongoDatabasesToQueryResult, @@ -672,6 +673,65 @@ test("parseMongoWriteCommand accepts unquoted insert and update commands", () => }); }); +test("parseMongoWriteCommand reads bulkWrite with every operation kind", () => { + const command = parseMongoWriteCommand(`db.orders.bulkWrite([ + { insertOne: { document: { sku: "A1", stock: 1 } } }, + { updateOne: { filter: { sku: "A1" }, update: { $inc: { stock: 1 } }, upsert: true } }, + { updateMany: { filter: { archived: true }, update: [{ $set: { stock: 0 } }] } }, + { replaceOne: { filter: { sku: "B2" }, replacement: { sku: "B2", stock: 9 } } }, + { deleteOne: { filter: { sku: "C3" } } }, + { deleteMany: { filter: { stock: { $lt: 0 } } } } + ], { ordered: false })`); + assert.ok(command && command.kind === "bulkWrite"); + assert.equal(command.collection, "orders"); + assert.equal(command.options, '{ "ordered": false }'); + assert.deepEqual( + (JSON.parse(command.operations) as Array>).map((entry) => Object.keys(entry)[0]), + ["insertOne", "updateOne", "updateMany", "replaceOne", "deleteOne", "deleteMany"], + ); + assert.equal(parseMongoWriteCommand('db["my-coll"].bulkWrite([{ deleteOne: { filter: { a: 1 } } }]);')?.kind, "bulkWrite"); +}); + +test("bulkWrite operations are validated and the problem is named", () => { + for (const [source, expected] of [ + ["db.orders.bulkWrite()", /array of operations and optional options/], + ["db.orders.bulkWrite({})", /requires an array of operations/], + ["db.orders.bulkWrite([])", /at least one operation/], + ["db.orders.bulkWrite([1])", /operation 1 must be a document/], + ["db.orders.bulkWrite([{ insertOne: {}, deleteOne: {} }])", /exactly one operation key/], + ["db.orders.bulkWrite([{ upsertOne: { document: {} } }])", /unsupported operation upsertOne/], + ["db.orders.bulkWrite([{ insertOne: { doc: {} } }])", /unsupported field doc/], + ["db.orders.bulkWrite([{ insertOne: {} }])", /requires a document document/], + ["db.orders.bulkWrite([{ updateOne: { filter: {}, update: { a: 1 } } }])", /must use operators such as \$set/], + ["db.orders.bulkWrite([{ updateOne: { filter: {}, update: { $set: { a: 1 } }, upsert: 1 } }])", /upsert must be a boolean/], + ["db.orders.bulkWrite([{ updateOne: { filter: {}, update: { $set: { a: 1 } }, collation: {} } }])", /unsupported field collation/], + ["db.orders.bulkWrite([{ replaceOne: { filter: {}, replacement: { $set: { a: 1 } } } }])", /must not contain update operators such as \$set/], + ["db.orders.bulkWrite([{ deleteOne: { filter: [] } }])", /field filter must be a document/], + ["db.orders.bulkWrite([{ deleteOne: { filter: {} } }], { ordered: 1 })", /ordered option must be a boolean/], + ["db.orders.bulkWrite([{ deleteOne: { filter: {} } }], { writeConcern: {} })", /Unsupported bulkWrite\(\) option: writeConcern/], + ] as const) { + assert.equal(parseMongoWriteCommand(source), null, source); + assert.match(describeMongoCommandParseFailure(source), expected, source); + } +}); + +test("evaluateMongoWriteSafety treats a bulkWrite as risky as its least-bounded operation", () => { + const policy = { allowWrites: true, allowDangerous: false } as Parameters[1]; + const bounded = parseMongoWriteCommand("db.orders.bulkWrite([{ insertOne: { document: { a: 1 } } }, { deleteOne: { filter: { sku: 'X' } } }])")!; + const unbounded = parseMongoWriteCommand("db.orders.bulkWrite([{ updateOne: { filter: { sku: 'X' }, update: { $set: { a: 1 } } } }, { deleteMany: { filter: {} } }])")!; + assert.equal(evaluateMongoWriteSafety(bounded, policy).allowed, true); + assert.equal(evaluateMongoWriteSafety(unbounded, policy).allowed, false); + assert.equal(evaluateMongoWriteSafety(unbounded, { ...policy, allowDangerous: true }).allowed, true); +}); + +test("mongoBulkWriteToQueryResult renders one row of counts", () => { + const result = mongoBulkWriteToQueryResult({ inserted_count: 2, matched_count: 3, modified_count: 1, deleted_count: 4, upserted_count: 1 }, 12.6); + assert.deepEqual(result.columns, ["insertedCount", "matchedCount", "modifiedCount", "deletedCount", "upsertedCount"]); + assert.deepEqual(result.rows, [[2, 3, 1, 4, 1]]); + assert.equal(result.affected_rows, 8); + assert.equal(result.execution_time_ms, 13); +}); + test("parseMongoWriteCommand reads replaceOne as a filtered replace", () => { assert.deepEqual(parseMongoWriteCommand('db.orders.replaceOne({_id: ObjectId("507f1f77bcf86cd799439011")}, {name: "new", tags: []}, {upsert: true})'), { kind: "replace", @@ -1249,8 +1309,8 @@ test("describeMongoCommandParseFailure explains the argument shape a known metho }); test("describeMongoCommandParseFailure names unsupported methods and points at alternatives", () => { - const bulkWrite = describeMongoCommandParseFailure("db.c.bulkWrite([{insertOne: {document: {a: 1}}}])"); - assert.match(bulkWrite, /^Collection method bulkWrite\(\) is not supported\. Supported collection methods: find, findOne/); + const rename = describeMongoCommandParseFailure('db.c.renameCollection("d")'); + assert.match(rename, /^Collection method renameCollection\(\) 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/); @@ -1261,7 +1321,7 @@ test("describeMongoCommandParseFailure names unsupported methods and points at a assert.match(describeMongoCommandParseFailure("show collections"), /listed in the sidebar/); // Leading comments do not hide the command shape. - assert.match(describeMongoCommandParseFailure("// note\ndb.c.bulkWrite([])"), /bulkWrite\(\) is not supported/); + assert.match(describeMongoCommandParseFailure("// note\ndb.c.mapReduce()"), /mapReduce\(\) is not supported/); }); test("describeMongoCommandParseFailure reports unclosed delimiters and shell hints", () => { diff --git a/src-tauri/src/commands/mcp_bridge.rs b/src-tauri/src/commands/mcp_bridge.rs index 537573b608..b7de3655c2 100644 --- a/src-tauri/src/commands/mcp_bridge.rs +++ b/src-tauri/src/commands/mcp_bridge.rs @@ -931,6 +931,25 @@ pub(crate) async fn ensure_mcp_mongo_filtered_write_allowed_by_id( } } +/// A bulk write is as dangerous as its least-bounded operation. +pub(crate) async fn ensure_mcp_mongo_bulk_write_allowed_by_id( + state: &Arc, + connection_id: &str, + database: &str, + operations_json: &str, +) -> Result<(), String> { + let command = dbx_core::mongo_shell::MongoCommand::BulkWrite { + collection: String::new(), + operations: operations_json.to_string(), + options: None, + }; + if command.has_effectively_unbounded_filter() { + ensure_mcp_dangerous_write_allowed_by_id(state, connection_id, database, "BulkWrite").await + } else { + ensure_mcp_write_allowed_by_id(state, connection_id, database, "BulkWrite").await + } +} + pub(crate) async fn ensure_mcp_mongo_aggregate_allowed_by_id( state: &Arc, connection_id: &str, diff --git a/src-tauri/src/commands/mongo_cmd.rs b/src-tauri/src/commands/mongo_cmd.rs index 4a3e6f55d4..ec9621910b 100644 --- a/src-tauri/src/commands/mongo_cmd.rs +++ b/src-tauri/src/commands/mongo_cmd.rs @@ -575,6 +575,37 @@ pub async fn mongo_replace_document( .await } +#[tauri::command] +pub async fn mongo_bulk_write( + state: State<'_, Arc>, + connection_id: String, + database: String, + collection: String, + operations_json: String, + options_json: Option, + mcp_request: Option, +) -> Result { + if mcp_request == Some(true) { + crate::commands::mcp_bridge::ensure_mcp_mongo_bulk_write_allowed_by_id( + state.inner(), + &connection_id, + &database, + &operations_json, + ) + .await?; + } + ensure_connection_writable(&state, &connection_id, "BulkWrite").await?; + dbx_core::mongo_ops::mongo_bulk_write_core( + &state, + &connection_id, + &database, + &collection, + &operations_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 1c44b2d8d6..5566f697db 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -2342,6 +2342,7 @@ pub fn run() { commands::mongo_cmd::mongo_update_document, commands::mongo_cmd::mongo_update_documents, commands::mongo_cmd::mongo_replace_document, + commands::mongo_cmd::mongo_bulk_write, commands::document_cmd::document_delete_document, commands::document_cmd::document_save_meilisearch_batch, commands::document_cmd::meilisearch_search_documents,