From c30f6821d3170485084116499c28c6eeefab3dc8 Mon Sep 17 00:00:00 2001 From: Loc Nguyen <113083148+thailoc-dev@users.noreply.github.com> Date: Mon, 21 Sep 2026 17:34:15 +0700 Subject: [PATCH] feat(mongodb): collection stats and create database on legacy agent --- .../database/databaseFeatureSupport.spec.ts | 5 +- .../lib/database/databaseNamespaceCreation.ts | 2 - crates/dbx-core/src/query/mongo_ops.rs | 14 +++- .../live_mongodb_legacy_agent_commands.rs | 71 ++++++++++++++++++- crates/dbx-drivers/src/db/mongo_driver.rs | 63 ++++++++++------ 5 files changed, 126 insertions(+), 29 deletions(-) diff --git a/apps/desktop/src/lib/__tests__/database/databaseFeatureSupport.spec.ts b/apps/desktop/src/lib/__tests__/database/databaseFeatureSupport.spec.ts index 633ff84425..35672c38ba 100644 --- a/apps/desktop/src/lib/__tests__/database/databaseFeatureSupport.spec.ts +++ b/apps/desktop/src/lib/__tests__/database/databaseFeatureSupport.spec.ts @@ -376,8 +376,9 @@ describe("database namespace creation", () => { expect(connectionNamespaceCreationTarget({ db_type: "duckdb" })).toBe("attach"); expect(connectionNamespaceCreationTarget({ db_type: "sqlite" })).toBe("attach"); expect(connectionNamespaceCreationTarget({ db_type: "mongodb" })).toBe("special"); - expect(connectionNamespaceCreationTarget({ db_type: "mongodb", driver_profile: "mongodb-legacy" })).toBeNull(); - expect(connectionNamespaceCreationTarget({ db_type: "mongodb", driver_profile: "legacy" })).toBeNull(); + // The Legacy Agent creates databases through runCommand, so it gets the same flow. + expect(connectionNamespaceCreationTarget({ db_type: "mongodb", driver_profile: "mongodb-legacy" })).toBe("special"); + expect(connectionNamespaceCreationTarget({ db_type: "mongodb", driver_profile: "legacy" })).toBe("special"); }); it("hides persistent SQLite attachment for memory and SQLCipher connections", () => { diff --git a/apps/desktop/src/lib/database/databaseNamespaceCreation.ts b/apps/desktop/src/lib/database/databaseNamespaceCreation.ts index e2d5599bbd..f8208cdc58 100644 --- a/apps/desktop/src/lib/database/databaseNamespaceCreation.ts +++ b/apps/desktop/src/lib/database/databaseNamespaceCreation.ts @@ -1,5 +1,4 @@ import type { ConnectionConfig, DatabaseType, TreeNodeType } from "@/types/database"; -import { isMongoLegacyDriverProfile } from "@/lib/mongo/mongoCapabilities"; import { connectionIsEffectivelyReadOnly } from "@/lib/database/readOnlyWriteAccess"; export type DatabaseNamespaceCreationTarget = "database" | "schema" | "attach" | "special"; @@ -114,7 +113,6 @@ function namespaceCreationMatrixEntry(connection: NonNullable mongo_driver::create_database(client, database).await, - PoolKind::Agent(_) => Err("MongoDB legacy agent does not support create database".to_string()), + PoolKind::Agent(client) => { + let database = mongo_driver::validate_create_database_name(database)?; + let command = mongodb::bson::doc! { "create": mongo_driver::CREATE_DATABASE_PLACEHOLDER_COLLECTION }; + agent_run_command_document(client, database, command, "create database").await.map(|_| ()) + } _ => Err("Not a MongoDB connection".to_string()), } } @@ -240,7 +244,11 @@ pub async fn mongo_collection_stats_core( let pool = state.pool_handle(connection_id).await.ok_or("Not found")?; match &pool { PoolKind::MongoDb(client) => mongo_driver::collection_stats(client, database, collection, scale).await, - PoolKind::Agent(_) => Err("MongoDB legacy agent does not support collection stats helpers".to_string()), + PoolKind::Agent(client) => { + let command = mongo_driver::collection_stats_command(collection, scale.as_ref())?; + let response = agent_run_command_document(client, database, command, "collection stats").await?; + Ok(mongo_driver::collection_stats_result_from_document(&response)) + } _ => Err("Not a MongoDB connection".to_string()), } } @@ -548,7 +556,7 @@ pub async fn mongo_distinct_core( PoolKind::Agent(client) => { let command = mongo_driver::distinct_command(collection, field, filter)?; let response = agent_run_command_document(client, database, command, "distinct").await?; - mongo_driver::distinct_response_result(&response) + mongo_driver::distinct_response_result(response) } _ => Err("Not a MongoDB connection".to_string()), } 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 0d3abfa2dd..3b207d3c8a 100644 --- a/crates/dbx-core/tests/live_mongodb_legacy_agent_commands.rs +++ b/crates/dbx-core/tests/live_mongodb_legacy_agent_commands.rs @@ -3,7 +3,10 @@ use dbx_core::{ connection::{AppState, PoolKind}, models::connection::ConnectionConfig, - mongo_ops::{execute_mongo_command_core, mongo_run_command_core}, + mongo_ops::{ + execute_mongo_command_core, mongo_collection_stats_core, mongo_create_database_core, mongo_list_databases_core, + mongo_run_command_core, + }, mongo_shell, storage::Storage, }; @@ -201,3 +204,69 @@ async fn distinct_runs_over_the_legacy_agent() { command(&state, id, &database, doc! { "dropDatabase": 1 }).await; } + +#[tokio::test] +#[ignore = "opt-in: DBX_MONGO_LEGACY_DUMP_TEST_HOST (host:port, MongoDB 3.6+ without auth) and an installed MongoDB Legacy Agent; creates temporary databases"] +async fn collection_stats_and_create_database_run_over_the_legacy_agent() { + let endpoint = std::env::var("DBX_MONGO_LEGACY_DUMP_TEST_HOST").expect("DBX_MONGO_LEGACY_DUMP_TEST_HOST"); + let (host, port) = endpoint.split_once(':').expect("host:port"); + let files = tempfile::tempdir().unwrap(); + let database = format!("dbx_legacy_admin_{}", uuid::Uuid::new_v4().simple()); + let state = AppState::new(Storage::open(&files.path().join("storage.db")).await.unwrap()); + let id = "legacy-admin-test"; + let config: ConnectionConfig = serde_json::from_value(serde_json::json!({ "id": id, "name": "Legacy admin test", "db_type": "mongodb", "host": host, "port": port.parse::().unwrap(), "username": "", "password": "", "database": database, "driver_profile": "mongodb-legacy" })).unwrap(); + state.configs.write().await.insert(id.into(), config); + let key = state.get_or_create_pool(id, Some(&database)).await.unwrap(); + assert!(matches!(state.pool_handle(&key).await, Some(PoolKind::Agent(_))), "test must run over the legacy agent"); + + command( + &state, + id, + &database, + doc! { "insert": "orders", "documents": [ { "_id": 1, "n": 1 }, { "_id": 2, "n": 2 }, { "_id": 3, "n": 3 } ] }, + ) + .await; + command( + &state, + id, + &database, + doc! { "createIndexes": "orders", "indexes": [ { "key": { "n": 1 }, "name": "n_1" } ] }, + ) + .await; + + 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:?}"); + assert!(stats.size.as_u64().is_some_and(|size| size > 0), "{stats:?}"); + // scale divides byte sizes; count and index count are unaffected. + let scaled = mongo_collection_stats_core(&state, id, &database, "orders", Some(serde_json::Number::from(1024))) + .await + .unwrap(); + assert_eq!(scaled.count, serde_json::json!(3)); + assert!(scaled.storage_size.as_u64().unwrap() <= stats.storage_size.as_u64().unwrap()); + let missing = mongo_collection_stats_core(&state, id, &database, "no_such_collection", None).await.unwrap_err(); + assert!( + missing.contains("not found") || missing.contains("ns not found") || missing.contains("NamespaceNotFound"), + "{missing}" + ); + + // Create database: the new database becomes visible with its placeholder collection. + let created = format!("{database}_created"); + mongo_create_database_core(&state, id, &created).await.unwrap(); + let databases = mongo_list_databases_core(&state, id).await.unwrap(); + assert!(databases.iter().any(|name| name == &created), "{databases:?}"); + let listed = command(&state, id, &created, doc! { "listCollections": 1 }).await; + let names: Vec = listed + .get_document("cursor") + .unwrap() + .get_array("firstBatch") + .unwrap() + .iter() + .map(|b| b.as_document().unwrap().get_str("name").unwrap().to_string()) + .collect(); + assert_eq!(names, vec!["dbx_init"]); + assert!(mongo_create_database_core(&state, id, " ").await.unwrap_err().contains("Database name is required")); + + command(&state, id, &created, doc! { "dropDatabase": 1 }).await; + command(&state, id, &database, doc! { "dropDatabase": 1 }).await; +} diff --git a/crates/dbx-drivers/src/db/mongo_driver.rs b/crates/dbx-drivers/src/db/mongo_driver.rs index 3ba27d3d28..8621bd763d 100644 --- a/crates/dbx-drivers/src/db/mongo_driver.rs +++ b/crates/dbx-drivers/src/db/mongo_driver.rs @@ -297,7 +297,7 @@ pub async fn list_object_statistics(client: &Client, database: &str) -> Result) -> Document { +/// The `collStats` command behind [`collection_stats`], for drivers that only expose `runCommand`. +pub fn collection_stats_command(collection: &str, scale: Option<&serde_json::Number>) -> Result { + let collection = collection.trim(); + if collection.is_empty() { + return Err("Collection name is required".to_string()); + } let mut command = doc! { "collStats": collection }; if let Some(scale) = scale { command.insert("scale", json_value_to_bson(&serde_json::Value::Number(scale.clone()))); } - command + Ok(command) } -fn collection_stats_result_from_document(result: &Document) -> MongoCollectionStatsResult { +pub fn collection_stats_result_from_document(result: &Document) -> MongoCollectionStatsResult { MongoCollectionStatsResult { count: collection_stats_field(result, "count"), size: collection_stats_field(result, "size"), @@ -750,12 +755,21 @@ fn mongo_namespace_missing(error: &str) -> bool { lower.contains("ns not found") || lower.contains("namespace not found") } +/// MongoDB has no "create database"; a database exists once it holds a collection, so both +/// drivers create this placeholder. +pub const CREATE_DATABASE_PLACEHOLDER_COLLECTION: &str = "dbx_init"; + pub async fn create_database(client: &Client, database: &str) -> Result<(), String> { + let database = validate_create_database_name(database)?; + client.database(database).create_collection(CREATE_DATABASE_PLACEHOLDER_COLLECTION).await.map_err(|e| e.to_string()) +} + +pub fn validate_create_database_name(database: &str) -> Result<&str, String> { let database = database.trim(); if database.is_empty() { return Err("Database name is required".to_string()); } - client.database(database).create_collection("dbx_init").await.map_err(|e| e.to_string()) + Ok(database) } pub async fn drop_database(client: &Client, database: &str) -> Result<(), String> { @@ -1689,7 +1703,12 @@ fn aggregate_options_explain(options: &Document) -> Result { /// array fields contribute their elements rather than the whole array, and the server may /// answer from an index with a DISTINCT_SCAN. Values are returned in the `documents` slot as /// bare scalars, which `mongoDocumentsToQueryResult` already renders as a single column. -fn distinct_filter_document(filter: Option<&str>) -> Result { +/// Validates the field name and parses the filter shared by the native `distinct` helper and the +/// command builder, so both drivers reject the same inputs with the same message. +fn distinct_inputs(field: &str, filter: Option<&str>) -> Result { + if field.trim().is_empty() { + return Err("Distinct field name is required".to_string()); + } match filter { Some(f) if !f.trim().is_empty() => { let json: serde_json::Value = serde_json::from_str(f).map_err(|e| format!("Invalid filter JSON: {e}"))?; @@ -1716,17 +1735,14 @@ fn distinct_values_result(values: Vec) -> MongoDocumentResult { /// The `distinct` server command behind `distinct()`, for drivers that only expose `runCommand` /// (the legacy agent). Same field and filter validation as [`distinct`]. pub fn distinct_command(collection: &str, field: &str, filter: Option<&str>) -> Result { - if field.trim().is_empty() { - return Err("Distinct field name is required".to_string()); - } - let query = distinct_filter_document(filter)?; + let query = distinct_inputs(field, filter)?; Ok(doc! { "distinct": collection, "key": field, "query": query }) } /// Turns a `distinct` response (`{ values: [...] }`) into the same result shape as [`distinct`]. -pub fn distinct_response_result(response: &Document) -> Result { - match response.get("values") { - Some(Bson::Array(values)) => Ok(distinct_values_result(values.clone())), +pub fn distinct_response_result(mut response: Document) -> Result { + match response.remove("values") { + Some(Bson::Array(values)) => Ok(distinct_values_result(values)), other => Err(format!("Unexpected distinct values: {other:?}")), } } @@ -1738,10 +1754,7 @@ pub async fn distinct( field: &str, filter: Option<&str>, ) -> Result { - if field.trim().is_empty() { - return Err("Distinct field name is required".to_string()); - } - let filter_doc = distinct_filter_document(filter)?; + let filter_doc = distinct_inputs(field, filter)?; let col = client.database(database).collection::(collection); let values = col.distinct(field, filter_doc).await.map_err(|e| e.to_string())?; Ok(distinct_values_result(values)) @@ -4989,10 +5002,18 @@ mod tests { #[test] fn collection_stats_command_serializes_scale() { - let command = collection_stats_command_document("users", Some(&serde_json::Number::from(1024))); + let command = collection_stats_command("users", Some(&serde_json::Number::from(1024))).unwrap(); assert_eq!(command.get_str("collStats").unwrap(), "users"); assert!(matches!(command.get("scale"), Some(Bson::Int64(1024)))); + assert!(!collection_stats_command("users", None).unwrap().contains_key("scale")); + assert!(collection_stats_command(" ", None).unwrap_err().contains("Collection name is required")); + } + + #[test] + fn create_database_name_is_trimmed_and_required() { + assert_eq!(validate_create_database_name(" app "), Ok("app")); + assert!(validate_create_database_name(" ").unwrap_err().contains("Database name is required")); } #[test] @@ -5128,10 +5149,10 @@ mod tests { assert!(distinct_command("users", " ", None).unwrap_err().contains("field name is required")); assert!(distinct_command("users", "status", Some("{not json")).unwrap_err().contains("Invalid filter JSON")); - let result = distinct_response_result(&doc! { "values": ["a", 2i64, Bson::Null], "ok": 1.0 }).unwrap(); + let result = distinct_response_result(doc! { "values": ["a", 2i64, Bson::Null], "ok": 1.0 }).unwrap(); assert_eq!(result.total, 3); assert_eq!(result.documents, vec![serde_json::json!("a"), serde_json::json!(2), serde_json::Value::Null]); - assert!(distinct_response_result(&doc! { "ok": 1.0 }).is_err()); + assert!(distinct_response_result(doc! { "ok": 1.0 }).is_err()); } #[test]