feat(mongodb): support bulkWrite

This commit is contained in:
Loc Nguyen authored and GitHub committed 2026-09-17 12:54:21 +08:00
1 parent 3f3c21aae9
commit f78dc1ba16
25 files changed
+1089 -18

No files matched your search

@@ -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,
@@ -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"]
}
@@ -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"
@@ -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<WriteModel<Document>> 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<String, Object> 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<String, Object> bulkWriteCounts(BulkWriteResult result) {
Map<String, Object> 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 `{ <op>: { ... } }`, as driver write models. */
static List<WriteModel<Document>> 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<WriteModel<Document>> 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<Document> 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<Document> 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<Document> 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);
@@ -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<WriteModel<Document>> 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<Document>) models.get(1)).getOptions().isUpsert());
assertTrue(models.get(2) instanceof UpdateManyModel);
assertNotNull(((UpdateManyModel<Document>) 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());
+1
View File
@@ -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");
+11
View File
@@ -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<MongoBulkWriteResult> {
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,
+11
View File
@@ -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<MongoBulkWriteResult> {
return invoke<MongoBulkWriteResult>("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<number>("mongo_replace_document", {
connectionId,
@@ -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<string, readonly MongoArgRole[]> = {
find: ["filter", "projection", "options"],
@@ -193,6 +195,7 @@ const METHOD_ARG_ROLES: Record<string, readonly MongoArgRole[]> = {
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 `{ <op>: { filter, update, … } }` entries; completing inside them is a follow-up.
case "operations":
case "options":
return { mode: "none" };
default:
+126 -2
View File
@@ -49,6 +49,7 @@ const COLLECTION_METHOD_SHAPES: Record<string, MongoMethodShape> = {
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<string, unknown>).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<string, unknown> {
return isRecord(value) && Object.keys(value).length > 0;
}
const BULK_WRITE_FIELDS: Record<string, readonly string[]> = {
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<string, unknown> => 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 `{ <op>: { ... } }` 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<string, unknown>).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<string, unknown>;
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<Record<string, { filter?: unknown }>>;
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":
+5
View File
@@ -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))));
@@ -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"]
}
@@ -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"
+16 -3
View File
@@ -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<i64>, 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<T: DeserializeOwned + Send + 'static>(&mut self, params: Value) -> Result<T, String> {
self.call_mongo_method(MongoAgentMethod::BulkWrite, params).await
}
pub async fn mongo_delete_document<T: DeserializeOwned + Send + 'static>(
&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");
+147
View File
@@ -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<bool>,
}
/// 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<MongoBulkWriteResult, String> {
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::<Document>(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::<Result<Vec<_>, _>>())
.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(
+58
View File
@@ -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<mongo_driver::MongoBulkWriteResult, String> {
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,
+275
View File
@@ -68,6 +68,8 @@ pub enum MongoCommand {
Update { collection: String, filter: String, update: String, options: Option<String>, many: bool },
#[serde(rename = "replace")]
Replace { collection: String, filter: String, replacement: String, options: Option<String> },
#[serde(rename = "bulkWrite")]
BulkWrite { collection: String, operations: String, options: Option<String> },
#[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<bool>, array_filters: Option<Value> },
UpdateMany { filter: Value, update: Value, upsert: Option<bool>, array_filters: Option<Value> },
ReplaceOne { filter: Value, replacement: Value, upsert: Option<bool> },
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
/// `{ <op>: { ... } }` document with exactly the fields that operation takes.
pub fn parse_bulk_write_operations(operations_json: &str) -> Result<Vec<BulkWriteOperation>, 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<BulkWriteOperation, String> {
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::<Vec<_>>().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<Value, String> {
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<String> {
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<MongoCommand, String> {
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<String> {
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 [
+19
View File
@@ -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(
+1
View File
@@ -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))
+40
View File
@@ -269,6 +269,16 @@ pub struct MongoReplaceDocumentRequest {
pub options_json: Option<String>,
}
#[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<String>,
}
#[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<Arc<WebState>>,
headers: HeaderMap,
Json(req): Json<MongoBulkWriteRequest>,
) -> Result<Json<dbx_core::db::mongo_driver::MongoBulkWriteResult>, 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<Arc<WebState>>,
headers: HeaderMap,
+2 -1
View File
@@ -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.
+63 -3
View File
@@ -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<Record<string, unknown>>).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<typeof evaluateMongoWriteSafety>[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", () => {
+19
View File
@@ -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<AppState>,
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<AppState>,
connection_id: &str,
+31
View File
@@ -575,6 +575,37 @@ pub async fn mongo_replace_document(
.await
}
#[tauri::command]
pub async fn mongo_bulk_write(
state: State<'_, Arc<AppState>>,
connection_id: String,
database: String,
collection: String,
operations_json: String,
options_json: Option<String>,
mcp_request: Option<bool>,
) -> Result<dbx_core::db::mongo_driver::MongoBulkWriteResult, String> {
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<AppState>>,
+1
View File
@@ -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,