fix(mongodb): infer collection fields for AI metadata tools

Closes #9398

crates/dbx-core/src/ai/agent_tools.rs

@@ -344,7 +344,7 @@ pub fn read_only_tools(db_type: DatabaseType) -> Vec<ToolDefinition> {
     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<String, String> {
     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<String> = Vec::new();
         if col.is_primary_key {

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<Vec<db::ColumnInfo>, 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;

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<Vec<ColumnInfo>, 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<Vec<ColumnInfo>, 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<String, BTreeSet<&str>> = 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::<Vec<_>>().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<Value>) -> 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::<BTreeMap<_, _>>();
+        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::<serde_json::Map<_, _>>();
+        assert_eq!(columns_from_sample(&result(vec![Value::Object(fields)])).unwrap().len(), MAX_FIELDS);
+    }
+}

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<AppState>, 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<AppState>,
+    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<AppState>, 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::<Vec<_>>();
+    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::<std::collections::BTreeMap<_, _>>();
+    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::<u16>().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;
+}