feat(mongodb): stream legacy agent exports via find cursors

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", "mongo_insert_documents", "mongo_replace_document"],
+  "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", "mongo_replace_document", "mongo_find_cursor"],
   "capabilities": ["connect", "test_connection", "metadata", "query", "paged_query", "transaction", "ddl"],
   "defaultSqlCapabilities": ["connect", "test_connection", "metadata", "query", "paged_query", "transaction", "ddl"],
   "commonMethods": [
@@ -42,6 +42,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", "insert_documents", "update_document", "update_documents", "replace_document", "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", "replace_document", "delete_document", "delete_documents", "run_command", "start_find_cursor", "fetch_find_cursor", "close_find_cursor"],
   "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", "mongo_insert_documents", "mongo_replace_document", "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", "mongo_replace_document", "mongo_find_cursor", "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": [
@@ -46,7 +46,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", "insert_documents", "update_document", "update_documents", "replace_document", "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", "replace_document", "delete_document", "delete_documents", "run_command", "start_find_cursor", "fetch_find_cursor", "close_find_cursor"],
   "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

@@ -1042,6 +1042,7 @@ pub enum AgentCapability {
     MongoRunCommand,
     MongoInsertDocuments,
     MongoReplaceDocument,
+    MongoFindCursor,
     MultiSession,
     StructuredErrorV1,
 }
@@ -1124,7 +1125,7 @@ fn parse_agent_rpc_error_header(header: &str) -> (Option<i64>, String) {
 }
 
 impl AgentCapability {
-    pub const ALL: [Self; 25] = [
+    pub const ALL: [Self; 26] = [
         Self::Connect,
         Self::TestConnection,
         Self::Metadata,
@@ -1148,6 +1149,7 @@ impl AgentCapability {
         Self::MongoRunCommand,
         Self::MongoInsertDocuments,
         Self::MongoReplaceDocument,
+        Self::MongoFindCursor,
         Self::MultiSession,
         Self::StructuredErrorV1,
     ];
@@ -1177,6 +1179,7 @@ impl AgentCapability {
             Self::MongoRunCommand => "mongo_run_command",
             Self::MongoInsertDocuments => "mongo_insert_documents",
             Self::MongoReplaceDocument => "mongo_replace_document",
+            Self::MongoFindCursor => "mongo_find_cursor",
             Self::MultiSession => "multi_session",
             Self::StructuredErrorV1 => "structured_error_v1",
         }
@@ -1383,10 +1386,13 @@ pub enum MongoAgentMethod {
     DeleteDocument,
     DeleteDocuments,
     RunCommand,
+    StartFindCursor,
+    FetchFindCursor,
+    CloseFindCursor,
 }
 
 impl MongoAgentMethod {
-    pub const ALL: [Self; 23] = [
+    pub const ALL: [Self; 26] = [
         Self::ListDatabases,
         Self::ListCollections,
         Self::FindDocuments,
@@ -1410,6 +1416,9 @@ impl MongoAgentMethod {
         Self::DeleteDocument,
         Self::DeleteDocuments,
         Self::RunCommand,
+        Self::StartFindCursor,
+        Self::FetchFindCursor,
+        Self::CloseFindCursor,
     ];
 
     pub fn as_str(self) -> &'static str {
@@ -1437,6 +1446,9 @@ impl MongoAgentMethod {
             Self::DeleteDocument => "delete_document",
             Self::DeleteDocuments => "delete_documents",
             Self::RunCommand => "run_command",
+            Self::StartFindCursor => "start_find_cursor",
+            Self::FetchFindCursor => "fetch_find_cursor",
+            Self::CloseFindCursor => "close_find_cursor",
         }
     }
 }
@@ -2795,6 +2807,27 @@ impl AgentDriverClient {
         self.call_mongo_method(MongoAgentMethod::InsertDocuments, params).await
     }
 
+    pub async fn mongo_start_find_cursor<T: DeserializeOwned + Send + 'static>(
+        &mut self,
+        params: Value,
+    ) -> Result<T, String> {
+        self.call_mongo_method(MongoAgentMethod::StartFindCursor, params).await
+    }
+
+    pub async fn mongo_fetch_find_cursor<T: DeserializeOwned + Send + 'static>(
+        &mut self,
+        params: Value,
+    ) -> Result<T, String> {
+        self.call_mongo_method(MongoAgentMethod::FetchFindCursor, params).await
+    }
+
+    pub async fn mongo_close_find_cursor<T: DeserializeOwned + Send + 'static>(
+        &mut self,
+        params: Value,
+    ) -> Result<T, String> {
+        self.call_mongo_method(MongoAgentMethod::CloseFindCursor, params).await
+    }
+
     pub async fn mongo_update_document<T: DeserializeOwned + Send + 'static>(
         &mut self,
         params: Value,
@@ -2972,6 +3005,7 @@ pub fn agent_supports_capability(handshake: Option<&AgentHandshake>, capability:
             | AgentCapability::MongoRunCommand
             | AgentCapability::MongoInsertDocuments
             | AgentCapability::MongoReplaceDocument
+            | AgentCapability::MongoFindCursor
     ) {
         return handshake.map(|value| value.supports(capability)).unwrap_or(false);
     }
@@ -5006,9 +5040,10 @@ for line in sys.stdin:
         assert_eq!(AgentCapability::MongoRunCommand.as_str(), "mongo_run_command");
         assert_eq!(AgentCapability::MongoInsertDocuments.as_str(), "mongo_insert_documents");
         assert_eq!(AgentCapability::MongoReplaceDocument.as_str(), "mongo_replace_document");
+        assert_eq!(AgentCapability::MongoFindCursor.as_str(), "mongo_find_cursor");
         assert_eq!(AgentCapability::MultiSession.as_str(), "multi_session");
         assert_eq!(AgentCapability::StructuredErrorV1.as_str(), "structured_error_v1");
-        assert_eq!(AgentCapability::ALL.len(), 25);
+        assert_eq!(AgentCapability::ALL.len(), 26);
     }
 
     #[test]
@@ -5073,6 +5108,9 @@ for line in sys.stdin:
         assert_eq!(MongoAgentMethod::ReplaceDocument.as_str(), "replace_document");
         assert_eq!(MongoAgentMethod::DeleteDocument.as_str(), "delete_document");
         assert_eq!(MongoAgentMethod::DeleteDocuments.as_str(), "delete_documents");
+        assert_eq!(MongoAgentMethod::StartFindCursor.as_str(), "start_find_cursor");
+        assert_eq!(MongoAgentMethod::FetchFindCursor.as_str(), "fetch_find_cursor");
+        assert_eq!(MongoAgentMethod::CloseFindCursor.as_str(), "close_find_cursor");
     }
 
     #[test]
@@ -5138,6 +5176,9 @@ for line in sys.stdin:
         let _mongo_clone_collection = AgentDriverClient::mongo_clone_collection::<serde_json::Value>;
         let _mongo_drop_database = AgentDriverClient::mongo_drop_database::<serde_json::Value>;
         let _mongo_insert_document = AgentDriverClient::mongo_insert_document::<serde_json::Value>;
+        let _mongo_start_find_cursor = AgentDriverClient::mongo_start_find_cursor::<serde_json::Value>;
+        let _mongo_fetch_find_cursor = AgentDriverClient::mongo_fetch_find_cursor::<serde_json::Value>;
+        let _mongo_close_find_cursor = AgentDriverClient::mongo_close_find_cursor::<serde_json::Value>;
         let _mongo_update_document = AgentDriverClient::mongo_update_document::<serde_json::Value>;
         let _mongo_update_documents = AgentDriverClient::mongo_update_documents::<serde_json::Value>;
         let _mongo_delete_document = AgentDriverClient::mongo_delete_document::<serde_json::Value>;
@@ -5389,6 +5430,8 @@ for line in sys.stdin:
         assert!(!agent_supports_capability(Some(&handshake), AgentCapability::MongoRunCommand));
         assert!(!agent_supports_capability(None, AgentCapability::MongoInsertDocuments));
         assert!(!agent_supports_capability(Some(&handshake), AgentCapability::MongoInsertDocuments));
+        assert!(!agent_supports_capability(None, AgentCapability::MongoFindCursor));
+        assert!(!agent_supports_capability(Some(&handshake), AgentCapability::MongoFindCursor));
 
         let mongo_handshake =
             AgentHandshake { capabilities: vec![AgentCapability::MongoDropDatabase.as_str().to_string()], ..handshake };
@@ -5414,6 +5457,12 @@ for line in sys.stdin:
             Some(&mongo_insert_documents_handshake),
             AgentCapability::MongoInsertDocuments
         ));
+
+        let mongo_find_cursor_handshake = AgentHandshake {
+            capabilities: vec![AgentCapability::MongoFindCursor.as_str().to_string()],
+            ..mongo_insert_documents_handshake
+        };
+        assert!(agent_supports_capability(Some(&mongo_find_cursor_handshake), AgentCapability::MongoFindCursor));
     }
 
     #[test]

crates/dbx-core/src/db/mongo_driver.rs

@@ -2649,6 +2649,41 @@ fn is_retryable_mongo_write_code(code: i32) -> bool {
     !matches!(code, 11000 | 11001 | 12582)
 }
 
+/// Reads the `insert_documents` result of a MongoDB legacy agent. The agent reports a partly
+/// applied batch as a success carrying one entry per rejected document, so callers must inspect
+/// [`MongoInsertOutcome::errors`] instead of trusting a bare affected-row count.
+pub fn agent_insert_outcome(result: &serde_json::Value) -> Result<MongoInsertOutcome, String> {
+    let inserted = result
+        .get("affected_rows")
+        .and_then(serde_json::Value::as_u64)
+        .ok_or_else(|| "MongoDB Legacy Agent returned an invalid insertMany result".to_string())?;
+    let errors = result
+        .get("errors")
+        .and_then(serde_json::Value::as_array)
+        .map(|errors| errors.iter().map(agent_insert_error).collect())
+        .unwrap_or_default();
+    Ok(MongoInsertOutcome { inserted, errors })
+}
+
+/// One rejected document, or a batch-wide rejection the agent could not attribute to a document
+/// (a write concern failure) — those carry no index and keep [`MongoBulkWriteError::index`] `None`.
+fn agent_insert_error(value: &serde_json::Value) -> MongoBulkWriteError {
+    // A rejection is never dropped for lack of a message: dropping it would make a document the
+    // server refused look like a successful insert.
+    let message = value
+        .get("message")
+        .and_then(serde_json::Value::as_str)
+        .unwrap_or("MongoDB Legacy Agent rejected a document without a message")
+        .to_string();
+    let code = value.get("code").and_then(serde_json::Value::as_i64).map(|code| code as i32);
+    MongoBulkWriteError {
+        message,
+        index: value.get("index").and_then(serde_json::Value::as_u64).map(|index| index as usize),
+        code,
+        retryable: code.is_some_and(is_retryable_mongo_write_code),
+    }
+}
+
 #[allow(clippy::too_many_arguments)]
 pub async fn for_each_find_document(
     client: &Client,

crates/dbx-core/src/mongo_ops.rs

@@ -439,11 +439,18 @@ pub async fn mongo_find_documents_extended_json_core(
     }
 }
 
-fn is_unknown_agent_method_error(error: &str, method: &str) -> bool {
+pub(crate) fn is_unknown_agent_method_error(error: &str, method: &str) -> bool {
     let lower = error.to_ascii_lowercase();
     lower.contains(method) && (lower.contains("unknown method") || lower.contains("method not found"))
 }
 
+/// True only when a legacy agent refused a bulk write because it has no `insertMany` at all, which
+/// is the one failure a document-by-document retry can still get through. It must stay precise:
+/// every other legacy-agent failure, a rejected document included, has to surface as-is.
+pub(crate) fn is_legacy_agent_insert_many_unsupported(error: &str) -> bool {
+    error.contains("does not support insertMany") || error.contains("does not support bulk insertMany")
+}
+
 pub async fn mongo_aggregate_documents_core(
     state: &AppState,
     connection_id: &str,
@@ -730,10 +737,18 @@ pub async fn mongo_insert_documents_core(
                     "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())
+            let outcome = mongo_driver::agent_insert_outcome(&result)?;
+            // The agent reports a partly applied batch as a success, so a rejected document has to
+            // fail the call here exactly like the native driver's insert_many does.
+            if let Some(first) = outcome.errors.first() {
+                return Err(format!(
+                    "MongoDB Legacy Agent rejected {} of {} documents: {}",
+                    outcome.errors.len(),
+                    documents.len(),
+                    first.message
+                ));
+            }
+            Ok(outcome.inserted)
         }
         _ => Err("Not a MongoDB connection".to_string()),
     }
@@ -1758,6 +1773,46 @@ for line in sys.stdin:
         assert_eq!(affected, 2);
     }
 
+    #[cfg(unix)]
+    #[tokio::test]
+    async fn mongo_insert_many_fails_when_the_legacy_agent_rejects_a_document() {
+        let documents = r#"[{"_id":1},{"_id":1}]"#;
+        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,
+                "errors": [{ "index": 1, "code": 11000, "message": "E11000 duplicate key" }],
+            }),
+            &["mongo_insert_documents"],
+        )
+        .await;
+
+        let error = mongo_insert_documents_core(&state, "legacy", "app", "users", documents).await.unwrap_err();
+
+        // The agent reports a partly applied batch as a success, so a rejected document has to
+        // fail the call here rather than being rounded up to "inserted".
+        assert_eq!(error, "MongoDB Legacy Agent rejected 1 of 2 documents: E11000 duplicate key");
+    }
+
+    #[test]
+    fn legacy_insert_many_heuristic_covers_only_the_missing_capability() {
+        assert!(is_legacy_agent_insert_many_unsupported(
+            "MongoDB Legacy Agent does not support insertMany; upgrade or reinstall the MongoDB Legacy driver"
+        ));
+        assert!(is_legacy_agent_insert_many_unsupported(
+            "MongoDB legacy agent does not support bulk insertMany/insertOne writes"
+        ));
+        // A rejected document must not be retried one by one: the rest of the batch is written.
+        assert!(!is_legacy_agent_insert_many_unsupported(
+            "MongoDB Legacy Agent rejected 1 of 2 documents: E11000 duplicate key"
+        ));
+    }
+
     #[cfg(unix)]
     #[tokio::test]
     async fn mongo_insert_many_requires_an_explicit_legacy_agent_capability() {

crates/dbx-core/src/transfer.rs

@@ -5421,7 +5421,7 @@ async fn insert_mongo_documents_for_transfer(
     let docs_json = serde_json::to_string(documents).map_err(|e| format!("Failed to encode MongoDB documents: {e}"))?;
     match crate::mongo_ops::mongo_insert_documents_core(state, connection_id, database, collection, &docs_json).await {
         Ok(count) => Ok(count),
-        Err(error) if error.to_ascii_lowercase().contains("legacy agent") => {
+        Err(error) if crate::mongo_ops::is_legacy_agent_insert_many_unsupported(&error) => {
             let mut inserted = 0;
             for document in documents {
                 let doc_json =
@@ -5457,7 +5457,7 @@ async fn insert_mongo_documents_extended_json_for_transfer(
     .await
     {
         Ok(count) => Ok(count),
-        Err(error) if error.to_ascii_lowercase().contains("legacy agent") => {
+        Err(error) if crate::mongo_ops::is_legacy_agent_insert_many_unsupported(&error) => {
             insert_mongo_documents_for_transfer(state, connection_id, database, collection, documents).await
         }
         Err(error) => Err(error),