fix(mongodb): support legacy agent insertMany writes

Refs: #6032

crates/dbx-core/assets/agent-protocol-v1.json

@@ -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,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

@@ -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

@@ -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() {