diff --git a/agents/common/src/main/java/com/dbx/agent/AgentProtocol.java b/agents/common/src/main/java/com/dbx/agent/AgentProtocol.java index 9afc9de99..e0fb754b8 100644 --- a/agents/common/src/main/java/com/dbx/agent/AgentProtocol.java +++ b/agents/common/src/main/java/com/dbx/agent/AgentProtocol.java @@ -44,6 +44,8 @@ public final class AgentProtocol { */ public static final String MONGO_METHOD_FIND_DOCUMENTS_EXTENDED_JSON = "find_documents_extended_json"; public static final String MONGO_METHOD_SERVER_VERSION = "server_version"; + public static final String MONGO_METHOD_CREATE_INDEX = "create_index"; + public static final String MONGO_METHOD_DROP_INDEXES = "drop_indexes"; public static final String MONGO_METHOD_INSERT_DOCUMENT = "insert_document"; public static final String MONGO_METHOD_UPDATE_DOCUMENT = "update_document"; public static final String MONGO_METHOD_UPDATE_DOCUMENTS = "update_documents"; @@ -122,6 +124,8 @@ public final class AgentProtocol { MONGO_METHOD_FIND_DOCUMENTS, MONGO_METHOD_FIND_DOCUMENTS_EXTENDED_JSON, MONGO_METHOD_SERVER_VERSION, + MONGO_METHOD_CREATE_INDEX, + MONGO_METHOD_DROP_INDEXES, MONGO_METHOD_INSERT_DOCUMENT, MONGO_METHOD_UPDATE_DOCUMENT, MONGO_METHOD_UPDATE_DOCUMENTS, diff --git a/agents/common/src/main/resources/agent-protocol-v1.json b/agents/common/src/main/resources/agent-protocol-v1.json index d9f88a126..28535c1f2 100644 --- a/agents/common/src/main/resources/agent-protocol-v1.json +++ b/agents/common/src/main/resources/agent-protocol-v1.json @@ -70,6 +70,8 @@ "find_documents", "find_documents_extended_json", "server_version", + "create_index", + "drop_indexes", "insert_document", "update_document", "update_documents", 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 73843d864..9d7a8ff5e 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 @@ -36,9 +36,11 @@ import java.util.Base64; import java.util.Collection; import java.util.Collections; import java.util.Date; +import java.util.HashSet; import java.util.LinkedHashMap; import java.util.List; import java.util.Map; +import java.util.Set; import javax.net.ssl.KeyManager; import javax.net.ssl.KeyManagerFactory; import javax.net.ssl.SSLContext; @@ -441,6 +443,137 @@ public final class MongoAgent { return version; } + private static Object createIndex(JsonObject params) { + MongoClient c = requireClient(); + String database = params.get("database").getAsString(); + String collection = params.get("collection").getAsString(); + Document keys = requiredDocument(params, "keys_json", "Index keys"); + if (keys.isEmpty()) { + throw new IllegalArgumentException("Index keys are required"); + } + + Document index = new Document("key", keys); + Document options = documentOrNull(params, "options_json"); + if (options != null) { + index.putAll(options); + } + String name = index.getString("name"); + if (name == null || name.isBlank()) { + name = defaultIndexName(keys); + index.put("name", name); + } + + c.getDatabase(database).runCommand( + new Document("createIndexes", collection) + .append("indexes", Collections.singletonList(index)) + ); + return Collections.singletonMap("name", name); + } + + private static Document requiredDocument(JsonObject params, String key, String label) { + Document document = documentOrNull(params, key); + if (document == null) { + throw new IllegalArgumentException(label + " are required"); + } + return document; + } + + private static String defaultIndexName(Document keys) { + List parts = new ArrayList<>(); + for (Map.Entry entry : keys.entrySet()) { + parts.add(entry.getKey() + "_" + String.valueOf(entry.getValue())); + } + return String.join("_", parts); + } + + private static Object dropIndexes(JsonObject params) { + MongoClient c = requireClient(); + String database = params.get("database").getAsString(); + String collection = params.get("collection").getAsString(); + String indexesJson = stringOrNull(params, "indexes_json"); + boolean single = params.has("single") && !params.get("single").isJsonNull() && params.get("single").getAsBoolean(); + Object index = parseDropIndexesValue(indexesJson, single); + + List before = listIndexInfos(c, database, collection); + c.getDatabase(database).runCommand(new Document("dropIndexes", collection).append("index", index)); + List after = listIndexInfos(c, database, collection); + List droppedNames = diffDroppedIndexNames(before, after); + Map result = new LinkedHashMap<>(); + result.put("dropped_names", droppedNames); + result.put("affected_rows", droppedNames.size()); + return result; + } + + private static Object parseDropIndexesValue(String indexesJson, boolean single) { + if (indexesJson == null || indexesJson.isBlank()) { + if (single) { + throw new IllegalArgumentException("dropIndex requires a string index name or JSON document"); + } + return "*"; + } + + JsonElement value = JsonParser.parseString(indexesJson); + if (value.isJsonPrimitive() && value.getAsJsonPrimitive().isString()) { + String name = value.getAsString(); + if (name.isBlank()) { + throw new IllegalArgumentException("Index name is required"); + } + if (single && "*".equals(name)) { + throw new IllegalArgumentException("dropIndex does not accept \"*\"; use dropIndexes() or dropIndexes(\"*\") instead"); + } + return name; + } + if (value.isJsonObject()) { + JsonObject object = value.getAsJsonObject(); + if (object.size() == 0) { + throw new IllegalArgumentException("Index specification is required"); + } + return Document.parse(indexesJson); + } + if (value.isJsonArray()) { + if (single) { + throw new IllegalArgumentException("dropIndex only accepts a string index name or JSON document; arrays are not supported"); + } + List names = new ArrayList<>(); + value.getAsJsonArray().forEach(item -> { + if (!item.isJsonPrimitive() || !item.getAsJsonPrimitive().isString() || item.getAsString().isBlank()) { + throw new IllegalArgumentException("dropIndexes only accepts arrays of string index names"); + } + names.add(item.getAsString()); + }); + if (names.isEmpty()) { + throw new IllegalArgumentException("dropIndexes only accepts non-empty string arrays"); + } + return names; + } + if (single) { + throw new IllegalArgumentException("dropIndex only accepts a string index name or JSON document"); + } + throw new IllegalArgumentException("dropIndexes only accepts a string index name, JSON document, or string array"); + } + + private static List listIndexInfos(MongoClient c, String database, String collection) { + List result = new ArrayList<>(); + for (Document index : c.getDatabase(database).getCollection(collection).listIndexes()) { + result.add(indexInfoFromDocument(index)); + } + return result; + } + + private static List diffDroppedIndexNames(List before, List after) { + Set remaining = new HashSet<>(); + for (IndexInfo index : after) { + remaining.add(index.getName()); + } + List droppedNames = new ArrayList<>(); + for (IndexInfo index : before) { + if (!remaining.contains(index.getName())) { + droppedNames.add(index.getName()); + } + } + return droppedNames; + } + private static Object insertDocument(JsonObject params) { MongoClient c = requireClient(); String database = params.get("database").getAsString(); @@ -673,6 +806,8 @@ public final class MongoAgent { case AgentProtocol.MONGO_METHOD_FIND_DOCUMENTS -> findDocuments(params); case AgentProtocol.MONGO_METHOD_FIND_DOCUMENTS_EXTENDED_JSON -> findDocumentsExtendedJson(params); case AgentProtocol.MONGO_METHOD_SERVER_VERSION -> serverVersion(params); + case AgentProtocol.MONGO_METHOD_CREATE_INDEX -> createIndex(params); + case AgentProtocol.MONGO_METHOD_DROP_INDEXES -> dropIndexes(params); case AgentProtocol.MONGO_METHOD_INSERT_DOCUMENT -> insertDocument(params); case AgentProtocol.MONGO_METHOD_UPDATE_DOCUMENT -> updateDocument(params); case AgentProtocol.MONGO_METHOD_UPDATE_DOCUMENTS -> updateDocuments(params); 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 421abf030..bf1812603 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 @@ -148,6 +148,34 @@ class MongoAgentTest { assertFalse(json.getAsJsonObject("error").get("message").getAsString().contains("Unknown method")); } + @Test + void createIndexMethodIsRecognizedOverJsonRpc() { + String response = MongoAgent.handleRequest( + "{\"jsonrpc\":\"2.0\",\"id\":12,\"method\":\"create_index\"," + + "\"params\":{\"database\":\"app\",\"collection\":\"orders\"," + + "\"keys_json\":\"{\\\"email\\\":1}\",\"options_json\":\"{\\\"name\\\":\\\"email_1\\\",\\\"background\\\":true}\"}}"); + + JsonObject json = JsonParser.parseString(response).getAsJsonObject(); + assertEquals(12, json.get("id").getAsInt()); + assertEquals("Not connected", json.getAsJsonObject("error").get("message").getAsString()); + assertFalse(json.getAsJsonObject("error").get("message").getAsString().contains("Unknown method")); + assertTrue(AgentProtocol.MONGO_LEGACY_METHODS.contains(AgentProtocol.MONGO_METHOD_CREATE_INDEX)); + } + + @Test + void dropIndexesMethodIsRecognizedOverJsonRpc() { + String response = MongoAgent.handleRequest( + "{\"jsonrpc\":\"2.0\",\"id\":13,\"method\":\"drop_indexes\"," + + "\"params\":{\"database\":\"app\",\"collection\":\"orders\"," + + "\"indexes_json\":\"\\\"email_1\\\"\",\"single\":true}}"); + + JsonObject json = JsonParser.parseString(response).getAsJsonObject(); + assertEquals(13, json.get("id").getAsInt()); + assertEquals("Not connected", json.getAsJsonObject("error").get("message").getAsString()); + assertFalse(json.getAsJsonObject("error").get("message").getAsString().contains("Unknown method")); + assertTrue(AgentProtocol.MONGO_LEGACY_METHODS.contains(AgentProtocol.MONGO_METHOD_DROP_INDEXES)); + } + @Test void updateDocumentsMethodIsRecognizedOverJsonRpc() { String response = MongoAgent.handleRequest( diff --git a/crates/dbx-core/assets/agent-protocol-v1.json b/crates/dbx-core/assets/agent-protocol-v1.json index d9f88a126..28535c1f2 100644 --- a/crates/dbx-core/assets/agent-protocol-v1.json +++ b/crates/dbx-core/assets/agent-protocol-v1.json @@ -70,6 +70,8 @@ "find_documents", "find_documents_extended_json", "server_version", + "create_index", + "drop_indexes", "insert_document", "update_document", "update_documents", diff --git a/crates/dbx-core/src/db/agent_driver.rs b/crates/dbx-core/src/db/agent_driver.rs index 96ddcde53..a3a24922a 100644 --- a/crates/dbx-core/src/db/agent_driver.rs +++ b/crates/dbx-core/src/db/agent_driver.rs @@ -254,6 +254,8 @@ pub enum MongoAgentMethod { FindDocuments, FindDocumentsExtendedJson, ServerVersion, + CreateIndex, + DropIndexes, InsertDocument, UpdateDocument, UpdateDocuments, @@ -262,12 +264,14 @@ pub enum MongoAgentMethod { } impl MongoAgentMethod { - pub const ALL: [Self; 10] = [ + pub const ALL: [Self; 12] = [ Self::ListDatabases, Self::ListCollections, Self::FindDocuments, Self::FindDocumentsExtendedJson, Self::ServerVersion, + Self::CreateIndex, + Self::DropIndexes, Self::InsertDocument, Self::UpdateDocument, Self::UpdateDocuments, @@ -282,6 +286,8 @@ impl MongoAgentMethod { Self::FindDocuments => "find_documents", Self::FindDocumentsExtendedJson => "find_documents_extended_json", Self::ServerVersion => "server_version", + Self::CreateIndex => "create_index", + Self::DropIndexes => "drop_indexes", Self::InsertDocument => "insert_document", Self::UpdateDocument => "update_document", Self::UpdateDocuments => "update_documents", @@ -1004,6 +1010,20 @@ impl AgentDriverClient { self.call_mongo_method(MongoAgentMethod::ServerVersion, mongo_database_params(database)).await } + pub async fn mongo_create_index( + &mut self, + params: Value, + ) -> Result { + self.call_mongo_method(MongoAgentMethod::CreateIndex, params).await + } + + pub async fn mongo_drop_indexes( + &mut self, + params: Value, + ) -> Result { + self.call_mongo_method(MongoAgentMethod::DropIndexes, params).await + } + pub async fn mongo_insert_document( &mut self, params: Value, @@ -1676,6 +1696,8 @@ mod tests { assert_eq!(MongoAgentMethod::FindDocuments.as_str(), "find_documents"); assert_eq!(MongoAgentMethod::FindDocumentsExtendedJson.as_str(), "find_documents_extended_json"); assert_eq!(MongoAgentMethod::ServerVersion.as_str(), "server_version"); + assert_eq!(MongoAgentMethod::CreateIndex.as_str(), "create_index"); + assert_eq!(MongoAgentMethod::DropIndexes.as_str(), "drop_indexes"); assert_eq!(MongoAgentMethod::InsertDocument.as_str(), "insert_document"); assert_eq!(MongoAgentMethod::UpdateDocument.as_str(), "update_document"); assert_eq!(MongoAgentMethod::UpdateDocuments.as_str(), "update_documents"); @@ -1720,6 +1742,8 @@ mod tests { let _mongo_find_documents_extended_json = AgentDriverClient::mongo_find_documents_extended_json::; let _mongo_server_version = AgentDriverClient::mongo_server_version::; + let _mongo_create_index = AgentDriverClient::mongo_create_index::; + let _mongo_drop_indexes = AgentDriverClient::mongo_drop_indexes::; let _mongo_insert_document = AgentDriverClient::mongo_insert_document::; let _mongo_update_document = AgentDriverClient::mongo_update_document::; let _mongo_update_documents = AgentDriverClient::mongo_update_documents::; diff --git a/crates/dbx-core/src/mongo_ops.rs b/crates/dbx-core/src/mongo_ops.rs index 2969e9a6a..8f74d099e 100644 --- a/crates/dbx-core/src/mongo_ops.rs +++ b/crates/dbx-core/src/mongo_ops.rs @@ -197,7 +197,18 @@ pub async fn mongo_create_index_core( PoolKind::MongoDb(client) => { mongo_driver::create_index(client, database, collection, keys_json, options_json).await } - PoolKind::Agent(_) => Err("MongoDB legacy agent does not support createIndex".to_string()), + PoolKind::Agent(client) => { + let mut client = client.lock().await; + let result: serde_json::Value = client + .mongo_create_index(serde_json::json!({ + "database": database, + "collection": collection, + "keys_json": keys_json, + "options_json": options_json, + })) + .await?; + Ok(result.get("name").and_then(|value| value.as_str()).unwrap_or("").to_string()) + } _ => Err("Not a MongoDB connection".to_string()), } } @@ -216,7 +227,17 @@ pub async fn mongo_drop_indexes_core( PoolKind::MongoDb(client) => { mongo_driver::drop_indexes(client, database, collection, indexes_json, single).await } - PoolKind::Agent(_) => Err("MongoDB legacy agent does not support dropIndex/dropIndexes".to_string()), + PoolKind::Agent(client) => { + let mut client = client.lock().await; + client + .mongo_drop_indexes(serde_json::json!({ + "database": database, + "collection": collection, + "indexes_json": indexes_json, + "single": single, + })) + .await + } _ => Err("Not a MongoDB connection".to_string()), } }