github/dbx
4
0
Fork 0
mirror of https://github.com/t8y2/dbx.git synced 2026-10-20 00:17:44 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(mongodb): support legacy agent insertMany writes

Browse Source 
Refs: #6032
This commit is contained in:
t8y2
2026-08-12 16:20:25 +00:00
parent 7dabb9e4fb
commit d0e027bb17
9 changed files with 316 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
agents/common/src/main/java/com/dbx/agent/AgentProtocol.java
+7 -1
View File
@@ -64,6 +64,7 @@ public final class AgentProtocol {
public static final String MONGO_METHOD_CLONE_COLLECTION = "clone_collection";
public static final String MONGO_METHOD_DROP_DATABASE = "drop_database";
public static final String MONGO_METHOD_INSERT_DOCUMENT = "insert_document";
public static final String MONGO_METHOD_INSERT_DOCUMENTS = "insert_documents";
public static final String MONGO_METHOD_UPDATE_DOCUMENT = "update_document";
public static final String MONGO_METHOD_UPDATE_DOCUMENTS = "update_documents";
public static final String MONGO_METHOD_DELETE_DOCUMENT = "delete_document";
@@ -122,6 +123,7 @@ public final class AgentProtocol {
public static final String CAPABILITY_MONGO_DROP_DATABASE = "mongo_drop_database";
public static final String CAPABILITY_MONGO_CLONE_COLLECTION = "mongo_clone_collection";
public static final String CAPABILITY_MONGO_RUN_COMMAND = "mongo_run_command";
public static final String CAPABILITY_MONGO_INSERT_DOCUMENTS = "mongo_insert_documents";
public static final String CAPABILITY_MULTI_SESSION = "multi_session";
public static final String CAPABILITY_STRUCTURED_ERROR_V1 = "structured_error_v1";
@@ -156,7 +158,8 @@ public final class AgentProtocol {
CAPABILITY_ETCD_AUTH,
CAPABILITY_MONGO_DROP_DATABASE,
CAPABILITY_MONGO_CLONE_COLLECTION,
CAPABILITY_MONGO_RUN_COMMAND
CAPABILITY_MONGO_RUN_COMMAND,
CAPABILITY_MONGO_INSERT_DOCUMENTS
));
public static final List<String> MULTI_SESSION_CAPABILITIES;
@@ -216,12 +219,14 @@ public final class AgentProtocol {
mongoCapabilities.add(CAPABILITY_MONGO_DROP_DATABASE);
mongoCapabilities.add(CAPABILITY_MONGO_CLONE_COLLECTION);
mongoCapabilities.add(CAPABILITY_MONGO_RUN_COMMAND);
mongoCapabilities.add(CAPABILITY_MONGO_INSERT_DOCUMENTS);
MONGO_LEGACY_CAPABILITIES = Collections.unmodifiableList(mongoCapabilities);
List<String> mongoMultiSessionCapabilities = new java.util.ArrayList<>(MULTI_SESSION_CAPABILITIES);
mongoMultiSessionCapabilities.add(CAPABILITY_MONGO_DROP_DATABASE);
mongoMultiSessionCapabilities.add(CAPABILITY_MONGO_CLONE_COLLECTION);
mongoMultiSessionCapabilities.add(CAPABILITY_MONGO_RUN_COMMAND);
mongoMultiSessionCapabilities.add(CAPABILITY_MONGO_INSERT_DOCUMENTS);
MONGO_LEGACY_MULTI_SESSION_CAPABILITIES = Collections.unmodifiableList(mongoMultiSessionCapabilities);
List<String> jdbcCapabilities = new java.util.ArrayList<>(MULTI_SESSION_CAPABILITIES);
@@ -260,6 +265,7 @@ public final class AgentProtocol {
MONGO_METHOD_CLONE_COLLECTION,
MONGO_METHOD_DROP_DATABASE,
MONGO_METHOD_INSERT_DOCUMENT,
MONGO_METHOD_INSERT_DOCUMENTS,
MONGO_METHOD_UPDATE_DOCUMENT,
MONGO_METHOD_UPDATE_DOCUMENTS,
MONGO_METHOD_DELETE_DOCUMENT,
agents/common/src/main/resources/agent-protocol-v1.json
+2 -2
View File
@@ -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"],
"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"],
"capabilities": ["connect", "test_connection", "metadata", "query", "paged_query", "transaction", "ddl"],
"defaultSqlCapabilities": ["connect", "test_connection", "metadata", "query", "paged_query", "transaction", "ddl"],
"commonMethods": [
@@ -39,6 +39,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", "update_document", "update_documents", "delete_document", "delete_documents", "run_command"],
"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", "delete_document", "delete_documents", "run_command"],
"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"]
}
agents/common/src/main/resources/agent-protocol-v2.json
+2 -2
View File
@@ -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", "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", "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": [
@@ -43,7 +43,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", "update_document", "update_documents", "delete_document", "delete_documents", "run_command"],
"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", "delete_document", "delete_documents", "run_command"],
"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"
agents/drivers/mongodb/src/main/java/com/dbx/agent/mongodb/MongoAgent.java
+26
View File
@@ -1467,6 +1467,31 @@ public final class MongoAgent {
return Collections.singletonMap("inserted_id", insertedId);
}
private static Object insertDocuments(JsonObject params) {
String docsJson = params.get("docs_json").getAsString();
JsonElement parsed = JsonParser.parseString(docsJson);
if (!parsed.isJsonArray()) {
throw new IllegalArgumentException("MongoDB insertMany documents must be a JSON array");
}
List<Document> documents = new ArrayList<>();
for (JsonElement item : parsed.getAsJsonArray()) {
if (!item.isJsonObject()) {
throw new IllegalArgumentException("Each MongoDB insertMany document must be an object");
}
documents.add(documentForWrite(item.toString()));
}
if (documents.isEmpty()) {
return Collections.singletonMap("affected_rows", 0);
}
MongoClient client = requireClient();
String database = params.get("database").getAsString();
String collection = params.get("collection").getAsString();
client.getDatabase(database).getCollection(collection).insertMany(documents);
return Collections.singletonMap("affected_rows", documents.size());
}
static Object parseId(String id) {
String stringId = decodeStringDocumentId(id);
if (stringId != null) {
@@ -1830,6 +1855,7 @@ public final class MongoAgent {
case AgentProtocol.MONGO_METHOD_CLONE_COLLECTION -> cloneCollection(params);
case AgentProtocol.MONGO_METHOD_DROP_DATABASE -> dropDatabase(params);
case AgentProtocol.MONGO_METHOD_INSERT_DOCUMENT -> insertDocument(params);
case AgentProtocol.MONGO_METHOD_INSERT_DOCUMENTS -> insertDocuments(params);
case AgentProtocol.MONGO_METHOD_UPDATE_DOCUMENT -> updateDocument(params);
case AgentProtocol.MONGO_METHOD_UPDATE_DOCUMENTS -> updateDocuments(params);
case AgentProtocol.MONGO_METHOD_DELETE_DOCUMENT -> deleteDocument(params);
agents/drivers/mongodb/src/test/java/com/dbx/agent/mongodb/MongoAgentTest.java
+122
View File
@@ -127,6 +127,7 @@ class MongoAgentTest {
assertTrue(containsCapability(result.getAsJsonArray("capabilities"), AgentProtocol.CAPABILITY_MONGO_DROP_DATABASE));
assertTrue(containsCapability(result.getAsJsonArray("capabilities"), AgentProtocol.CAPABILITY_MONGO_CLONE_COLLECTION));
assertTrue(containsCapability(result.getAsJsonArray("capabilities"), AgentProtocol.CAPABILITY_MONGO_RUN_COMMAND));
assertTrue(containsCapability(result.getAsJsonArray("capabilities"), AgentProtocol.CAPABILITY_MONGO_INSERT_DOCUMENTS));
}
@Test
@@ -148,6 +149,7 @@ class MongoAgentTest {
assertTrue(containsCapability(result.getAsJsonArray("capabilities"), AgentProtocol.CAPABILITY_MONGO_DROP_DATABASE));
assertTrue(containsCapability(result.getAsJsonArray("capabilities"), AgentProtocol.CAPABILITY_MONGO_CLONE_COLLECTION));
assertTrue(containsCapability(result.getAsJsonArray("capabilities"), AgentProtocol.CAPABILITY_MONGO_RUN_COMMAND));
assertTrue(containsCapability(result.getAsJsonArray("capabilities"), AgentProtocol.CAPABILITY_MONGO_INSERT_DOCUMENTS));
}
@Test
@@ -874,6 +876,69 @@ class MongoAgentTest {
assertFalse(json.getAsJsonObject("error").get("message").getAsString().contains("Unknown method"));
}
@Test
void insertDocumentsRpcInsertsTheReportedArrayInOneRequest() {
List<List<Document>> batches = new ArrayList<>();
List<String> calls = new ArrayList<>();
MongoClient client = recordingInsertMongoClient(calls, batches);
String documents = "[{\"type\":999,\"refid\":\"11\",\"externalId\":\"67ab71a2a5d1c681530dc61c\",\"__v\":0},"
+ "{\"type\":999,\"refid\":\"12\",\"externalId\":\"67fd678f595634e3a3094f5b\",\"__v\":0}]";
JsonObject response = insertDocumentsRpc(client, 40, "user", documents);
assertFalse(response.has("error"), response.toString());
assertEquals(2, response.getAsJsonObject("result").get("affected_rows").getAsInt());
assertEquals(List.of("getDatabase:app", "getCollection:user", "insertMany:2"), calls);
assertEquals(1, batches.size());
assertEquals("11", batches.get(0).get(0).getString("refid"));
assertEquals("12", batches.get(0).get(1).getString("refid"));
assertTrue(AgentProtocol.MONGO_LEGACY_METHODS.contains(AgentProtocol.MONGO_METHOD_INSERT_DOCUMENTS));
assertTrue(AgentProtocol.MONGO_LEGACY_CAPABILITIES.contains(AgentProtocol.CAPABILITY_MONGO_INSERT_DOCUMENTS));
}
@Test
void insertDocumentsRpcUsesInsertManyForOneDocument() {
List<List<Document>> batches = new ArrayList<>();
List<String> calls = new ArrayList<>();
MongoClient client = recordingInsertMongoClient(calls, batches);
JsonObject response = insertDocumentsRpc(client, 41, "users", "[{\"name\":\"Ada\"}]");
assertFalse(response.has("error"), response.toString());
assertEquals(1, response.getAsJsonObject("result").get("affected_rows").getAsInt());
assertEquals(List.of("getDatabase:app", "getCollection:users", "insertMany:1"), calls);
assertEquals(new Document("name", "Ada"), batches.get(0).get(0));
}
@Test
void insertDocumentsRpcTreatsAnEmptyArrayAsANoOp() {
List<List<Document>> batches = new ArrayList<>();
List<String> calls = new ArrayList<>();
MongoClient client = recordingInsertMongoClient(calls, batches);
JsonObject response = insertDocumentsRpc(client, 42, "users", "[]");
assertFalse(response.has("error"), response.toString());
assertEquals(0, response.getAsJsonObject("result").get("affected_rows").getAsInt());
assertTrue(calls.isEmpty());
assertTrue(batches.isEmpty());
}
@Test
void insertDocumentsRpcRejectsInvalidBatchesBeforeDatabaseAccess() {
for (String documents : List.of("{", "{\"name\":\"Ada\"}", "[1]", "[{\"name\":\"Ada\"},null]")) {
List<List<Document>> batches = new ArrayList<>();
List<String> calls = new ArrayList<>();
MongoClient client = recordingInsertMongoClient(calls, batches);
JsonObject response = insertDocumentsRpc(client, 43, "users", documents);
assertTrue(response.has("error"), documents + ": " + response);
assertTrue(calls.isEmpty(), documents + ": " + calls);
assertTrue(batches.isEmpty(), documents + ": " + batches);
}
}
@Test
void updateDocumentsRpcUsesDocumentOverloads() {
List<String> calls = new ArrayList<>();
@@ -1307,6 +1372,63 @@ class MongoAgentTest {
assertEquals(1, response.getAsJsonObject("result").get("modified_count").getAsLong());
}
private static JsonObject insertDocumentsRpc(
MongoClient client, int id, String collection, String documents
) {
JsonObject params = new JsonObject();
params.addProperty("database", "app");
params.addProperty("collection", collection);
params.addProperty("docs_json", documents);
JsonObject request = new JsonObject();
request.addProperty("jsonrpc", "2.0");
request.addProperty("id", id);
request.addProperty("method", "insert_documents");
request.add("params", params);
return JsonParser.parseString(MongoAgent.handleRequest(request.toString(), client)).getAsJsonObject();
}
@SuppressWarnings("unchecked")
private static MongoClient recordingInsertMongoClient(
List<String> calls, List<List<Document>> batches
) {
MongoCollection<Document> collection = (MongoCollection<Document>) Proxy.newProxyInstance(
MongoCollection.class.getClassLoader(),
new Class<?>[] {MongoCollection.class},
(proxy, method, args) -> {
if ("insertMany".equals(method.getName())) {
List<Document> batch = ((List<Document>) args[0]).stream().map(Document::new).toList();
batches.add(batch);
calls.add("insertMany:" + batch.size());
return 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())) {
calls.add("getCollection:" + args[0]);
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())) {
calls.add("getDatabase:" + args[0]);
return database;
}
throw new UnsupportedOperationException(method.getName());
}
);
}
@SuppressWarnings("unchecked")
private static MongoClient recordingMongoClient(List<String> calls) {
MongoCollection<Document> collection = (MongoCollection<Document>) Proxy.newProxyInstance(
crates/dbx-core/assets/agent-protocol-v1.json
+2 -2
View File
@@ -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"],
"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"],
"capabilities": ["connect", "test_connection", "metadata", "query", "paged_query", "transaction", "ddl"],
"defaultSqlCapabilities": ["connect", "test_connection", "metadata", "query", "paged_query", "transaction", "ddl"],
"commonMethods": [
@@ -39,6 +39,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", "update_document", "update_documents", "delete_document", "delete_documents", "run_command"],
"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", "delete_document", "delete_documents", "run_command"],
"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"]
}
crates/dbx-core/assets/agent-protocol-v2.json
+2 -2
View File
@@ -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", "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", "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": [
@@ -43,7 +43,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", "update_document", "update_documents", "delete_document", "delete_documents", "run_command"],
"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", "delete_document", "delete_documents", "run_command"],
"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"
crates/dbx-core/src/db/agent_driver.rs
+30 -3
View File
@@ -974,6 +974,7 @@ pub enum AgentCapability {
MongoDropDatabase,
MongoCloneCollection,
MongoRunCommand,
MongoInsertDocuments,
MultiSession,
StructuredErrorV1,
}
@@ -1056,7 +1057,7 @@ fn parse_agent_rpc_error_header(header: &str) -> (Option<i64>, String) {
}
impl AgentCapability {
pub const ALL: [Self; 23] = [
pub const ALL: [Self; 24] = [
Self::Connect,
Self::TestConnection,
Self::Metadata,
@@ -1078,6 +1079,7 @@ impl AgentCapability {
Self::MongoDropDatabase,
Self::MongoCloneCollection,
Self::MongoRunCommand,
Self::MongoInsertDocuments,
Self::MultiSession,
Self::StructuredErrorV1,
];
@@ -1105,6 +1107,7 @@ impl AgentCapability {
Self::MongoDropDatabase => "mongo_drop_database",
Self::MongoCloneCollection => "mongo_clone_collection",
Self::MongoRunCommand => "mongo_run_command",
Self::MongoInsertDocuments => "mongo_insert_documents",
Self::MultiSession => "multi_session",
Self::StructuredErrorV1 => "structured_error_v1",
}
@@ -1293,6 +1296,7 @@ pub enum MongoAgentMethod {
CloneCollection,
DropDatabase,
InsertDocument,
InsertDocuments,
UpdateDocument,
UpdateDocuments,
DeleteDocument,
@@ -1301,7 +1305,7 @@ pub enum MongoAgentMethod {
}
impl MongoAgentMethod {
pub const ALL: [Self; 21] = [
pub const ALL: [Self; 22] = [
Self::ListDatabases,
Self::ListCollections,
Self::FindDocuments,
@@ -1318,6 +1322,7 @@ impl MongoAgentMethod {
Self::CloneCollection,
Self::DropDatabase,
Self::InsertDocument,
Self::InsertDocuments,
Self::UpdateDocument,
Self::UpdateDocuments,
Self::DeleteDocument,
@@ -1343,6 +1348,7 @@ impl MongoAgentMethod {
Self::CloneCollection => "clone_collection",
Self::DropDatabase => "drop_database",
Self::InsertDocument => "insert_document",
Self::InsertDocuments => "insert_documents",
Self::UpdateDocument => "update_document",
Self::UpdateDocuments => "update_documents",
Self::DeleteDocument => "delete_document",
@@ -2601,6 +2607,13 @@ impl AgentDriverClient {
self.call_mongo_method(MongoAgentMethod::InsertDocument, params).await
}
pub async fn mongo_insert_documents<T: DeserializeOwned + Send + 'static>(
&mut self,
params: Value,
) -> Result<T, String> {
self.call_mongo_method(MongoAgentMethod::InsertDocuments, params).await
}
pub async fn mongo_update_document<T: DeserializeOwned + Send + 'static>(
&mut self,
params: Value,
@@ -2769,6 +2782,7 @@ pub fn agent_supports_capability(handshake: Option<&AgentHandshake>, capability:
| AgentCapability::MongoDropDatabase
| AgentCapability::MongoCloneCollection
| AgentCapability::MongoRunCommand
| AgentCapability::MongoInsertDocuments
) {
return handshake.map(|value| value.supports(capability)).unwrap_or(false);
}
@@ -4358,9 +4372,10 @@ for line in sys.stdin:
assert_eq!(AgentCapability::MongoDropDatabase.as_str(), "mongo_drop_database");
assert_eq!(AgentCapability::MongoCloneCollection.as_str(), "mongo_clone_collection");
assert_eq!(AgentCapability::MongoRunCommand.as_str(), "mongo_run_command");
assert_eq!(AgentCapability::MongoInsertDocuments.as_str(), "mongo_insert_documents");
assert_eq!(AgentCapability::MultiSession.as_str(), "multi_session");
assert_eq!(AgentCapability::StructuredErrorV1.as_str(), "structured_error_v1");
assert_eq!(AgentCapability::ALL.len(), 23);
assert_eq!(AgentCapability::ALL.len(), 24);
}
#[test]
@@ -4415,6 +4430,7 @@ for line in sys.stdin:
assert_eq!(MongoAgentMethod::CloneCollection.as_str(), "clone_collection");
assert_eq!(MongoAgentMethod::DropDatabase.as_str(), "drop_database");
assert_eq!(MongoAgentMethod::InsertDocument.as_str(), "insert_document");
assert_eq!(MongoAgentMethod::InsertDocuments.as_str(), "insert_documents");
assert_eq!(MongoAgentMethod::UpdateDocument.as_str(), "update_document");
assert_eq!(MongoAgentMethod::UpdateDocuments.as_str(), "update_documents");
assert_eq!(MongoAgentMethod::DeleteDocument.as_str(), "delete_document");
@@ -4719,6 +4735,8 @@ for line in sys.stdin:
assert!(!agent_supports_capability(Some(&handshake), AgentCapability::MongoCloneCollection));
assert!(!agent_supports_capability(None, AgentCapability::MongoRunCommand));
assert!(!agent_supports_capability(Some(&handshake), AgentCapability::MongoRunCommand));
assert!(!agent_supports_capability(None, AgentCapability::MongoInsertDocuments));
assert!(!agent_supports_capability(Some(&handshake), AgentCapability::MongoInsertDocuments));
let mongo_handshake =
AgentHandshake { capabilities: vec![AgentCapability::MongoDropDatabase.as_str().to_string()], ..handshake };
@@ -4735,6 +4753,15 @@ for line in sys.stdin:
..mongo_clone_handshake
};
assert!(agent_supports_capability(Some(&mongo_run_command_handshake), AgentCapability::MongoRunCommand));
let mongo_insert_documents_handshake = AgentHandshake {
capabilities: vec![AgentCapability::MongoInsertDocuments.as_str().to_string()],
..mongo_run_command_handshake
};
assert!(agent_supports_capability(
Some(&mongo_insert_documents_handshake),
AgentCapability::MongoInsertDocuments
));
}
#[test]
crates/dbx-core/src/mongo_ops.rs
+123 -1
View File
@@ -657,7 +657,38 @@ pub async fn mongo_insert_documents_core(
let connections = state.connections.read().await;
match connections.get(connection_id).ok_or("Not found")? {
PoolKind::MongoDb(client) => mongo_driver::insert_documents(client, database, collection, docs_json).await,
PoolKind::Agent(_) => Err("MongoDB legacy agent does not support bulk insertMany/insertOne writes".to_string()),
PoolKind::Agent(client) => {
let documents: serde_json::Value =
serde_json::from_str(docs_json).map_err(|error| format!("Invalid JSON: {error}"))?;
let documents = documents.as_array().ok_or_else(|| {
"MongoDB legacy agent does not support bulk insertMany/insertOne writes; insertMany requires an array"
.to_string()
})?;
if documents.iter().any(|document| !document.is_object()) {
return Err("Each MongoDB insertMany document must be an object".to_string());
}
if documents.is_empty() {
return Ok(0);
}
let mut client = client.lock().await;
if !client.supports_capability(AgentCapability::MongoInsertDocuments) {
return Err(
"MongoDB Legacy Agent does not support insertMany; upgrade or reinstall the MongoDB Legacy driver"
.to_string(),
);
}
let result: serde_json::Value = client
.mongo_insert_documents(serde_json::json!({
"database": database,
"collection": collection,
"docs_json": docs_json,
}))
.await?;
result
.get("affected_rows")
.and_then(serde_json::Value::as_u64)
.ok_or_else(|| "MongoDB Legacy Agent returned an invalid insertMany result".to_string())
}
_ => Err("Not a MongoDB connection".to_string()),
}
}
@@ -1429,6 +1460,97 @@ for line in sys.stdin:
assert!(!error.contains("Unknown method"), "{error}");
}
#[cfg(unix)]
#[tokio::test]
async fn mongo_insert_many_routes_legacy_connections_to_one_agent_call() {
let documents = r#"[{"type":999,"refid":"11"},{"type":999,"refid":"12"}]"#;
let (state, _directory) = legacy_mongo_state_with_capabilities(
"insert_documents",
serde_json::json!({
"database": "app",
"collection": "user",
"docs_json": documents,
}),
serde_json::json!({ "affected_rows": 2 }),
&["mongo_insert_documents"],
)
.await;
let affected = mongo_insert_documents_core(&state, "legacy", "app", "user", documents).await.unwrap();
assert_eq!(affected, 2);
}
#[cfg(unix)]
#[tokio::test]
async fn mongo_insert_many_requires_an_explicit_legacy_agent_capability() {
let documents = r#"[{"name":"Ada"}]"#;
let (state, _directory) = legacy_mongo_state_with_capabilities(
"insert_documents",
serde_json::json!({
"database": "app",
"collection": "users",
"docs_json": documents,
}),
serde_json::json!({ "affected_rows": 1 }),
&[],
)
.await;
let error = mongo_insert_documents_core(&state, "legacy", "app", "users", documents).await.unwrap_err();
assert!(error.contains("upgrade or reinstall"), "{error}");
assert!(!error.contains("unexpected MongoDB RPC"), "{error}");
assert!(!error.contains("Unknown method"), "{error}");
}
#[cfg(unix)]
#[tokio::test]
async fn mongo_insert_many_rejects_invalid_batches_and_skips_empty_batches_before_dispatch() {
let (state, _directory) = legacy_mongo_state_with_capabilities(
"insert_documents",
serde_json::json!({
"database": "app",
"collection": "users",
"docs_json": [],
}),
serde_json::json!({ "affected_rows": 1 }),
&[AgentCapability::MongoInsertDocuments.as_str()],
)
.await;
let affected = mongo_insert_documents_core(&state, "legacy", "app", "users", "[]").await.unwrap();
let error = mongo_insert_documents_core(&state, "legacy", "app", "users", r#"[{"name":"Ada"},null]"#)
.await
.unwrap_err();
assert_eq!(affected, 0);
assert!(error.contains("must be an object"), "{error}");
assert!(!error.contains("unexpected MongoDB RPC"), "{error}");
}
#[cfg(unix)]
#[tokio::test]
async fn mongo_insert_one_keeps_the_existing_legacy_agent_behavior() {
let document = r#"{"name":"Ada"}"#;
let (state, _directory) = legacy_mongo_state_with_capabilities(
"insert_documents",
serde_json::json!({
"database": "app",
"collection": "users",
"docs_json": document,
}),
serde_json::json!({ "affected_rows": 1 }),
&[AgentCapability::MongoInsertDocuments.as_str()],
)
.await;
let error = mongo_insert_documents_core(&state, "legacy", "app", "users", document).await.unwrap_err();
assert!(error.contains("insertMany/insertOne"), "{error}");
assert!(!error.contains("unexpected MongoDB RPC"), "{error}");
}
#[cfg(unix)]
#[tokio::test]
async fn mongo_drop_collection_routes_legacy_connections_to_the_agent() {