fix(mongodb): use dedicated count for legacy collections

src-tauri/src/commands/document_cmd.rs

@@ -6,7 +6,11 @@ use crate::commands::connection::{ensure_connection_writable, AppState};
 use dbx_core::db::mongo_driver::MongoDocumentResult;
 use dbx_core::document_ops::CollectionInfo;
 
-async fn run_cancellable<T, F>(state: &Arc<AppState>, execution_id: Option<String>, future: F) -> Result<T, String>
+pub(crate) async fn run_cancellable<T, F>(
+    state: &Arc<AppState>,
+    execution_id: Option<String>,
+    future: F,
+) -> Result<T, String>
 where
     F: Future<Output = Result<T, String>>,
 {

src-tauri/src/commands/mcp_bridge.rs

@@ -62,6 +62,16 @@ struct MongoFindDocumentsRequest {
     sort: Option<String>,
 }
 
+#[derive(Deserialize)]
+struct MongoCountDocumentsRequest {
+    connection_name: String,
+    connection_id: Option<String>,
+    database: Option<String>,
+    collection: String,
+    filter: Option<String>,
+    mode: Option<String>,
+}
+
 #[derive(Deserialize)]
 struct MongoServerVersionRequest {
     connection_name: String,
@@ -214,6 +224,8 @@ pub fn start(app_handle: AppHandle, state: Arc<AppState>, data_dir: PathBuf) {
                     handle_describe_table_data(&st, body, &mut stream).await;
                 } else if first_line.starts_with("POST /data/mongo/list-collections") {
                     handle_mongo_list_collections_data(&st, body, &mut stream).await;
+                } else if first_line.starts_with("POST /data/mongo/count-documents") {
+                    handle_mongo_count_documents_data(&st, body, &mut stream).await;
                 } else if first_line.starts_with("POST /data/mongo/find-documents") {
                     handle_mongo_find_documents_data(&st, body, &mut stream).await;
                 } else if first_line.starts_with("POST /data/mongo/server-version") {
@@ -518,6 +530,34 @@ async fn handle_mongo_find_documents_data(state: &Arc<AppState>, body: &str, str
     }
 }
 
+async fn handle_mongo_count_documents_data(state: &Arc<AppState>, body: &str, stream: &mut tokio::net::TcpStream) {
+    let req: MongoCountDocumentsRequest = match serde_json::from_str(body) {
+        Ok(r) => r,
+        Err(_) => {
+            respond_error(stream, "400 Bad Request", "Invalid JSON").await;
+            return;
+        }
+    };
+    let Some((pool_key, database, _connection_id)) =
+        resolve_mongo_pool_key(state, req.connection_id.as_deref(), &req.connection_name, req.database, stream).await
+    else {
+        return;
+    };
+    match dbx_core::mongo_ops::mongo_count_documents_core(
+        state,
+        &pool_key,
+        &database,
+        &req.collection,
+        req.filter.as_deref(),
+        req.mode.as_deref(),
+    )
+    .await
+    {
+        Ok(total) => respond_json(stream, &total).await,
+        Err(e) => respond_error(stream, "500 Internal Server Error", &e).await,
+    }
+}
+
 async fn handle_mongo_server_version_data(state: &Arc<AppState>, body: &str, stream: &mut tokio::net::TcpStream) {
     let req: MongoServerVersionRequest = match serde_json::from_str(body) {
         Ok(r) => r,

src-tauri/src/commands/mongo_cmd.rs

@@ -110,6 +110,32 @@ pub async fn mongo_find_documents(
     .await
 }
 
+#[tauri::command]
+pub async fn mongo_count_documents(
+    state: State<'_, Arc<AppState>>,
+    connection_id: String,
+    database: String,
+    collection: String,
+    filter: Option<String>,
+    mode: Option<String>,
+    execution_id: Option<String>,
+) -> Result<u64, String> {
+    let app = state.inner().clone();
+    crate::commands::document_cmd::run_cancellable(
+        &app,
+        execution_id,
+        dbx_core::mongo_ops::mongo_count_documents_core(
+            &app,
+            &connection_id,
+            &database,
+            &collection,
+            filter.as_deref(),
+            mode.as_deref(),
+        ),
+    )
+    .await
+}
+
 #[tauri::command]
 pub async fn mongo_server_version(
     state: State<'_, Arc<AppState>>,