feat(mongodb): support bulkWrite
This commit is contained in:
1 parent
3f3c21aae9
commit
f78dc1ba16
25 files changed
+1089
-18
No files matched your search
@@ -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 [
|
||||
|
||||
Reference in new issue
Block a user