feat(mongodb): support bulkWrite
This commit is contained in:
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());
|
||||
|
||||
@@ -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");
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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":
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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");
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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 [
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -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
@@ -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
@@ -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", () => {
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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>>,
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in new issue
Block a user