feat(mongodb): support bulkWrite

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

No files matched your search

@@ -2,7 +2,7 @@
"protocolVersion": 1,
"handshakeMethod": "handshake",
"handshakeResponseFields": ["protocolVersion", "agentProtocolVersion", "capabilities"],
"allCapabilities": ["connect", "test_connection", "metadata", "query", "paged_query", "transaction", "ddl", "kv", "kv_ttl", "kv_cas", "kv_list_values", "kv_status", "kv_history", "etcd_compaction", "etcd_defrag", "etcd_watch", "etcd_lease", "etcd_auth", "mongo_drop_database", "mongo_clone_collection", "mongo_run_command", "mongo_insert_documents", "mongo_replace_document", "mongo_find_cursor"],
"allCapabilities": ["connect", "test_connection", "metadata", "query", "paged_query", "transaction", "ddl", "kv", "kv_ttl", "kv_cas", "kv_list_values", "kv_status", "kv_history", "etcd_compaction", "etcd_defrag", "etcd_watch", "etcd_lease", "etcd_auth", "mongo_drop_database", "mongo_clone_collection", "mongo_run_command", "mongo_insert_documents", "mongo_replace_document", "mongo_bulk_write", "mongo_find_cursor"],
"capabilities": ["connect", "test_connection", "metadata", "query", "paged_query", "transaction", "ddl"],
"defaultSqlCapabilities": ["connect", "test_connection", "metadata", "query", "paged_query", "transaction", "ddl"],
"commonMethods": [
@@ -42,6 +42,6 @@
"disconnect",
"shutdown"
],
"mongoLegacyMethods": ["list_databases", "list_collections", "find_documents", "find_one", "explain_find", "aggregate_documents", "find_documents_extended_json", "count_documents", "server_version", "create_index", "create_user", "drop_indexes", "drop_collection", "clone_collection", "drop_database", "insert_document", "insert_documents", "update_document", "update_documents", "replace_document", "delete_document", "delete_documents", "run_command", "start_find_cursor", "fetch_find_cursor", "close_find_cursor"],
"mongoLegacyMethods": ["list_databases", "list_collections", "find_documents", "find_one", "explain_find", "aggregate_documents", "find_documents_extended_json", "count_documents", "server_version", "create_index", "create_user", "drop_indexes", "drop_collection", "clone_collection", "drop_database", "insert_document", "insert_documents", "update_document", "update_documents", "replace_document", "bulk_write", "delete_document", "delete_documents", "run_command", "start_find_cursor", "fetch_find_cursor", "close_find_cursor"],
"kvMethods": ["kv_list_prefix", "kv_get", "kv_put", "kv_delete", "kv_rename", "kv_history", "kv_status", "etcd_compact", "etcd_defrag", "etcd_watch_start", "etcd_watch_poll", "etcd_watch_stop", "etcd_lease_list", "etcd_lease_get", "etcd_lease_grant", "etcd_lease_keepalive_once", "etcd_lease_revoke", "etcd_auth_user_list", "etcd_auth_user_get", "etcd_auth_user_add", "etcd_auth_user_delete", "etcd_auth_user_change_password", "etcd_auth_user_grant_role", "etcd_auth_user_revoke_role", "etcd_auth_role_list", "etcd_auth_role_get", "etcd_auth_role_add", "etcd_auth_role_delete", "etcd_auth_role_grant_permission", "etcd_auth_role_revoke_permission"]
}
@@ -2,7 +2,7 @@
"protocolVersion": 2,
"handshakeMethod": "handshake",
"handshakeResponseFields": ["protocolVersion", "agentProtocolVersion", "capabilities"],
"allCapabilities": ["connect", "test_connection", "metadata", "query", "paged_query", "transaction", "ddl", "kv", "kv_ttl", "kv_cas", "kv_list_values", "kv_status", "kv_history", "etcd_compaction", "etcd_defrag", "etcd_watch", "etcd_lease", "etcd_auth", "mongo_drop_database", "mongo_clone_collection", "mongo_run_command", "mongo_insert_documents", "mongo_replace_document", "mongo_find_cursor", "multi_session", "structured_error_v1"],
"allCapabilities": ["connect", "test_connection", "metadata", "query", "paged_query", "transaction", "ddl", "kv", "kv_ttl", "kv_cas", "kv_list_values", "kv_status", "kv_history", "etcd_compaction", "etcd_defrag", "etcd_watch", "etcd_lease", "etcd_auth", "mongo_drop_database", "mongo_clone_collection", "mongo_run_command", "mongo_insert_documents", "mongo_replace_document", "mongo_bulk_write", "mongo_find_cursor", "multi_session", "structured_error_v1"],
"capabilities": ["connect", "test_connection", "metadata", "query", "paged_query", "transaction", "ddl", "multi_session"],
"defaultSqlCapabilities": ["connect", "test_connection", "metadata", "query", "paged_query", "transaction", "ddl", "multi_session", "structured_error_v1"],
"commonMethods": [
@@ -46,7 +46,7 @@
"disconnect",
"shutdown"
],
"mongoLegacyMethods": ["list_databases", "list_collections", "find_documents", "find_one", "explain_find", "aggregate_documents", "find_documents_extended_json", "count_documents", "server_version", "create_index", "create_user", "drop_indexes", "drop_collection", "clone_collection", "drop_database", "insert_document", "insert_documents", "update_document", "update_documents", "replace_document", "delete_document", "delete_documents", "run_command", "start_find_cursor", "fetch_find_cursor", "close_find_cursor"],
"mongoLegacyMethods": ["list_databases", "list_collections", "find_documents", "find_one", "explain_find", "aggregate_documents", "find_documents_extended_json", "count_documents", "server_version", "create_index", "create_user", "drop_indexes", "drop_collection", "clone_collection", "drop_database", "insert_document", "insert_documents", "update_document", "update_documents", "replace_document", "bulk_write", "delete_document", "delete_documents", "run_command", "start_find_cursor", "fetch_find_cursor", "close_find_cursor"],
"kvMethods": ["kv_list_prefix", "kv_get", "kv_put", "kv_delete", "kv_rename", "kv_history", "kv_status", "etcd_compact", "etcd_defrag", "etcd_watch_start", "etcd_watch_poll", "etcd_watch_stop", "etcd_lease_list", "etcd_lease_get", "etcd_lease_grant", "etcd_lease_keepalive_once", "etcd_lease_revoke", "etcd_auth_user_list", "etcd_auth_user_get", "etcd_auth_user_add", "etcd_auth_user_delete", "etcd_auth_user_change_password", "etcd_auth_user_grant_role", "etcd_auth_user_revoke_role", "etcd_auth_role_list", "etcd_auth_role_get", "etcd_auth_role_add", "etcd_auth_role_delete", "etcd_auth_role_grant_permission", "etcd_auth_role_revoke_permission"],
"sessionField": "agentSessionId",
"cursorSessionField": "sessionId"
+16 -3
View File
@@ -1042,6 +1042,7 @@ pub enum AgentCapability {
MongoRunCommand,
MongoInsertDocuments,
MongoReplaceDocument,
MongoBulkWrite,
MongoFindCursor,
MultiSession,
StructuredErrorV1,
@@ -1125,7 +1126,7 @@ fn parse_agent_rpc_error_header(header: &str) -> (Option<i64>, String) {
}
impl AgentCapability {
pub const ALL: [Self; 26] = [
pub const ALL: [Self; 27] = [
Self::Connect,
Self::TestConnection,
Self::Metadata,
@@ -1149,6 +1150,7 @@ impl AgentCapability {
Self::MongoRunCommand,
Self::MongoInsertDocuments,
Self::MongoReplaceDocument,
Self::MongoBulkWrite,
Self::MongoFindCursor,
Self::MultiSession,
Self::StructuredErrorV1,
@@ -1179,6 +1181,7 @@ impl AgentCapability {
Self::MongoRunCommand => "mongo_run_command",
Self::MongoInsertDocuments => "mongo_insert_documents",
Self::MongoReplaceDocument => "mongo_replace_document",
Self::MongoBulkWrite => "mongo_bulk_write",
Self::MongoFindCursor => "mongo_find_cursor",
Self::MultiSession => "multi_session",
Self::StructuredErrorV1 => "structured_error_v1",
@@ -1383,6 +1386,7 @@ pub enum MongoAgentMethod {
UpdateDocument,
UpdateDocuments,
ReplaceDocument,
BulkWrite,
DeleteDocument,
DeleteDocuments,
RunCommand,
@@ -1392,7 +1396,7 @@ pub enum MongoAgentMethod {
}
impl MongoAgentMethod {
pub const ALL: [Self; 26] = [
pub const ALL: [Self; 27] = [
Self::ListDatabases,
Self::ListCollections,
Self::FindDocuments,
@@ -1413,6 +1417,7 @@ impl MongoAgentMethod {
Self::UpdateDocument,
Self::UpdateDocuments,
Self::ReplaceDocument,
Self::BulkWrite,
Self::DeleteDocument,
Self::DeleteDocuments,
Self::RunCommand,
@@ -1443,6 +1448,7 @@ impl MongoAgentMethod {
Self::UpdateDocument => "update_document",
Self::UpdateDocuments => "update_documents",
Self::ReplaceDocument => "replace_document",
Self::BulkWrite => "bulk_write",
Self::DeleteDocument => "delete_document",
Self::DeleteDocuments => "delete_documents",
Self::RunCommand => "run_command",
@@ -2849,6 +2855,10 @@ impl AgentDriverClient {
self.call_mongo_method(MongoAgentMethod::ReplaceDocument, params).await
}
pub async fn mongo_bulk_write<T: DeserializeOwned + Send + 'static>(&mut self, params: Value) -> Result<T, String> {
self.call_mongo_method(MongoAgentMethod::BulkWrite, params).await
}
pub async fn mongo_delete_document<T: DeserializeOwned + Send + 'static>(
&mut self,
params: Value,
@@ -3005,6 +3015,7 @@ pub fn agent_supports_capability(handshake: Option<&AgentHandshake>, capability:
| AgentCapability::MongoRunCommand
| AgentCapability::MongoInsertDocuments
| AgentCapability::MongoReplaceDocument
| AgentCapability::MongoBulkWrite
| AgentCapability::MongoFindCursor
) {
return handshake.map(|value| value.supports(capability)).unwrap_or(false);
@@ -5040,10 +5051,11 @@ for line in sys.stdin:
assert_eq!(AgentCapability::MongoRunCommand.as_str(), "mongo_run_command");
assert_eq!(AgentCapability::MongoInsertDocuments.as_str(), "mongo_insert_documents");
assert_eq!(AgentCapability::MongoReplaceDocument.as_str(), "mongo_replace_document");
assert_eq!(AgentCapability::MongoBulkWrite.as_str(), "mongo_bulk_write");
assert_eq!(AgentCapability::MongoFindCursor.as_str(), "mongo_find_cursor");
assert_eq!(AgentCapability::MultiSession.as_str(), "multi_session");
assert_eq!(AgentCapability::StructuredErrorV1.as_str(), "structured_error_v1");
assert_eq!(AgentCapability::ALL.len(), 26);
assert_eq!(AgentCapability::ALL.len(), 27);
}
#[test]
@@ -5106,6 +5118,7 @@ for line in sys.stdin:
assert_eq!(MongoAgentMethod::UpdateDocument.as_str(), "update_document");
assert_eq!(MongoAgentMethod::UpdateDocuments.as_str(), "update_documents");
assert_eq!(MongoAgentMethod::ReplaceDocument.as_str(), "replace_document");
assert_eq!(MongoAgentMethod::BulkWrite.as_str(), "bulk_write");
assert_eq!(MongoAgentMethod::DeleteDocument.as_str(), "delete_document");
assert_eq!(MongoAgentMethod::DeleteDocuments.as_str(), "delete_documents");
assert_eq!(MongoAgentMethod::StartFindCursor.as_str(), "start_find_cursor");
+147
View File
@@ -2120,6 +2120,153 @@ pub async fn update_documents(
Ok(result.modified_count)
}
/// Counts reported by `bulkWrite`, mirroring the shell's `BulkWriteResult`.
#[derive(Debug, Default, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct MongoBulkWriteResult {
pub inserted_count: u64,
pub matched_count: u64,
pub modified_count: u64,
pub deleted_count: u64,
pub upserted_count: u64,
}
#[derive(Default, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
struct MongoBulkWriteOptions {
ordered: Option<bool>,
}
/// bulkWrite(operations[, { ordered }]).
///
/// The Rust driver's `bulk_write` is client-level and needs MongoDB 8.0, so the
/// operations run here one at a time in order, which is what the shell's
/// collection-level bulkWrite guarantees anyway (it is not transactional). With
/// `ordered: true` (the default) the first failure stops the batch; with
/// `ordered: false` every operation is attempted and the failures are reported together.
pub async fn bulk_write(
client: &Client,
database: &str,
collection: &str,
operations_json: &str,
options_json: Option<&str>,
) -> Result<MongoBulkWriteResult, String> {
use crate::mongo_shell::BulkWriteOperation;
let operations = crate::mongo_shell::parse_bulk_write_operations(operations_json)?;
let ordered = match options_json.filter(|value| !value.trim().is_empty()) {
Some(raw) => {
let options: MongoBulkWriteOptions =
serde_json::from_str(raw).map_err(|e| format!("Invalid bulkWrite options: {e}"))?;
options.ordered.unwrap_or(true)
}
None => true,
};
let col = client.database(database).collection::<Document>(collection);
let mut result = MongoBulkWriteResult::default();
let mut failures = Vec::new();
for (index, operation) in operations.iter().enumerate() {
let outcome: Result<(), String> = async {
match operation {
BulkWriteOperation::InsertOne { document } => {
let document = json_object_to_document(document).map_err(|e| format!("Invalid document: {e}"))?;
col.insert_one(document).await.map_err(|e| e.to_string())?;
result.inserted_count += 1;
}
BulkWriteOperation::UpdateOne { filter, update, upsert, array_filters }
| BulkWriteOperation::UpdateMany { filter, update, upsert, array_filters } => {
let many = matches!(operation, BulkWriteOperation::UpdateMany { .. });
let filter = json_filter_to_document(filter).map_err(|e| format!("Invalid filter: {e}"))?;
let update = json_update_to_modifications(update).map_err(|e| format!("Invalid update: {e}"))?;
let array_filters = array_filters
.as_ref()
.and_then(serde_json::Value::as_array)
.map(|filters| filters.iter().map(json_filter_to_document).collect::<Result<Vec<_>, _>>())
.transpose()
.map_err(|e| format!("Invalid arrayFilters: {e}"))?;
let update_result = if many {
let mut action = col.update_many(filter, update);
if let Some(upsert) = upsert {
action = action.upsert(*upsert);
}
if let Some(filters) = array_filters {
action = action.array_filters(filters);
}
action.await.map_err(|e| e.to_string())?
} else {
let mut action = col.update_one(filter, update);
if let Some(upsert) = upsert {
action = action.upsert(*upsert);
}
if let Some(filters) = array_filters {
action = action.array_filters(filters);
}
action.await.map_err(|e| e.to_string())?
};
result.matched_count += update_result.matched_count;
result.modified_count += update_result.modified_count;
result.upserted_count += u64::from(update_result.upserted_id.is_some());
}
BulkWriteOperation::ReplaceOne { filter, replacement, upsert } => {
let filter = json_filter_to_document(filter).map_err(|e| format!("Invalid filter: {e}"))?;
let replacement =
json_object_to_document(replacement).map_err(|e| format!("Invalid replacement: {e}"))?;
let mut action = col.replace_one(filter, replacement);
if let Some(upsert) = upsert {
action = action.upsert(*upsert);
}
let replace_result = action.await.map_err(|e| e.to_string())?;
result.matched_count += replace_result.matched_count;
result.modified_count += replace_result.modified_count;
result.upserted_count += u64::from(replace_result.upserted_id.is_some());
}
BulkWriteOperation::DeleteOne { filter } | BulkWriteOperation::DeleteMany { filter } => {
let many = matches!(operation, BulkWriteOperation::DeleteMany { .. });
let filter = json_filter_to_document(filter).map_err(|e| format!("Invalid filter: {e}"))?;
let delete_result = if many {
col.delete_many(filter).await.map_err(|e| e.to_string())?
} else {
col.delete_one(filter).await.map_err(|e| e.to_string())?
};
result.deleted_count += delete_result.deleted_count;
}
}
Ok(())
}
.await;
if let Err(error) = outcome {
let message = format!("operation {} ({}): {error}", index + 1, operation.kind());
if ordered {
return Err(format!(
"bulkWrite stopped at {message}. Completed before it: {}",
describe_bulk_counts(&result)
));
}
failures.push(message);
}
}
if failures.is_empty() {
Ok(result)
} else {
Err(format!(
"bulkWrite finished with {} failed operation(s): {}. Completed: {}",
failures.len(),
failures.join("; "),
describe_bulk_counts(&result)
))
}
}
fn describe_bulk_counts(result: &MongoBulkWriteResult) -> String {
format!(
"inserted {}, matched {}, modified {}, deleted {}, upserted {}",
result.inserted_count, result.matched_count, result.modified_count, result.deleted_count, result.upserted_count
)
}
/// replaceOne(filter, replacement[, { upsert }]): the whole document is swapped, so
/// unlike update there are no array filters to apply.
pub async fn replace_document(
+58
View File
@@ -858,6 +858,58 @@ pub async fn mongo_replace_document_core(
}
}
pub async fn mongo_bulk_write_core(
state: &AppState,
connection_id: &str,
database: &str,
collection: &str,
operations_json: &str,
options_json: Option<&str>,
) -> Result<mongo_driver::MongoBulkWriteResult, String> {
ensure_document_pool(state, connection_id).await?;
let pool = state.pool_handle(connection_id).await.ok_or("Not found")?;
match &pool {
PoolKind::MongoDb(client) => {
mongo_driver::bulk_write(client, database, collection, operations_json, options_json).await
}
PoolKind::Agent(client) => {
let mut client = client.lock().await;
if !client.supports_capability(AgentCapability::MongoBulkWrite) {
return Err(
"MongoDB Legacy Agent does not support bulkWrite; upgrade or reinstall the MongoDB Legacy driver"
.to_string(),
);
}
client
.mongo_bulk_write(serde_json::json!({
"database": database,
"collection": collection,
"operations_json": operations_json,
"options_json": options_json,
}))
.await
}
_ => Err("Not a MongoDB connection".to_string()),
}
}
/// One row of counts, the way the shell prints a `BulkWriteResult`.
pub fn mongo_bulk_write_query_result(result: &mongo_driver::MongoBulkWriteResult) -> QueryResult {
let columns = ["insertedCount", "matchedCount", "modifiedCount", "deletedCount", "upsertedCount"];
let values = [
result.inserted_count,
result.matched_count,
result.modified_count,
result.deleted_count,
result.upserted_count,
];
query_result(
columns.iter().map(ToString::to_string).collect(),
vec![values.iter().map(|value| serde_json::Value::from(*value)).collect()],
result.inserted_count + result.modified_count + result.deleted_count + result.upserted_count,
)
}
pub async fn mongo_delete_document_core(
state: &AppState,
connection_id: &str,
@@ -1085,6 +1137,12 @@ pub async fn execute_mongo_command_core(
let affected = mongo_insert_documents_core(state, connection_id, database, collection, documents).await?;
Ok(affected_query_result(affected))
}
MongoCommand::BulkWrite { collection, operations, options } => {
let result =
mongo_bulk_write_core(state, connection_id, database, collection, operations, options.as_deref())
.await?;
Ok(mongo_bulk_write_query_result(&result))
}
MongoCommand::Replace { collection, filter, replacement, options } => {
let affected = mongo_replace_document_core(
state,
+275
View File
@@ -68,6 +68,8 @@ pub enum MongoCommand {
Update { collection: String, filter: String, update: String, options: Option<String>, many: bool },
#[serde(rename = "replace")]
Replace { collection: String, filter: String, replacement: String, options: Option<String> },
#[serde(rename = "bulkWrite")]
BulkWrite { collection: String, operations: String, options: Option<String> },
#[serde(rename = "delete")]
Delete { collection: String, filter: String, many: bool },
#[serde(rename = "createIndex")]
@@ -101,6 +103,7 @@ impl MongoCommand {
| Self::Insert { .. }
| Self::Update { .. }
| Self::Replace { .. }
| Self::BulkWrite { .. }
| Self::Delete { .. }
| Self::CreateIndex { .. }
| Self::DropIndexes { .. }
@@ -125,6 +128,9 @@ impl MongoCommand {
| Self::FindOneAndUpdate { filter, .. }
| Self::FindOneAndReplace { filter, .. }
| Self::FindOneAndDelete { filter, .. } => is_empty_object(filter),
Self::BulkWrite { operations, .. } => {
bulk_write_filters(operations).iter().any(|filter| is_empty_object(filter))
}
_ => false,
}
}
@@ -137,11 +143,160 @@ impl MongoCommand {
| Self::FindOneAndUpdate { filter, .. }
| Self::FindOneAndReplace { filter, .. }
| Self::FindOneAndDelete { filter, .. } => mongo_filter_is_effectively_unbounded(filter),
Self::BulkWrite { operations, .. } => {
bulk_write_filters(operations).iter().any(|filter| mongo_filter_is_effectively_unbounded(filter))
}
_ => false,
}
}
}
/// One entry of a `bulkWrite([...])` array, validated but not yet converted to BSON.
#[derive(Debug, Clone, PartialEq)]
pub enum BulkWriteOperation {
InsertOne { document: Value },
UpdateOne { filter: Value, update: Value, upsert: Option<bool>, array_filters: Option<Value> },
UpdateMany { filter: Value, update: Value, upsert: Option<bool>, array_filters: Option<Value> },
ReplaceOne { filter: Value, replacement: Value, upsert: Option<bool> },
DeleteOne { filter: Value },
DeleteMany { filter: Value },
}
impl BulkWriteOperation {
pub fn kind(&self) -> &'static str {
match self {
Self::InsertOne { .. } => "insertOne",
Self::UpdateOne { .. } => "updateOne",
Self::UpdateMany { .. } => "updateMany",
Self::ReplaceOne { .. } => "replaceOne",
Self::DeleteOne { .. } => "deleteOne",
Self::DeleteMany { .. } => "deleteMany",
}
}
fn filter(&self) -> Option<&Value> {
match self {
Self::InsertOne { .. } => None,
Self::UpdateOne { filter, .. }
| Self::UpdateMany { filter, .. }
| Self::ReplaceOne { filter, .. }
| Self::DeleteOne { filter }
| Self::DeleteMany { filter } => Some(filter),
}
}
}
/// Validate a `bulkWrite` operations array the way the shell does: each entry is one
/// `{ <op>: { ... } }` document with exactly the fields that operation takes.
pub fn parse_bulk_write_operations(operations_json: &str) -> Result<Vec<BulkWriteOperation>, String> {
let Some(Value::Array(entries)) = parse_json_value(operations_json) else {
return Err("MongoDB bulkWrite() requires an array of operations.".to_string());
};
if entries.is_empty() {
return Err("MongoDB bulkWrite() requires at least one operation.".to_string());
}
entries.iter().enumerate().map(|(index, entry)| parse_bulk_write_operation(index, entry)).collect()
}
fn parse_bulk_write_operation(index: usize, entry: &Value) -> Result<BulkWriteOperation, String> {
let position = index + 1;
let Value::Object(wrapper) = entry else {
return Err(format!("MongoDB bulkWrite() operation {position} must be a document such as {{ insertOne: {{ document: {{ ... }} }} }}."));
};
let (kind, spec) = match wrapper.iter().collect::<Vec<_>>().as_slice() {
[(kind, Value::Object(spec))] => (kind.as_str(), spec),
[(kind, _)] => return Err(format!("MongoDB bulkWrite() operation {position} ({kind}) must be a document.")),
_ => return Err(format!("MongoDB bulkWrite() operation {position} must have exactly one operation key.")),
};
let allowed: &[&str] = match kind {
"insertOne" => &["document"],
"updateOne" | "updateMany" => &["filter", "update", "upsert", "arrayFilters"],
"replaceOne" => &["filter", "replacement", "upsert"],
"deleteOne" | "deleteMany" => &["filter"],
other => {
return Err(format!(
"MongoDB bulkWrite() operation {position} uses unsupported operation {other}; supported: insertOne, updateOne, updateMany, replaceOne, deleteOne, deleteMany."
))
}
};
if let Some(unknown) = spec.keys().find(|key| !allowed.contains(&key.as_str())) {
return Err(format!("MongoDB bulkWrite() operation {position} ({kind}) has unsupported field {unknown}."));
}
let document = |field: &str| -> Result<Value, String> {
match spec.get(field) {
Some(value @ Value::Object(_)) => Ok(value.clone()),
Some(_) => {
Err(format!("MongoDB bulkWrite() operation {position} ({kind}) field {field} must be a document."))
}
None => Err(format!("MongoDB bulkWrite() operation {position} ({kind}) requires a {field} document.")),
}
};
let upsert = match spec.get("upsert") {
None => None,
Some(Value::Bool(value)) => Some(*value),
Some(_) => return Err(format!("MongoDB bulkWrite() operation {position} ({kind}) upsert must be a boolean.")),
};
match kind {
"insertOne" => Ok(BulkWriteOperation::InsertOne { document: document("document")? }),
"updateOne" | "updateMany" => {
let update = match spec.get("update") {
Some(value @ Value::Object(object)) => {
if object.is_empty() || !object.keys().all(|key| key.starts_with('$')) {
return Err(format!(
"MongoDB bulkWrite() operation {position} ({kind}) update must use operators such as $set; use replaceOne for a whole document."
));
}
value.clone()
}
Some(value @ Value::Array(_)) => value.clone(),
Some(_) => {
return Err(format!(
"MongoDB bulkWrite() operation {position} ({kind}) update must be a document or pipeline."
))
}
None => return Err(format!("MongoDB bulkWrite() operation {position} ({kind}) requires an update.")),
};
let array_filters = match spec.get("arrayFilters") {
None => None,
Some(value @ Value::Array(_)) => Some(value.clone()),
Some(_) => {
return Err(format!(
"MongoDB bulkWrite() operation {position} ({kind}) arrayFilters must be an array."
))
}
};
let filter = document("filter")?;
Ok(if kind == "updateOne" {
BulkWriteOperation::UpdateOne { filter, update, upsert, array_filters }
} else {
BulkWriteOperation::UpdateMany { filter, update, upsert, array_filters }
})
}
"replaceOne" => {
let replacement = document("replacement")?;
if let Some(operator) =
replacement.as_object().and_then(|object| object.keys().find(|key| key.starts_with('$')))
{
return Err(format!(
"MongoDB bulkWrite() operation {position} (replaceOne) replacement must not contain update operators such as {operator}; use updateOne to modify fields."
));
}
Ok(BulkWriteOperation::ReplaceOne { filter: document("filter")?, replacement, upsert })
}
"deleteOne" => Ok(BulkWriteOperation::DeleteOne { filter: document("filter")? }),
_ => Ok(BulkWriteOperation::DeleteMany { filter: document("filter")? }),
}
}
/// Filters of every non-insert operation, as JSON text, for the safety checks.
fn bulk_write_filters(operations_json: &str) -> Vec<String> {
parse_bulk_write_operations(operations_json)
.map(|operations| operations.iter().filter_map(|op| op.filter().map(Value::to_string)).collect())
.unwrap_or_default()
}
pub fn validate_safety(
command: &MongoCommand,
allow_writes: bool,
@@ -614,6 +769,19 @@ pub fn parse(input: &str) -> Result<MongoCommand, String> {
return Ok(MongoCommand::Insert { collection, documents });
}
if let Some((args, tail)) = method_call(source, prefix_end, "bulkWrite") {
if !tail.is_empty() || !(1..=2).contains(&args.len()) {
return Err("MongoDB bulkWrite() requires an array of operations and optional options.".to_string());
}
let operations = normalized_json(&args[0])?;
parse_bulk_write_operations(&operations)?;
let options = args.get(1).filter(|arg| !arg.trim().is_empty()).map(|arg| normalized_json(arg)).transpose()?;
if let Some(options) = &options {
validate_bulk_write_options(options)?;
}
return Ok(MongoCommand::BulkWrite { collection, operations, options });
}
if let Some((args, tail)) = method_call(source, prefix_end, "replaceOne") {
if !tail.is_empty() || !(2..=3).contains(&args.len()) {
return Err(
@@ -1369,6 +1537,21 @@ fn parse_use_database(source: &str) -> Option<String> {
Some(database.to_string())
}
/// `bulkWrite` options: only `ordered` is honoured, so anything else is rejected rather than dropped.
fn validate_bulk_write_options(options: &str) -> Result<(), String> {
let Some(Value::Object(object)) = parse_json_value(options) else {
return Err("MongoDB bulkWrite() options must be a document.".to_string());
};
for (key, value) in &object {
match (key.as_str(), value) {
("ordered", Value::Bool(_)) => {}
("ordered", _) => return Err("MongoDB bulkWrite() ordered option must be a boolean.".to_string()),
(other, _) => return Err(format!("Unsupported MongoDB bulkWrite() option: {other}.")),
}
}
Ok(())
}
/// A replacement is a whole document; `{$set: ...}` here almost always means updateOne() was intended.
fn require_replacement_document(replacement: &str) -> Result<(), String> {
let Some(Value::Object(document)) = parse_json_value(replacement) else {
@@ -1944,6 +2127,98 @@ mod tests {
}
}
#[test]
fn parses_bulk_write_with_every_operation_kind() {
let command = parse(
r#"db.orders.bulkWrite([
{ insertOne: { document: { sku: "A1", stock: 1 } } },
{ updateOne: { filter: { sku: "A1" }, update: { $inc: { stock: 1 } }, upsert: true } },
{ updateMany: { filter: { archived: true }, update: { $set: { stock: 0 } }, arrayFilters: [] } },
{ replaceOne: { filter: { sku: "B2" }, replacement: { sku: "B2", stock: 9 } } },
{ deleteOne: { filter: { sku: "C3" } } },
{ deleteMany: { filter: { stock: { $lt: 0 } } } }
], { ordered: false })"#,
)
.unwrap();
let MongoCommand::BulkWrite { collection, operations, options } = &command else {
panic!("expected a bulkWrite command");
};
assert_eq!(collection, "orders");
assert_eq!(options.as_deref(), Some(r#"{"ordered":false}"#));
let kinds: Vec<_> = parse_bulk_write_operations(operations).unwrap().iter().map(|op| op.kind()).collect();
assert_eq!(kinds, ["insertOne", "updateOne", "updateMany", "replaceOne", "deleteOne", "deleteMany"]);
assert!(command.is_mutating());
assert!(!command.is_dangerous());
assert!(!command.has_effectively_unbounded_filter());
assert_eq!(validate_safety(&command, true, false, false), Ok(()));
assert_eq!(validate_safety(&command, false, false, false), Err(MongoSafetyError::WritesDisabled));
// Only an insert: no filters at all, so never unbounded.
let inserts = parse(r#"db.orders.bulkWrite([{ insertOne: { document: { a: 1 } } }])"#).unwrap();
assert!(!inserts.has_empty_filter());
}
#[test]
fn bulk_write_with_any_unbounded_filter_is_guarded_like_an_update() {
let command = parse(
r#"db.orders.bulkWrite([
{ updateOne: { filter: { sku: "A1" }, update: { $set: { a: 1 } } } },
{ deleteMany: { filter: {} } }
])"#,
)
.unwrap();
assert!(command.has_empty_filter());
assert!(command.has_effectively_unbounded_filter());
assert_eq!(validate_safety(&command, true, false, false), Err(MongoSafetyError::EmptyFilter));
}
#[test]
fn rejects_malformed_bulk_write_operations() {
for (source, expected) in [
("db.orders.bulkWrite()", "array of operations"),
("db.orders.bulkWrite({})", "array of operations"),
("db.orders.bulkWrite([])", "at least one"),
("db.orders.bulkWrite([1])", "operation 1 must be a document"),
("db.orders.bulkWrite([{ insertOne: {}, deleteOne: {} }])", "exactly one operation key"),
("db.orders.bulkWrite([{ upsertOne: { document: {} } }])", "unsupported operation upsertOne"),
("db.orders.bulkWrite([{ insertOne: { doc: {} } }])", "unsupported field doc"),
("db.orders.bulkWrite([{ insertOne: {} }])", "requires a document document"),
(
"db.orders.bulkWrite([{ updateOne: { filter: {}, update: { a: 1 } } }])",
"must use operators such as $set",
),
(
"db.orders.bulkWrite([{ updateOne: { filter: {}, update: { $set: { a: 1 } }, upsert: 1 } }])",
"upsert must be a boolean",
),
(
"db.orders.bulkWrite([{ updateOne: { filter: {}, update: { $set: { a: 1 } }, collation: {} } }])",
"unsupported field collation",
),
(
"db.orders.bulkWrite([{ replaceOne: { filter: {}, replacement: { $set: { a: 1 } } } }])",
"must not contain update operators such as $set",
),
("db.orders.bulkWrite([{ deleteOne: { filter: [] } }])", "field filter must be a document"),
(
"db.orders.bulkWrite([{ deleteOne: { filter: {} } }], { ordered: 1 })",
"ordered option must be a boolean",
),
(
"db.orders.bulkWrite([{ deleteOne: { filter: {} } }], { writeConcern: {} })",
"Unsupported MongoDB bulkWrite() option: writeConcern",
),
(
"db.orders.bulkWrite([{ deleteOne: { filter: {} } }]).limit(1)",
"array of operations and optional options",
),
] {
let error = parse(source).unwrap_err();
assert!(error.contains(expected), "{source}\n expected: {expected}\n got: {error}");
}
}
#[test]
fn parses_show_databases_aliases_as_read_only() {
for source in [
+19
View File
@@ -1860,6 +1860,25 @@ impl DbxBackend for WebBackend {
.map_err(|error| format!("Invalid MongoDB insert response: {error}"))?;
Ok(affected_query_result(affected_rows_from_value(&value)))
}
MongoCommand::BulkWrite { collection, operations, options } => {
let result: dbx_core::db::mongo_driver::MongoBulkWriteResult = self
.request(
reqwest::Method::POST,
"/api/mongo/bulk-write",
Some(json!({
"connectionId": connection_id,
"database": database,
"collection": collection,
"operationsJson": operations,
"optionsJson": options,
})),
)
.await?
.json()
.await
.map_err(|error| format!("Invalid MongoDB bulkWrite response: {error}"))?;
Ok(dbx_core::mongo_ops::mongo_bulk_write_query_result(&result))
}
MongoCommand::Replace { collection, filter, replacement, options } => {
let value: Value = self
.request(
+1
View File
@@ -968,6 +968,7 @@ async fn main() {
.route("/mongo/update-document", post(routes::mongo::update_document))
.route("/mongo/update-documents", post(routes::mongo::update_documents))
.route("/mongo/replace-document", post(routes::mongo::replace_document))
.route("/mongo/bulk-write", post(routes::mongo::bulk_write))
.route("/mongo/delete-document", post(routes::mongo::delete_document))
.route("/mongo/delete-documents", post(routes::mongo::delete_documents))
.route("/mongo/find-one-and-update", post(routes::mongo::find_one_and_update))
+40
View File
@@ -269,6 +269,16 @@ pub struct MongoReplaceDocumentRequest {
pub options_json: Option<String>,
}
#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct MongoBulkWriteRequest {
pub connection_id: String,
pub database: String,
pub collection: String,
pub operations_json: String,
pub options_json: Option<String>,
}
#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct MongoDeleteDocumentsRequest {
@@ -846,6 +856,36 @@ pub async fn replace_document(
Ok(Json(serde_json::json!({ "affected_rows": result })))
}
pub async fn bulk_write(
State(state): State<Arc<WebState>>,
headers: HeaderMap,
Json(req): Json<MongoBulkWriteRequest>,
) -> Result<Json<dbx_core::db::mongo_driver::MongoBulkWriteResult>, AppError> {
let database = super::mcp_policy::resolve_database(&state, &headers, &req.connection_id, &req.database).await?;
let command = dbx_core::mongo_shell::MongoCommand::BulkWrite {
collection: req.collection.clone(),
operations: req.operations_json.clone(),
options: None,
};
if command.has_effectively_unbounded_filter() {
super::mcp_policy::ensure_dangerous_write(&state, &headers, &req.connection_id, &database, "BulkWrite").await?;
} else {
super::mcp_policy::ensure_write(&state, &headers, &req.connection_id, &database, "BulkWrite").await?;
}
ensure_writable(&state.app, &req.connection_id, "BulkWrite").await?;
let result = dbx_core::mongo_ops::mongo_bulk_write_core(
&state.app,
&req.connection_id,
&database,
&req.collection,
&req.operations_json,
req.options_json.as_deref(),
)
.await
.map_err(AppError::from)?;
Ok(Json(result))
}
pub async fn find_one_and_update(
State(state): State<Arc<WebState>>,
headers: HeaderMap,