From 98d3209026bc152b5aa96219281378a105878e48 Mon Sep 17 00:00:00 2001 From: Loc Nguyen <113083148+thailoc-dev@users.noreply.github.com> Date: Mon, 21 Sep 2026 12:37:09 +0700 Subject: [PATCH] feat(mongodb): find and modify over the legacy agent --- crates/dbx-core/src/query/mongo_ops.rs | 49 ++++- ...ve_mongodb_legacy_agent_find_and_modify.rs | 154 ++++++++++++++++ crates/dbx-drivers/src/db/mongo_driver.rs | 167 ++++++++++++++++++ 3 files changed, 367 insertions(+), 3 deletions(-) create mode 100644 crates/dbx-core/tests/live_mongodb_legacy_agent_find_and_modify.rs diff --git a/crates/dbx-core/src/query/mongo_ops.rs b/crates/dbx-core/src/query/mongo_ops.rs index db32304509..5ef467d73d 100644 --- a/crates/dbx-core/src/query/mongo_ops.rs +++ b/crates/dbx-core/src/query/mongo_ops.rs @@ -11,6 +11,35 @@ use crate::types::QueryResult; pub const MONGO_SHOW_DATABASES_DATABASE: &str = "admin"; pub const MONGO_SHOW_DATABASES_COMMAND_JSON: &str = r#"{"listDatabases":1}"#; +/// Runs one server command through the legacy agent's generic `runCommand` and returns the +/// response document. The agent surfaces `ok: 0` as an error, so a returned document succeeded. +async fn agent_run_command_document( + client: &crate::db::agent_driver::PooledAgentClient, + database: &str, + command: mongodb::bson::Document, + label: &str, +) -> Result { + let mut client = client.lock().await; + if !client.supports_capability(AgentCapability::MongoRunCommand) { + return Err(format!( + "MongoDB Legacy Agent does not support {label}; upgrade or reinstall the MongoDB Legacy driver" + )); + } + let result: MongoDocumentResult = client + .mongo_run_command(serde_json::json!({ + "database": database, + "command_json": mongo_driver::document_to_canonical_extended_json(&command).to_string(), + })) + .await?; + let response = result + .extended_documents + .as_ref() + .and_then(|documents| documents.first()) + .or_else(|| result.documents.first()) + .ok_or_else(|| format!("MongoDB Legacy Agent returned an empty {label} response"))?; + mongo_driver::json_object_to_document_extended_json(response) +} + async fn ensure_document_pool(state: &AppState, connection_id: &str) -> Result<(), String> { state.get_or_create_pool(connection_id, None).await.map(|_| ()) } @@ -983,7 +1012,12 @@ pub async fn mongo_find_one_and_update_core( mongo_driver::find_one_and_update(client, database, collection, filter_json, update_json, options_json) .await } - PoolKind::Agent(_) => Err("MongoDB legacy agent does not support findOneAndUpdate".to_string()), + PoolKind::Agent(client) => { + let command = + mongo_driver::find_one_and_update_command(collection, filter_json, update_json, options_json)?; + let response = agent_run_command_document(client, database, command, "findOneAndUpdate").await?; + mongo_driver::find_and_modify_response_result(&response) + } _ => Err("Not a MongoDB connection".to_string()), } } @@ -1011,7 +1045,12 @@ pub async fn mongo_find_one_and_replace_core( ) .await } - PoolKind::Agent(_) => Err("MongoDB legacy agent does not support findOneAndReplace".to_string()), + PoolKind::Agent(client) => { + let command = + mongo_driver::find_one_and_replace_command(collection, filter_json, replacement_json, options_json)?; + let response = agent_run_command_document(client, database, command, "findOneAndReplace").await?; + mongo_driver::find_and_modify_response_result(&response) + } _ => Err("Not a MongoDB connection".to_string()), } } @@ -1030,7 +1069,11 @@ pub async fn mongo_find_one_and_delete_core( PoolKind::MongoDb(client) => { mongo_driver::find_one_and_delete(client, database, collection, filter_json, options_json).await } - PoolKind::Agent(_) => Err("MongoDB legacy agent does not support findOneAndDelete".to_string()), + PoolKind::Agent(client) => { + let command = mongo_driver::find_one_and_delete_command(collection, filter_json, options_json)?; + let response = agent_run_command_document(client, database, command, "findOneAndDelete").await?; + mongo_driver::find_and_modify_response_result(&response) + } _ => Err("Not a MongoDB connection".to_string()), } } diff --git a/crates/dbx-core/tests/live_mongodb_legacy_agent_find_and_modify.rs b/crates/dbx-core/tests/live_mongodb_legacy_agent_find_and_modify.rs new file mode 100644 index 0000000000..1b423a3d8a --- /dev/null +++ b/crates/dbx-core/tests/live_mongodb_legacy_agent_find_and_modify.rs @@ -0,0 +1,154 @@ +//! `findOneAndUpdate` / `findOneAndReplace` / `findOneAndDelete` over the MongoDB Legacy Agent, +//! driven through the same parse → dispatch path the query tab uses. +use dbx_core::{ + connection::{AppState, PoolKind}, + models::connection::ConnectionConfig, + mongo_ops::{execute_mongo_command_core, mongo_run_command_core}, + mongo_shell, + storage::Storage, +}; +use mongodb::bson::{doc, Document}; + +async fn command(state: &AppState, id: &str, database: &str, command: Document) -> Document { + let json = mongodb::bson::Bson::Document(command).into_canonical_extjson().to_string(); + let result = mongo_run_command_core(state, id, database, &json).await.unwrap(); + dbx_core::db::mongo_driver::json_object_to_document_extended_json(&result.extended_documents.unwrap()[0]).unwrap() +} + +/// `runCommand` replies come back through the agent as relaxed Extended JSON, so integers lose +/// their width; compare numerically. +fn number(document: &Document, key: &str) -> Option { + match document.get(key)? { + mongodb::bson::Bson::Int32(n) => Some(i64::from(*n)), + mongodb::bson::Bson::Int64(n) => Some(*n), + mongodb::bson::Bson::Double(n) => Some(*n as i64), + _ => None, + } +} + +async fn find_all(state: &AppState, id: &str, database: &str, collection: &str) -> Vec { + let response = command(state, id, database, doc! { "find": collection, "sort": { "_id": 1 } }).await; + response + .get_document("cursor") + .unwrap() + .get_array("firstBatch") + .unwrap() + .iter() + .map(|b| b.as_document().unwrap().clone()) + .collect() +} + +async fn shell(state: &AppState, id: &str, database: &str, text: &str) -> Result { + let parsed = mongo_shell::parse(text)?; + execute_mongo_command_core(state, id, database, &parsed, 100).await +} + +#[tokio::test] +#[ignore = "opt-in: DBX_MONGO_LEGACY_DUMP_TEST_HOST (host:port, MongoDB 3.6+ without auth) and an installed MongoDB Legacy Agent; creates a temporary database"] +async fn find_and_modify_commands_run_over_the_legacy_agent() { + let endpoint = std::env::var("DBX_MONGO_LEGACY_DUMP_TEST_HOST").expect("DBX_MONGO_LEGACY_DUMP_TEST_HOST"); + let (host, port) = endpoint.split_once(':').expect("host:port"); + let files = tempfile::tempdir().unwrap(); + let database = format!("dbx_legacy_fam_{}", uuid::Uuid::new_v4().simple()); + let state = AppState::new(Storage::open(&files.path().join("storage.db")).await.unwrap()); + let id = "legacy-find-and-modify-test"; + let config: ConnectionConfig = serde_json::from_value(serde_json::json!({ "id": id, "name": "Legacy findAndModify test", "db_type": "mongodb", "host": host, "port": port.parse::().unwrap(), "username": "", "password": "", "database": database, "driver_profile": "mongodb-legacy" })).unwrap(); + state.configs.write().await.insert(id.into(), config); + let key = state.get_or_create_pool(id, Some(&database)).await.unwrap(); + assert!(matches!(state.pool_handle(&key).await, Some(PoolKind::Agent(_))), "test must run over the legacy agent"); + + command(&state, id, &database, doc! { "insert": "users", "documents": [ + { "_id": 1, "name": "Ada", "score": 10i64, "tags": [ { "k": "a", "on": false }, { "k": "b", "on": false } ] }, + { "_id": 2, "name": "Grace", "score": 20i64 }, + ] }).await; + + // findOneAndUpdate: default returns the document before the change. + let before = shell(&state, id, &database, "db.users.findOneAndUpdate({_id: 1}, {$inc: {score: 5}})").await.unwrap(); + assert_eq!(before.rows.len(), 1, "{before:?}"); + assert!(format!("{:?}", before.rows[0]).contains("10"), "pre-change score expected: {:?}", before.rows[0]); + // …and with returnDocument: 'after', sort and projection: sort {_id: -1} must pick Grace. + let after = shell( + &state, + id, + &database, + r#"db.users.findOneAndUpdate({}, {$set: {level: 2}}, {sort: {_id: -1}, returnDocument: "after", projection: {_id: 1, name: 1, level: 1}})"#, + ) + .await + .unwrap(); + assert_eq!(after.rows.len(), 1); + let shown = format!("{:?}", after.rows[0]); + assert!(shown.contains("Grace") && shown.contains("2") && !shown.contains("score"), "{shown}"); + let docs = find_all(&state, id, &database, "users").await; + assert_eq!(number(&docs[0], "score"), Some(15)); + assert!(!docs[0].contains_key("level"), "sort honoured: Ada untouched"); + assert_eq!(number(&docs[1], "level"), Some(2)); + + // arrayFilters must flip only tag "a". + shell( + &state, + id, + &database, + r#"db.users.findOneAndUpdate({_id: 1}, {$set: {"tags.$[t].on": true}}, {arrayFilters: [{"t.k": "a"}]})"#, + ) + .await + .unwrap(); + let ada = &find_all(&state, id, &database, "users").await[0]; + assert_eq!(ada.get_array("tags").unwrap()[0].as_document().unwrap().get_bool("on"), Ok(true)); + assert_eq!(ada.get_array("tags").unwrap()[1].as_document().unwrap().get_bool("on"), Ok(false)); + + // A server-side failure comes back as the server's own message. + let server_error = shell( + &state, + id, + &database, + r#"db.users.findOneAndUpdate({_id: 2}, {$set: {"tags.$[t].on": true}}, {arrayFilters: [{"t.k": "a"}]})"#, + ) + .await + .unwrap_err(); + assert!(server_error.contains("must exist in the document"), "{server_error}"); + + // upsert creates when nothing matches. + let upserted = shell( + &state, + id, + &database, + r#"db.users.findOneAndUpdate({_id: 3}, {$set: {name: "Linus"}}, {upsert: true, new: true})"#, + ) + .await + .unwrap(); + assert_eq!(upserted.rows.len(), 1); + assert_eq!(find_all(&state, id, &database, "users").await.len(), 3); + + // No match, no upsert: empty result, not an error. + let missing = shell(&state, id, &database, "db.users.findOneAndUpdate({_id: 99}, {$set: {x: 1}})").await.unwrap(); + assert_eq!(missing.rows.len(), 0); + + // findOneAndReplace replaces the whole document. + let replaced = shell( + &state, + id, + &database, + r#"db.users.findOneAndReplace({_id: 2}, {name: "Grace H", level: 1}, {returnDocument: "after"})"#, + ) + .await + .unwrap(); + assert_eq!(replaced.rows.len(), 1); + let grace = &find_all(&state, id, &database, "users").await[1]; + assert_eq!(grace.get_str("name"), Ok("Grace H")); + assert!(!grace.contains_key("score"), "replacement drops old fields: {grace}"); + // A replacement carrying update operators is refused before reaching the server. + let refused = + shell(&state, id, &database, r#"db.users.findOneAndReplace({_id: 2}, {$set: {name: "x"}})"#).await.unwrap_err(); + assert!(refused.contains("update operators"), "{refused}"); + + // findOneAndDelete removes and returns the matched document (with projection and sort). + let deleted = + shell(&state, id, &database, r#"db.users.findOneAndDelete({}, {sort: {_id: -1}, projection: {name: 1}})"#) + .await + .unwrap(); + assert_eq!(deleted.rows.len(), 1); + assert!(format!("{:?}", deleted.rows[0]).contains("Linus"), "highest _id deleted first: {:?}", deleted.rows[0]); + assert_eq!(find_all(&state, id, &database, "users").await.len(), 2); + + command(&state, id, &database, doc! { "dropDatabase": 1 }).await; +} diff --git a/crates/dbx-drivers/src/db/mongo_driver.rs b/crates/dbx-drivers/src/db/mongo_driver.rs index 687c33c6d0..ac1f05ebc0 100644 --- a/crates/dbx-drivers/src/db/mongo_driver.rs +++ b/crates/dbx-drivers/src/db/mongo_driver.rs @@ -2496,6 +2496,116 @@ fn single_document_result(document: Option) -> MongoDocumentResult { } } +/// The `findAndModify` server command behind `findOneAndUpdate()`, for drivers that only expose +/// `runCommand` (the legacy agent). Same inputs and option validation as [`find_one_and_update`]. +pub fn find_one_and_update_command( + collection: &str, + filter_json: &str, + update_json: &str, + options_json: Option<&str>, +) -> Result { + let filter_value: serde_json::Value = + serde_json::from_str(filter_json).map_err(|e| format!("Invalid filter JSON: {e}"))?; + let update_value: serde_json::Value = + serde_json::from_str(update_json).map_err(|e| format!("Invalid update JSON: {e}"))?; + let filter = json_filter_to_document(&filter_value).map_err(|e| format!("Invalid filter: {e}"))?; + let update = match json_update_to_modifications(&update_value).map_err(|e| format!("Invalid update: {e}"))? { + UpdateModifications::Document(document) => Bson::Document(document), + UpdateModifications::Pipeline(stages) => Bson::Array(stages.into_iter().map(Bson::Document).collect()), + _ => return Err("Unsupported update modification".to_string()), + }; + let options: MongoFindOneAndUpdateOptions = parse_find_and_modify_options(options_json, "findOneAndUpdate")?; + let mut command = doc! { "findAndModify": collection, "query": filter, "update": update }; + if find_and_modify_returns_after(options.return_document.as_deref(), options.return_new_document, options.new)? { + command.insert("new", true); + } + if let Some(upsert) = options.upsert { + command.insert("upsert", upsert); + } + if let Some(projection) = parse_optional_document(options.projection.as_ref(), "projection")? { + command.insert("fields", projection); + } + if let Some(sort) = parse_optional_document(options.sort.as_ref(), "sort")? { + command.insert("sort", sort); + } + if let Some(array_filters) = find_and_modify_array_filters(options.array_filters.as_ref())? { + command.insert("arrayFilters", array_filters); + } + Ok(command) +} + +/// A replacement document whose keys start with `$` would be silently executed as +/// an update by the server, so both the native and legacy-agent `findOneAndReplace` +/// paths refuse it up front with the same message. +fn reject_update_operator_replacement(replacement: &Document) -> Result<(), String> { + if let Some(operator) = replacement.keys().find(|key| key.starts_with('$')) { + return Err(format!("Replacement document must not contain update operators such as {operator}")); + } + Ok(()) +} + +/// The `findAndModify` server command behind `findOneAndReplace()`. A replacement with update +/// operators is refused here because the server would silently treat it as an update instead. +pub fn find_one_and_replace_command( + collection: &str, + filter_json: &str, + replacement_json: &str, + options_json: Option<&str>, +) -> Result { + let filter_value: serde_json::Value = + serde_json::from_str(filter_json).map_err(|e| format!("Invalid filter JSON: {e}"))?; + let replacement_value: serde_json::Value = + serde_json::from_str(replacement_json).map_err(|e| format!("Invalid replacement JSON: {e}"))?; + let filter = json_filter_to_document(&filter_value).map_err(|e| format!("Invalid filter: {e}"))?; + let replacement = json_object_to_document(&replacement_value).map_err(|e| format!("Invalid replacement: {e}"))?; + reject_update_operator_replacement(&replacement)?; + let options: MongoFindOneAndReplaceOptions = parse_find_and_modify_options(options_json, "findOneAndReplace")?; + let mut command = doc! { "findAndModify": collection, "query": filter, "update": replacement }; + if find_and_modify_returns_after(options.return_document.as_deref(), options.return_new_document, options.new)? { + command.insert("new", true); + } + if let Some(upsert) = options.upsert { + command.insert("upsert", upsert); + } + if let Some(projection) = parse_optional_document(options.projection.as_ref(), "projection")? { + command.insert("fields", projection); + } + if let Some(sort) = parse_optional_document(options.sort.as_ref(), "sort")? { + command.insert("sort", sort); + } + Ok(command) +} + +/// The `findAndModify` server command behind `findOneAndDelete()`. +pub fn find_one_and_delete_command( + collection: &str, + filter_json: &str, + options_json: Option<&str>, +) -> Result { + let filter_value: serde_json::Value = + serde_json::from_str(filter_json).map_err(|e| format!("Invalid filter JSON: {e}"))?; + let filter = json_filter_to_document(&filter_value).map_err(|e| format!("Invalid filter: {e}"))?; + let options: MongoFindOneAndDeleteOptions = parse_find_and_modify_options(options_json, "findOneAndDelete")?; + let mut command = doc! { "findAndModify": collection, "query": filter, "remove": true }; + if let Some(projection) = parse_optional_document(options.projection.as_ref(), "projection")? { + command.insert("fields", projection); + } + if let Some(sort) = parse_optional_document(options.sort.as_ref(), "sort")? { + command.insert("sort", sort); + } + Ok(command) +} + +/// Turns a `findAndModify` response into the same shape the native helpers return: the matched +/// document (before or after the change, as requested) or nothing. +pub fn find_and_modify_response_result(response: &Document) -> Result { + match response.get("value") { + Some(Bson::Document(document)) => Ok(single_document_result(Some(document.clone()))), + Some(Bson::Null) | None => Ok(single_document_result(None)), + Some(other) => Err(format!("Unexpected findAndModify value: {other:?}")), + } +} + pub async fn find_one_and_update( client: &Client, database: &str, @@ -2546,6 +2656,7 @@ pub async fn find_one_and_replace( serde_json::from_str(replacement_json).map_err(|e| format!("Invalid replacement JSON: {e}"))?; let filter = json_filter_to_document(&filter_value).map_err(|e| format!("Invalid filter: {e}"))?; let replacement = json_object_to_document(&replacement_value).map_err(|e| format!("Invalid replacement: {e}"))?; + reject_update_operator_replacement(&replacement)?; let options: MongoFindOneAndReplaceOptions = parse_find_and_modify_options(options_json, "findOneAndReplace")?; let col = client.database(database).collection::(collection); let mut action = col.find_one_and_replace(filter, replacement); @@ -4939,6 +5050,62 @@ mod tests { assert!(find_and_modify_returns_after(Some("newest"), None, None).is_err()); } + #[test] + fn find_and_modify_commands_mirror_the_native_options() { + let update = find_one_and_update_command( + "users", + r#"{"_id":{"$oid":"64b1f0c2a1b2c3d4e5f60718"}}"#, + r#"{"$set":{"active":true}}"#, + Some(r#"{"returnDocument":"after","upsert":true,"projection":{"active":1},"sort":{"_id":1},"arrayFilters":[{"item.active":true}]}"#), + ) + .unwrap(); + assert_eq!(update.get_str("findAndModify"), Ok("users")); + assert_eq!( + update.get_document("query").unwrap().get_object_id("_id").unwrap().to_hex(), + "64b1f0c2a1b2c3d4e5f60718" + ); + assert_eq!(update.get_document("update"), Ok(&doc! { "$set": { "active": true } })); + assert_eq!(update.get_bool("new"), Ok(true)); + assert_eq!(update.get_bool("upsert"), Ok(true)); + assert_eq!(update.get_document("fields"), Ok(&doc! { "active": 1i64 })); + assert_eq!(update.get_document("sort"), Ok(&doc! { "_id": 1i64 })); + assert_eq!(update.get_array("arrayFilters").unwrap().len(), 1); + + // Default return is the pre-change document, and absent options stay absent. + let minimal = find_one_and_update_command("users", "{}", r#"{"$inc":{"n":1}}"#, None).unwrap(); + assert_eq!(minimal.keys().collect::>(), vec!["findAndModify", "query", "update"]); + + // A pipeline update is passed through as an array. + let pipeline = find_one_and_update_command("users", "{}", r#"[{"$set":{"n":1}}]"#, None).unwrap(); + assert!(pipeline.get_array("update").is_ok()); + + let replace = + find_one_and_replace_command("users", r#"{"_id":1}"#, r#"{"name":"Grace"}"#, Some(r#"{"new":true}"#)) + .unwrap(); + assert_eq!(replace.get_document("update"), Ok(&doc! { "name": "Grace" })); + assert_eq!(replace.get_bool("new"), Ok(true)); + assert!(!replace.contains_key("remove")); + // The server would run an operator document as an update; refuse it up front instead. + let refused = find_one_and_replace_command("users", "{}", r#"{"$set":{"name":"Grace"}}"#, None).unwrap_err(); + assert!(refused.contains("must not contain update operators such as $set"), "{refused}"); + + let delete = find_one_and_delete_command("users", r#"{"_id":1}"#, Some(r#"{"sort":{"_id":-1}}"#)).unwrap(); + assert_eq!(delete.get_bool("remove"), Ok(true)); + assert_eq!(delete.get_document("sort"), Ok(&doc! { "_id": -1i64 })); + assert!(!delete.contains_key("update")); + } + + #[test] + fn find_and_modify_response_maps_value_to_zero_or_one_document() { + let matched = + find_and_modify_response_result(&doc! { "value": { "_id": 1, "name": "Ada" }, "ok": 1.0 }).unwrap(); + assert_eq!(matched.total, 1); + assert_eq!(matched.documents[0]["name"], "Ada"); + let none = find_and_modify_response_result(&doc! { "value": Bson::Null, "ok": 1.0 }).unwrap(); + assert_eq!(none.total, 0); + assert!(find_and_modify_response_result(&doc! { "value": "oops", "ok": 1.0 }).is_err()); + } + #[test] fn single_document_result_has_zero_or_one_metadata() { let empty = single_document_result(None);