diff --git a/agents/drivers/mongodb/src/main/java/com/dbx/agent/mongodb/MongoAgent.java b/agents/drivers/mongodb/src/main/java/com/dbx/agent/mongodb/MongoAgent.java index 743775736b..a3d004f7fb 100644 --- a/agents/drivers/mongodb/src/main/java/com/dbx/agent/mongodb/MongoAgent.java +++ b/agents/drivers/mongodb/src/main/java/com/dbx/agent/mongodb/MongoAgent.java @@ -3,6 +3,7 @@ package com.dbx.agent.mongodb; import com.dbx.agent.AgentProtocol; import com.dbx.agent.IndexInfo; import com.google.gson.Gson; +import com.google.gson.JsonArray; import com.google.gson.JsonElement; import com.google.gson.JsonNull; import com.google.gson.JsonObject; @@ -13,7 +14,9 @@ import com.mongodb.MongoClientSettings; import com.mongodb.ServerAddress; import com.mongodb.client.MongoClient; import com.mongodb.client.MongoClients; +import com.mongodb.client.MongoCollection; import com.mongodb.client.model.UpdateOptions; +import com.mongodb.client.result.UpdateResult; import java.io.BufferedReader; import java.io.FileInputStream; import java.io.InputStream; @@ -711,13 +714,26 @@ public final class MongoAgent { var col = c.getDatabase(database).getCollection(collection); Document filter = documentForWrite(filterJson); - Document update = documentForWrite(updateJson); - requireBulkUpdateOperatorDocument(update); UpdateOptions options = updateOptionsForWrite(optionsJson); - var result = many ? col.updateMany(filter, update, options) : col.updateOne(filter, update, options); + var result = isUpdatePipelineJson(updateJson) + ? updateDocumentsWithPipeline(col, filter, updatePipelineForWrite(updateJson), options, many) + : updateDocumentsWithDocument(col, filter, documentForWrite(updateJson), options, many); return Collections.singletonMap("modified_count", result.getModifiedCount()); } + private static UpdateResult updateDocumentsWithDocument( + MongoCollection col, Document filter, Document update, + UpdateOptions options, boolean many) { + requireBulkUpdateOperatorDocument(update); + return many ? col.updateMany(filter, update, options) : col.updateOne(filter, update, options); + } + + private static UpdateResult updateDocumentsWithPipeline( + MongoCollection col, Document filter, List pipeline, + UpdateOptions options, boolean many) { + return many ? col.updateMany(filter, pipeline, options) : col.updateOne(filter, pipeline, options); + } + static UpdateOptions updateOptionsForWrite(String optionsJson) { UpdateOptions result = new UpdateOptions(); if (optionsJson == null || optionsJson.trim().isEmpty()) { @@ -752,6 +768,27 @@ public final class MongoAgent { return doc; } + private static boolean isUpdatePipelineJson(String updateJson) { + return updateJson.trim().startsWith("["); + } + + static List updatePipelineForWrite(String updateJson) { + JsonElement parsed = JsonParser.parseString(updateJson); + if (!parsed.isJsonArray()) { + throw new IllegalArgumentException("Update pipeline must be an array"); + } + JsonArray stages = parsed.getAsJsonArray(); + List pipeline = new ArrayList<>(stages.size()); + for (JsonElement stage : stages) { + if (!stage.isJsonObject()) { + // The Java driver pipeline overload accepts BSON stages, not scalar array entries. + throw new IllegalArgumentException("Each update pipeline stage must be an object"); + } + pipeline.add(documentForWrite(stage.toString())); + } + return pipeline; + } + private static Document replacementDocument(Document doc) { Document replacement = new Document(doc); replacement.remove("_id"); @@ -980,28 +1017,41 @@ public final class MongoAgent { } static String handleRequest(String line) { - JsonObject req = JsonParser.parseString(line).getAsJsonObject(); - JsonElement id = req.get("id"); - String method = req.get("method").getAsString(); - JsonObject params = req.has("params") && req.get("params").isJsonObject() - ? req.getAsJsonObject("params") - : new JsonObject(); + return handleRequest(line, null); + } - JsonObject response = new JsonObject(); - response.addProperty("jsonrpc", "2.0"); - response.add("id", id); - - try { - Object result = dispatch(method, params); - response.add("result", GSON.toJsonTree(result)); - } catch (Exception e) { - JsonObject error = new JsonObject(); - error.addProperty("code", -1); - error.addProperty("message", e.getMessage() == null ? "Unknown error" : e.getMessage()); - response.add("error", error); + static String handleRequest(String line, MongoClient client) { + if (client != null) { + CURRENT_CLIENT.set(client); } + try { + JsonObject req = JsonParser.parseString(line).getAsJsonObject(); + JsonElement id = req.get("id"); + String method = req.get("method").getAsString(); + JsonObject params = req.has("params") && req.get("params").isJsonObject() + ? req.getAsJsonObject("params") + : new JsonObject(); - return GSON.toJson(response); + JsonObject response = new JsonObject(); + response.addProperty("jsonrpc", "2.0"); + response.add("id", id); + + try { + Object result = dispatch(method, params); + response.add("result", GSON.toJsonTree(result)); + } catch (Exception e) { + JsonObject error = new JsonObject(); + error.addProperty("code", -1); + error.addProperty("message", e.getMessage() == null ? "Unknown error" : e.getMessage()); + response.add("error", error); + } + + return GSON.toJson(response); + } finally { + if (client != null) { + CURRENT_CLIENT.remove(); + } + } } public static void main(String[] args) throws Exception { diff --git a/agents/drivers/mongodb/src/test/java/com/dbx/agent/mongodb/MongoAgentTest.java b/agents/drivers/mongodb/src/test/java/com/dbx/agent/mongodb/MongoAgentTest.java index d7f8a93df7..65428267cc 100644 --- a/agents/drivers/mongodb/src/test/java/com/dbx/agent/mongodb/MongoAgentTest.java +++ b/agents/drivers/mongodb/src/test/java/com/dbx/agent/mongodb/MongoAgentTest.java @@ -12,15 +12,22 @@ import com.google.gson.JsonArray; import com.google.gson.JsonObject; import com.google.gson.JsonParser; import com.mongodb.MongoClientSettings; +import com.mongodb.client.MongoClient; +import com.mongodb.client.MongoCollection; +import com.mongodb.client.MongoDatabase; import com.mongodb.client.model.UpdateOptions; +import com.mongodb.client.result.UpdateResult; import java.io.FileInputStream; +import java.lang.reflect.Proxy; import java.nio.file.Files; import java.nio.file.Path; import java.security.KeyStore; import java.security.PrivateKey; +import java.util.ArrayList; import java.util.Base64; import java.util.Collections; import java.util.Date; +import java.util.List; import org.bson.Document; import org.bson.types.ObjectId; import org.junit.jupiter.api.BeforeAll; @@ -255,6 +262,38 @@ class MongoAgentTest { assertFalse(json.getAsJsonObject("error").get("message").getAsString().contains("Unknown method")); } + @Test + void updateDocumentsRpcUsesDocumentOverloads() { + List calls = new ArrayList<>(); + MongoClient client = recordingMongoClient(calls); + + assertRpcModifiedCount(client, 20, "{\"$set\":{\"status\":\"done\"}}", false); + assertRpcModifiedCount(client, 21, "{\"$unset\":{\"legacy\":1}}", true); + + assertEquals(List.of("updateOne:document", "updateMany:document"), calls); + } + + @Test + void updateDocumentsRpcUsesPipelineOverloads() { + List calls = new ArrayList<>(); + MongoClient client = recordingMongoClient(calls); + + assertRpcModifiedCount(client, 22, "[{\"$set\":{\"status\":\"$source\"}}]", false); + assertRpcModifiedCount(client, 23, "[{\"$unset\":\"legacy\"}]", true); + + assertEquals(List.of("updateOne:pipeline", "updateMany:pipeline"), calls); + } + + @Test + void updatePipelineRejectsNonDocumentStages() { + IllegalArgumentException error = assertThrows( + IllegalArgumentException.class, + () -> MongoAgent.updatePipelineForWrite("[{\"$set\":{\"a\":1}}, 2]") + ); + + assertEquals("Each update pipeline stage must be an object", error.getMessage()); + } + @Test void parsesArrayFiltersUpdateOption() { UpdateOptions options = MongoAgent.updateOptionsForWrite( @@ -551,6 +590,63 @@ class MongoAgentTest { // ─── helpers ─── + private static void assertRpcModifiedCount(MongoClient client, int id, String updateJson, boolean many) { + JsonObject params = new JsonObject(); + params.addProperty("database", "app"); + params.addProperty("collection", "orders"); + params.addProperty("filter_json", "{}"); + params.addProperty("update_json", updateJson); + params.addProperty("many", many); + + JsonObject request = new JsonObject(); + request.addProperty("jsonrpc", "2.0"); + request.addProperty("id", id); + request.addProperty("method", "update_documents"); + request.add("params", params); + + JsonObject response = JsonParser.parseString(MongoAgent.handleRequest(request.toString(), client)).getAsJsonObject(); + assertFalse(response.has("error"), response.toString()); + assertEquals(1, response.getAsJsonObject("result").get("modified_count").getAsLong()); + } + + @SuppressWarnings("unchecked") + private static MongoClient recordingMongoClient(List calls) { + MongoCollection collection = (MongoCollection) Proxy.newProxyInstance( + MongoCollection.class.getClassLoader(), + new Class[] {MongoCollection.class}, + (proxy, method, args) -> { + if ("updateOne".equals(method.getName()) || "updateMany".equals(method.getName())) { + calls.add(method.getName() + ":" + (args[1] instanceof List ? "pipeline" : "document")); + return UpdateResult.acknowledged(1, 1L, null); + } + throw new UnsupportedOperationException(method.getName()); + } + ); + MongoDatabase database = (MongoDatabase) Proxy.newProxyInstance( + MongoDatabase.class.getClassLoader(), + new Class[] {MongoDatabase.class}, + (proxy, method, args) -> { + if ("getCollection".equals(method.getName())) { + return collection; + } + throw new UnsupportedOperationException(method.getName()); + } + ); + return (MongoClient) Proxy.newProxyInstance( + MongoClient.class.getClassLoader(), + new Class[] {MongoClient.class}, + (proxy, method, args) -> { + if ("getDatabase".equals(method.getName())) { + return database; + } + if ("close".equals(method.getName())) { + return null; + } + throw new UnsupportedOperationException(method.getName()); + } + ); + } + private static JsonObject minimalConnection() { JsonObject conn = new JsonObject(); conn.addProperty("host", "127.0.0.1"); diff --git a/crates/dbx-core/src/db/mongo_driver.rs b/crates/dbx-core/src/db/mongo_driver.rs index a1b5a15058..3b0e675528 100644 --- a/crates/dbx-core/src/db/mongo_driver.rs +++ b/crates/dbx-core/src/db/mongo_driver.rs @@ -1,6 +1,6 @@ use mongodb::{ bson::{doc, oid::ObjectId, Bson, DateTime, Document}, - options::{ClientOptions, GridFsBucketOptions, IndexOptions}, + options::{ClientOptions, GridFsBucketOptions, IndexOptions, UpdateModifications}, Client, Database, IndexModel, }; use serde::{Deserialize, Serialize}; @@ -1038,7 +1038,7 @@ pub async fn update_documents( 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 = json_object_to_document(&update_value).map_err(|e| format!("Invalid update: {e}"))?; + let update = json_update_to_modifications(&update_value).map_err(|e| format!("Invalid update: {e}"))?; let array_filters = parse_update_array_filters(options_json)?; let col = client.database(database).collection::(collection); let result = if many { @@ -1185,7 +1185,7 @@ pub async fn find_one_and_update( 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 = json_object_to_document(&update_value).map_err(|e| format!("Invalid update: {e}"))?; + let update = json_update_to_modifications(&update_value).map_err(|e| format!("Invalid update: {e}"))?; let options: MongoFindOneAndUpdateOptions = parse_find_and_modify_options(options_json, "findOneAndUpdate")?; let col = client.database(database).collection::(collection); let mut action = col.find_one_and_update(filter, update); @@ -1391,6 +1391,22 @@ pub fn json_object_to_document(value: &serde_json::Value) -> Result Result { + match json_value_to_bson(value) { + Bson::Document(document) => Ok(UpdateModifications::Document(document)), + Bson::Array(stages) => stages + .into_iter() + .enumerate() + .map(|(index, stage)| match stage { + Bson::Document(document) => Ok(document), + other => Err(format!("Update pipeline stage {} must be a JSON object, got {other:?}", index + 1)), + }) + .collect::, _>>() + .map(UpdateModifications::Pipeline), + other => Err(format!("Expected a JSON object or pipeline array, got {other:?}")), + } +} + fn json_object_to_document_extended_json(value: &serde_json::Value) -> Result { match Bson::try_from(value.clone()).map_err(|e| e.to_string())? { Bson::Document(doc) => Ok(doc), @@ -1660,6 +1676,47 @@ fn expand_object_id_string_array(items: &[serde_json::Value]) -> Bson { mod tests { use super::*; + #[test] + fn update_modifications_accept_document_and_pipeline() { + let document = json_update_to_modifications(&serde_json::json!({ "$set": { "status": "done" } })).unwrap(); + match document { + mongodb::options::UpdateModifications::Document(document) => { + assert_eq!(document, doc! { "$set": { "status": "done" } }); + } + other => panic!("expected document update, got {other:?}"), + } + + let owner_id = ObjectId::parse_str("507f1f77bcf86cd799439011").unwrap(); + let pipeline = json_update_to_modifications(&serde_json::json!([ + { "$set": { "update_date": { "$add": ["$update_date", 1000] } } }, + { "$set": { "owner_id": { "$oid": owner_id.to_hex() } } } + ])) + .unwrap(); + match pipeline { + mongodb::options::UpdateModifications::Pipeline(stages) => { + assert_eq!( + stages, + vec![ + doc! { "$set": { "update_date": { "$add": ["$update_date", 1000_i64] } } }, + doc! { "$set": { "owner_id": owner_id } }, + ] + ); + } + other => panic!("expected pipeline update, got {other:?}"), + } + } + + #[test] + fn update_modifications_reject_invalid_shapes() { + let stage_error = + json_update_to_modifications(&serde_json::json!([{ "$set": { "status": "done" } }, "invalid"])) + .unwrap_err(); + assert!(stage_error.contains("stage 2")); + + let value_error = json_update_to_modifications(&serde_json::json!("invalid")).unwrap_err(); + assert!(value_error.contains("object or pipeline array")); + } + #[test] fn update_options_parse_array_filters() { let filters = parse_update_array_filters(Some(r#"{"arrayFilters":[{"item.id":322678},{"item.active":true}]}"#))