From d7e1b47aa46bedf5b05427a5a96dc3afa6d6fe83 Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Sun, 20 Sep 2026 08:40:13 +0800 Subject: [PATCH] fix(mongodb): infer collection fields for AI metadata tools Closes #9398 --- crates/dbx-core/src/ai/agent_tools.rs | 28 ++- crates/dbx-core/src/schema/mod.rs | 4 + crates/dbx-core/src/schema/mongodb_columns.rs | 158 ++++++++++++ .../dbx-core/tests/mongodb_column_metadata.rs | 237 ++++++++++++++++++ 4 files changed, 420 insertions(+), 7 deletions(-) create mode 100644 crates/dbx-core/src/schema/mongodb_columns.rs create mode 100644 crates/dbx-core/tests/mongodb_column_metadata.rs diff --git a/crates/dbx-core/src/ai/agent_tools.rs b/crates/dbx-core/src/ai/agent_tools.rs index ab14f0d6b2..db0ecc2b77 100644 --- a/crates/dbx-core/src/ai/agent_tools.rs +++ b/crates/dbx-core/src/ai/agent_tools.rs @@ -344,7 +344,7 @@ pub fn read_only_tools(db_type: DatabaseType) -> Vec { if is_vector_db(db_type) { vec![list_collections_tool(), get_current_time_tool()] } else { - vec![list_databases_tool(), list_tables_tool(), get_columns_tool(), get_current_time_tool()] + vec![list_databases_tool(), list_tables_tool(), get_columns_tool(db_type), get_current_time_tool()] } } @@ -355,7 +355,7 @@ pub fn all_tools(db_type: DatabaseType, sql_permissions: AgentSqlPermissions) -> if is_vector_db(db_type) { return vec![list_collections_tool(), browse_collection_tool(), get_current_time_tool()]; } - let mut tools = vec![list_databases_tool(), list_tables_tool(), get_columns_tool(), get_current_time_tool()]; + let mut tools = vec![list_databases_tool(), list_tables_tool(), get_columns_tool(db_type), get_current_time_tool()]; if db_type == DatabaseType::MongoDb { tools.push(mongo_execute_query_tool(sql_permissions)); } else if supports_sql_query(db_type) { @@ -426,13 +426,18 @@ fn list_tables_tool() -> ToolDefinition { } /// get_columns tool definition. -fn get_columns_tool() -> ToolDefinition { +fn get_columns_tool(db_type: DatabaseType) -> ToolDefinition { ToolDefinition { name: "get_columns", - description: + description: if db_type == DatabaseType::MongoDb { + "Sample up to 100 documents from a MongoDB collection and infer up to 512 top-level field names and types. \ + The sample may be smaller and is not a complete schema or a guarantee of required fields. \ + Nested documents and arrays remain object and array fields; numeric BSON types are reported as number." + } else { "Get column definitions for a table: names, types, primary keys, nullable, defaults, and comments. \ Use this when the user asks about table structure, column details, or field information — \ - even if some schema context was provided, this tool returns the authoritative and complete column list.", + even if some schema context was provided, this tool returns the authoritative and complete column list." + }, parameters: json!({ "type": "object", "properties": { @@ -751,7 +756,7 @@ async fn execute_get_columns( connection_id: &str, database: &str, default_schema: Option<&str>, - _db_type: &DatabaseType, + db_type: &DatabaseType, ) -> Result { let database = effective_database(tool_call, database); let table = tool_call @@ -780,11 +785,20 @@ async fn execute_get_columns( .map_err(|e| format!("Failed to get columns for {table}: {e}"))?; if columns.is_empty() { + if *db_type == DatabaseType::MongoDb { + return Ok(format!( + "No fields could be inferred from collection '{table}': the sample contains no documents or fields." + )); + } return Ok(format!("No columns found for table '{table}'.")); } let mut lines = Vec::new(); - lines.push(format!("Columns of {table}:")); + lines.push(if *db_type == DatabaseType::MongoDb { + format!("Sampled fields of {table} (up to 100 documents and 512 top-level fields; not a complete schema or a guarantee of required fields):") + } else { + format!("Columns of {table}:") + }); for col in &columns { let mut flags: Vec = Vec::new(); if col.is_primary_key { diff --git a/crates/dbx-core/src/schema/mod.rs b/crates/dbx-core/src/schema/mod.rs index f7d078d0af..24292e88a9 100644 --- a/crates/dbx-core/src/schema/mod.rs +++ b/crates/dbx-core/src/schema/mod.rs @@ -18,6 +18,7 @@ use std::sync::{Arc, OnceLock}; use std::time::{Duration, Instant}; mod kingbase; +mod mongodb_columns; macro_rules! extract_pool { ($pool:expr, $variant:ident) => { @@ -7001,6 +7002,9 @@ pub async fn get_columns_core_for_session( table: &str, client_session_id: Option<&str>, ) -> Result, String> { + if connection_config(state, connection_id).await.is_some_and(|config| config.db_type == DatabaseType::MongoDb) { + return Box::pin(mongodb_columns::get_columns(state, connection_id, database, table)).await; + } if client_session_id.is_none() { let metadata_session = EphemeralAgentMetadataSession::open(state, connection_id, Some(database), "columns").await; diff --git a/crates/dbx-core/src/schema/mongodb_columns.rs b/crates/dbx-core/src/schema/mongodb_columns.rs new file mode 100644 index 0000000000..8d0d411aef --- /dev/null +++ b/crates/dbx-core/src/schema/mongodb_columns.rs @@ -0,0 +1,158 @@ +use std::collections::{BTreeMap, BTreeSet}; + +use mongodb::bson::Bson; +use serde_json::{json, Value}; + +use crate::connection::AppState; +use crate::db::mongo_driver::MongoDocumentResult; +use crate::db::ColumnInfo; + +const SAMPLE_SIZE: usize = 100; +const MAX_FIELDS: usize = 512; + +pub(super) async fn get_columns( + state: &AppState, + connection_id: &str, + database: &str, + collection: &str, +) -> Result, String> { + let command = json!({ + "find": collection, + "filter": {}, + "limit": SAMPLE_SIZE, + "batchSize": SAMPLE_SIZE, + "singleBatch": true, + "maxTimeMS": 5000, + }); + let result = crate::mongo_ops::mongo_run_command_core(state, connection_id, database, &command.to_string()).await?; + columns_from_sample(&result) +} + +fn columns_from_sample(result: &MongoDocumentResult) -> Result, String> { + let documents = result + .extended_documents + .as_ref() + .and_then(|documents| documents.first()) + .and_then(|response| response.get("cursor")) + .and_then(|cursor| cursor.get("firstBatch")) + .and_then(Value::as_array) + .ok_or("MongoDB field sampling returned no Extended JSON cursor batch")?; + let sample_count = documents.len().min(SAMPLE_SIZE); + let mut fields: BTreeMap> = BTreeMap::new(); + for document in documents.iter().take(SAMPLE_SIZE) { + let document = Bson::try_from(document.clone()).map_err(|error| format!("Invalid MongoDB sample: {error}"))?; + let Bson::Document(document) = document else { + return Err("MongoDB field sampling returned a non-document value".to_string()); + }; + for (name, value) in document { + if fields.len() < MAX_FIELDS || fields.contains_key(&name) { + fields.entry(name).or_default().insert(sample_type(&value)); + } + } + } + Ok(fields + .into_iter() + .map(|(name, types)| ColumnInfo { + is_primary_key: name == "_id", + name, + data_type: types.into_iter().collect::>().join(" | "), + resolved_schema: None, + is_nullable: true, + column_default: None, + is_unique: false, + extra: Some(format!("sampled from {sample_count} document(s); schema may be incomplete")), + comment: None, + numeric_precision: None, + numeric_scale: None, + character_maximum_length: None, + enum_values: None, + character_set: None, + collation: None, + }) + .collect()) +} + +fn sample_type(value: &Bson) -> &'static str { + match value { + Bson::Double(_) | Bson::Int32(_) | Bson::Int64(_) => "number", + Bson::String(_) => "string", + Bson::Array(_) => "array", + Bson::Document(_) => "object", + Bson::Boolean(_) => "boolean", + Bson::Null => "null", + Bson::RegularExpression(_) => "regex", + Bson::JavaScriptCode(_) | Bson::JavaScriptCodeWithScope(_) => "javascript", + Bson::Timestamp(_) => "timestamp", + Bson::Binary(_) => "binary", + Bson::ObjectId(_) => "objectId", + Bson::DateTime(_) => "date", + Bson::Symbol(_) => "symbol", + Bson::Decimal128(_) => "decimal128", + Bson::Undefined => "undefined", + Bson::MinKey => "minKey", + Bson::MaxKey => "maxKey", + Bson::DbPointer(_) => "dbPointer", + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn result(documents: Vec) -> MongoDocumentResult { + serde_json::from_value(json!({ + "documents": [], + "extended_documents": [{ "cursor": { "firstBatch": documents } }], + "total": 1, + })) + .unwrap() + } + + #[test] + fn mongodb_sample_preserves_special_types_and_mixed_optional_fields() { + let columns = columns_from_sample(&result(vec![ + json!({ + "_id": {"$oid": "000000000000000000000001"}, + "amount": {"$numberLong": "42"}, + "created": {"$date": "2026-01-01T00:00:00Z"}, + "profile": {"city": "sample"}, + "tags": ["one", 2], + "optional": null, + }), + json!({"amount": "unknown", "optional": true, "late": 1}), + ])) + .unwrap(); + let types = + columns.iter().map(|column| (column.name.as_str(), column.data_type.as_str())).collect::>(); + assert_eq!(types["_id"], "objectId"); + assert_eq!(types["amount"], "number | string"); + assert_eq!(types["created"], "date"); + assert_eq!(types["profile"], "object"); + assert_eq!(types["tags"], "array"); + assert_eq!(types["optional"], "boolean | null"); + assert_eq!(types["late"], "number"); + assert!(columns.iter().all(|column| column.is_nullable && column.column_default.is_none())); + } + + #[test] + fn mongodb_sample_handles_empty_and_rejects_malformed_batches() { + assert!(columns_from_sample(&result(vec![])).unwrap().is_empty()); + assert!(columns_from_sample(&result(vec![json!("not a document")])).is_err()); + let mut response = result(vec![]); + response.extended_documents = None; + assert!(columns_from_sample(&response).is_err()); + } + + #[test] + fn mongodb_sample_limits_documents_and_fields() { + let mut documents = vec![json!({"included": 1}); SAMPLE_SIZE]; + documents.push(json!({"outside_sample": 1})); + let columns = columns_from_sample(&result(documents)).unwrap(); + assert_eq!(columns.len(), 1); + assert_eq!(columns[0].name, "included"); + let fields = (0..MAX_FIELDS + 1) + .map(|index| (format!("field_{index}"), json!(index))) + .collect::>(); + assert_eq!(columns_from_sample(&result(vec![Value::Object(fields)])).unwrap().len(), MAX_FIELDS); + } +} diff --git a/crates/dbx-core/tests/mongodb_column_metadata.rs b/crates/dbx-core/tests/mongodb_column_metadata.rs new file mode 100644 index 0000000000..0569378382 --- /dev/null +++ b/crates/dbx-core/tests/mongodb_column_metadata.rs @@ -0,0 +1,237 @@ +#![recursion_limit = "256"] + +use std::sync::Arc; +use std::time::Duration; + +use dbx_core::agent_events::ToolCall; +use dbx_core::agent_tools::{all_tools, execute_tool, read_only_tools, AgentSqlPermissions}; +use dbx_core::connection::{AppState, PoolKind}; +use dbx_core::db::agent_driver::{AgentDriverClient, AgentLaunchSpec}; +use dbx_core::db::mongo_driver; +use dbx_core::models::connection::{ConnectionConfig, DatabaseType}; +use dbx_core::mongo_ops::mongo_run_command_core; +use dbx_core::schema::get_columns_core; +use dbx_core::storage::Storage; +use serde_json::json; + +async fn state_with_pool(pool: PoolKind, db_type: DatabaseType, database: &str) -> (Arc, tempfile::TempDir) { + let directory = tempfile::tempdir().unwrap(); + let storage = Storage::open(&directory.path().join("storage.db")).await.unwrap(); + let state = Arc::new(AppState::new(storage)); + let config: ConnectionConfig = serde_json::from_value(json!({ + "id": "metadata", "name": "Column metadata regression", "db_type": db_type, + "host": "127.0.0.1", "port": 1, "username": "", "password": "", + "database": "default_database", "connect_timeout_secs": 1, + })) + .unwrap(); + state.configs.write().await.insert(config.id.clone(), config); + state + .update_connection_pools(|connections| { + connections.insert("metadata".to_string(), pool.clone()); + connections.insert(format!("metadata:{database}"), pool); + }) + .await; + (state, directory) +} + +async fn columns_tool( + state: &Arc, + database: &str, + collection: &str, + db_type: DatabaseType, +) -> dbx_core::agent_events::ToolResult { + execute_tool( + &ToolCall { + id: "columns".to_string(), + name: "get_columns".to_string(), + arguments: json!({"table": collection, "database": database}), + provider_payload: None, + }, + state, + "metadata", + "default_database", + None, + &db_type, + AgentSqlPermissions::default(), + ) + .await +} + +#[cfg(unix)] +async fn scripted_state(mode: &str, db_type: DatabaseType) -> (Arc, tempfile::TempDir) { + let script = format!( + r#"import json, sys +mode = {mode:?} +print(json.dumps({{"ready": True}}), flush=True) +for line in sys.stdin: + request = json.loads(line) + response = {{"jsonrpc": "2.0", "id": request["id"]}} + try: + method = request["method"] + if method == "handshake": + result = {{"protocolVersion": 1, "agentProtocolVersion": 1, "capabilities": [] if mode == "unsupported" else ["mongo_run_command"]}} + elif method == "get_columns" and mode == "sql": + result = [{{"name": "id", "data_type": "int", "is_nullable": False, "is_primary_key": True}}] + elif method == "run_command": + params = request["params"] + assert params["database"] == "fixture", params + command = json.loads(params["command_json"]) + assert command == {{"find": "items", "filter": {{}}, "limit": 100, "batchSize": 100, "singleBatch": True, "maxTimeMS": 5000}}, command + if mode == "denied": + raise RuntimeError("not authorized to read collection") + documents = [] if mode == "empty" else [{{"_id": {{"$oid": "000000000000000000000001"}}, "amount": 42, "profile": {{"city": "private-value"}}, "tags": ["a"]}}, {{"amount": "unknown", "optional": None}}] + result = {{"documents": [], "extended_documents": [{{"cursor": {{"firstBatch": documents}}}}], "total": 1}} + else: + raise RuntimeError("Unknown method: " + method) + response["result"] = result + except Exception as error: + response["error"] = {{"code": -1, "message": str(error)}} + print(json.dumps(response), flush=True) +"# + ); + let mut client = + AgentDriverClient::spawn(AgentLaunchSpec::new("python3").with_args(["-u", "-c", &script])).await.unwrap(); + client.try_optional_handshake("test").await.unwrap(); + state_with_pool(PoolKind::agent(client), db_type, "fixture").await +} + +#[test] +fn mongodb_columns_tools_describe_sampling_without_changing_sql_tools() { + for tools in + [read_only_tools(DatabaseType::MongoDb), all_tools(DatabaseType::MongoDb, AgentSqlPermissions::default())] + { + let tool = tools.iter().find(|tool| tool.name == "get_columns").unwrap(); + assert!(tool.description.contains("100 documents")); + assert!(tool.description.contains("not a complete schema")); + assert!(!tool.description.contains("authoritative")); + assert!(tool.read_only); + } + let tools = read_only_tools(DatabaseType::Mysql); + assert!(tools + .iter() + .find(|tool| tool.name == "get_columns") + .unwrap() + .description + .contains("authoritative and complete")); +} + +#[cfg(unix)] +#[tokio::test] +async fn mongodb_legacy_columns_use_bounded_sampling_and_database_override() { + let (state, _directory) = scripted_state("sample", DatabaseType::MongoDb).await; + let columns = get_columns_core(&state, "metadata", "fixture", "", "items").await.unwrap(); + assert_eq!(columns.len(), 5); + assert_eq!(columns.iter().find(|column| column.name == "amount").unwrap().data_type, "number | string"); + assert!(columns.iter().all(|column| column.is_nullable)); + let result = columns_tool(&state, "fixture", "items", DatabaseType::MongoDb).await; + assert!(!result.is_error, "{}", result.content); + assert!(result.content.contains("Sampled fields of items")); + assert!(result.content.contains("not a complete schema")); + assert!(result.content.contains("profile: object")); + assert!(result.content.contains("tags: array")); + assert!(!result.content.contains("private-value")); + assert!(!result.content.contains("NOT NULL")); + state.remove_connection_pools("metadata").await; +} + +#[cfg(unix)] +#[tokio::test] +async fn mongodb_legacy_empty_sample_is_not_an_rpc_failure() { + let (state, _directory) = scripted_state("empty", DatabaseType::MongoDb).await; + let result = columns_tool(&state, "fixture", "items", DatabaseType::MongoDb).await; + assert!(!result.is_error, "{}", result.content); + assert!(result.content.contains("No fields could be inferred")); + state.remove_connection_pools("metadata").await; +} + +#[cfg(unix)] +#[tokio::test] +async fn mongodb_legacy_preserves_permission_errors_and_reports_unsupported_drivers() { + for (mode, expected) in [("denied", "not authorized"), ("unsupported", "upgrade or reinstall")] { + let (state, _directory) = scripted_state(mode, DatabaseType::MongoDb).await; + let result = columns_tool(&state, "fixture", "items", DatabaseType::MongoDb).await; + assert!(result.is_error); + assert!(result.content.contains(expected), "{}", result.content); + state.remove_connection_pools("metadata").await; + } +} + +#[cfg(unix)] +#[tokio::test] +async fn sql_agent_columns_keep_the_existing_rpc_and_output() { + let (state, _directory) = scripted_state("sql", DatabaseType::Mysql).await; + let result = columns_tool(&state, "fixture", "items", DatabaseType::Mysql).await; + assert!(!result.is_error, "{}", result.content); + assert_eq!(result.content, "Columns of items:\n - id: int (PK, NOT NULL)"); + state.remove_connection_pools("metadata").await; +} + +async fn verify_live_columns(pool: PoolKind) { + let database = format!("dbx_issue9398_{}", uuid::Uuid::new_v4().simple()); + let (state, _directory) = state_with_pool(pool, DatabaseType::MongoDb, &database).await; + let collection = "youzy_uyk_plandata_tagged_2026"; + let documents = (0..101) + .map(|index| { + let mut document = json!({ + "_id": {"$oid": format!("{:024x}", index + 1)}, + "amount": if index % 2 == 0 { json!(index) } else { json!("unknown") }, + "created": {"$date": "2026-01-01T00:00:00Z"}, + "profile": {"city": "sample"}, "tags": ["sample"], "optional": null, + }); + if index == 100 { + document["outside_sample"] = json!(true); + } + document + }) + .collect::>(); + mongo_run_command_core( + &state, + "metadata", + &database, + &json!({"insert": collection, "documents": documents}).to_string(), + ) + .await + .unwrap(); + mongo_run_command_core(&state, "metadata", &database, r#"{"create":"empty_collection"}"#).await.unwrap(); + let columns = Box::pin(get_columns_core(&state, "metadata", &database, "", collection)).await.unwrap(); + let types = columns + .iter() + .map(|column| (column.name.as_str(), column.data_type.as_str())) + .collect::>(); + assert_eq!(types["_id"], "objectId"); + assert_eq!(types["amount"], "number | string"); + assert_eq!(types["created"], "date"); + assert_eq!(types["profile"], "object"); + assert_eq!(types["tags"], "array"); + assert!(!types.contains_key("outside_sample")); + assert!(columns.iter().all(|column| column.extra.as_deref().unwrap().contains("100 document(s)"))); + let result = Box::pin(columns_tool(&state, &database, collection, DatabaseType::MongoDb)).await; + assert!(!result.is_error, "{}", result.content); + assert!(result.content.contains("not a complete schema")); + assert!(result.content.contains("amount: number | string")); + let empty = Box::pin(columns_tool(&state, &database, "empty_collection", DatabaseType::MongoDb)).await; + assert!(!empty.is_error, "{}", empty.content); + assert!(empty.content.contains("No fields could be inferred")); + mongo_run_command_core(&state, "metadata", &database, r#"{"dropDatabase":1}"#).await.unwrap(); + state.remove_connection_pools("metadata").await; +} + +#[tokio::test] +#[ignore = "requires DBX_LIVE_MONGODB_URL pointing to a test MongoDB server; creates and removes a temporary database"] +async fn live_native_mongodb_column_sampling() { + let url = std::env::var("DBX_LIVE_MONGODB_URL").expect("DBX_LIVE_MONGODB_URL"); + let client = mongo_driver::connect(&url, Duration::from_secs(10), Duration::from_secs(60)).await.unwrap(); + Box::pin(verify_live_columns(PoolKind::MongoDb(client))).await; +} + +#[tokio::test] +#[ignore = "requires DBX_LIVE_MONGODB_AGENT_JAR and DBX_LIVE_MONGODB_LEGACY_PORT via a localhost SSH tunnel; creates and removes a temporary database"] +async fn live_legacy_mongodb_column_sampling() { + let jar = std::env::var("DBX_LIVE_MONGODB_AGENT_JAR").expect("DBX_LIVE_MONGODB_AGENT_JAR"); + let port = + std::env::var("DBX_LIVE_MONGODB_LEGACY_PORT").expect("DBX_LIVE_MONGODB_LEGACY_PORT").parse::().unwrap(); + let mut client = AgentDriverClient::spawn(AgentLaunchSpec::new("java").with_args(["-jar", &jar])).await.unwrap(); + client.try_optional_handshake("test").await.unwrap(); + client.connect(json!({"connection": {"host": "127.0.0.1", "port": port, "database": "admin", "username": "", "password": ""}})).await.unwrap(); + Box::pin(verify_live_columns(PoolKind::agent(client))).await; +}