From def42987a4c6922a7e77624fb7df2fb9b8405110 Mon Sep 17 00:00:00 2001 From: Loc Nguyen <113083148+thailoc-dev@users.noreply.github.com> Date: Tue, 22 Sep 2026 13:22:15 +0700 Subject: [PATCH] feat(mongodb): fill object statistics over the legacy agent --- crates/dbx-core/src/query/document_ops.rs | 2 +- crates/dbx-core/src/query/mongo_ops.rs | 32 +++++++++++++++++++ crates/dbx-core/src/schema/mod.rs | 3 ++ .../live_mongodb_legacy_agent_commands.rs | 15 +++++++++ crates/dbx-drivers/src/db/mongo_driver.rs | 2 +- 5 files changed, 52 insertions(+), 2 deletions(-) diff --git a/crates/dbx-core/src/query/document_ops.rs b/crates/dbx-core/src/query/document_ops.rs index 81050a02b7..b9d63be881 100644 --- a/crates/dbx-core/src/query/document_ops.rs +++ b/crates/dbx-core/src/query/document_ops.rs @@ -107,7 +107,7 @@ fn sort_mongo_collection_specs( /// Decode both generations of the Legacy Agent response. Existing installed /// Agents return names, while current Agents opt into `name` + `kind` specs. -fn mongo_collection_specs_from_agent_response( +pub(crate) fn mongo_collection_specs_from_agent_response( value: serde_json::Value, ) -> Result, String> { let values = diff --git a/crates/dbx-core/src/query/mongo_ops.rs b/crates/dbx-core/src/query/mongo_ops.rs index c88114711f..efc0916b64 100644 --- a/crates/dbx-core/src/query/mongo_ops.rs +++ b/crates/dbx-core/src/query/mongo_ops.rs @@ -40,6 +40,38 @@ async fn agent_run_command_document( mongo_driver::json_object_to_document_extended_json(response) } +/// The database view's per-collection "rows"/"size" columns over the legacy agent: one +/// `collStats` per non-view collection through `runCommand`, mirroring +/// `mongo_driver::list_object_statistics`. A collection whose `collStats` fails is left out so +/// the view keeps that row blank instead of failing the whole listing. The agent connection is +/// a single RPC channel, so the round trips run sequentially. +pub async fn mongo_agent_list_object_statistics( + client: &crate::db::agent_driver::PooledAgentClient, + database: &str, +) -> Result, String> { + let database = database.trim(); + if database.is_empty() { + return Err("Database name is required".to_string()); + } + let specs = { + let mut client = client.lock().await; + if !client.supports_capability(AgentCapability::MongoRunCommand) { + return Ok(Vec::new()); + } + crate::document_ops::mongo_collection_specs_from_agent_response( + client.mongo_list_collection_specs(database).await?, + )? + }; + let mut statistics = Vec::with_capacity(specs.len()); + for spec in specs.into_iter().filter(|spec| spec.kind != mongo_driver::MongoCollectionKind::View) { + let Ok(command) = mongo_driver::collection_stats_command(&spec.name, None) else { continue }; + if let Ok(result) = agent_run_command_document(client, database, command, "collection stats").await { + statistics.push(mongo_driver::object_statistics_from_collection_stats(&spec.name, database, &result)); + } + } + Ok(statistics) +} + async fn ensure_document_pool(state: &AppState, connection_id: &str) -> Result<(), String> { state.get_or_create_pool(connection_id, None).await.map(|_| ()) } diff --git a/crates/dbx-core/src/schema/mod.rs b/crates/dbx-core/src/schema/mod.rs index 726130ebd5..6c81865de3 100644 --- a/crates/dbx-core/src/schema/mod.rs +++ b/crates/dbx-core/src/schema/mod.rs @@ -6318,6 +6318,9 @@ async fn list_object_statistics_once( } } if let Some(client) = extract_pool!(pool_handle.as_ref(), Agent) { + if db_config.as_ref().is_some_and(|config| config.db_type == DatabaseType::MongoDb) { + return crate::mongo_ops::mongo_agent_list_object_statistics(&client, database).await; + } if db_config.as_ref().is_some_and(|config| config.db_type == DatabaseType::Oracle) { return oracle_agent_list_object_statistics( client, diff --git a/crates/dbx-core/tests/live_mongodb_legacy_agent_commands.rs b/crates/dbx-core/tests/live_mongodb_legacy_agent_commands.rs index 3b207d3c8a..5bc6084b90 100644 --- a/crates/dbx-core/tests/live_mongodb_legacy_agent_commands.rs +++ b/crates/dbx-core/tests/live_mongodb_legacy_agent_commands.rs @@ -234,6 +234,21 @@ async fn collection_stats_and_create_database_run_over_the_legacy_agent() { ) .await; + // The database view's object statistics take the same numbers per collection, skipping views. + command(&state, id, &database, doc! { "create": "orders_view", "viewOn": "orders", "pipeline": [] }).await; + command(&state, id, &database, doc! { "create": "empty" }).await; + let mut objects = dbx_core::schema::list_object_statistics_core(&state, id, &database, "").await.unwrap(); + objects.sort_by(|a, b| a.name.cmp(&b.name)); + // `system.views` is a real collection and is listed on the native path too; the view is not. + let names: Vec<&str> = objects.iter().map(|o| o.name.as_str()).collect(); + assert!(!names.contains(&"orders_view"), "views own no storage and are skipped: {names:?}"); + let by_name = + |name: &str| objects.iter().find(|o| o.name == name).unwrap_or_else(|| panic!("{name} missing: {names:?}")); + assert_eq!(by_name("orders").estimated_rows, Some(3), "{objects:?}"); + assert!(by_name("orders").total_bytes.is_some_and(|bytes| bytes > 0), "{objects:?}"); + assert_eq!(by_name("orders").schema.as_deref(), Some(database.as_str())); + assert_eq!(by_name("empty").estimated_rows, Some(0), "{objects:?}"); + let stats = mongo_collection_stats_core(&state, id, &database, "orders", None).await.unwrap(); assert_eq!(stats.count, serde_json::json!(3), "{stats:?}"); assert_eq!(stats.nindexes, serde_json::json!(2), "{stats:?}"); diff --git a/crates/dbx-drivers/src/db/mongo_driver.rs b/crates/dbx-drivers/src/db/mongo_driver.rs index 8621bd763d..120dc56cce 100644 --- a/crates/dbx-drivers/src/db/mongo_driver.rs +++ b/crates/dbx-drivers/src/db/mongo_driver.rs @@ -311,7 +311,7 @@ pub async fn list_object_statistics(client: &Client, database: &str) -> Result ObjectStatistics { +pub fn object_statistics_from_collection_stats(name: &str, database: &str, result: &Document) -> ObjectStatistics { ObjectStatistics { name: name.to_string(), schema: Some(database.to_string()),