From 489da1b7dbd59969485472acbb95bbb0082eda26 Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Tue, 23 Jun 2026 19:47:10 +0800 Subject: [PATCH] fix(mongodb): support legacy indexes and driver mode --- .../com/dbx/agent/mongodb/MongoAgent.java | 48 +++++++++++++ .../com/dbx/agent/mongodb/MongoAgentTest.java | 31 +++++++++ .../connection/ConnectionDialog.vue | 1 + apps/desktop/src/stores/connectionStore.ts | 25 +++++++ src-tauri/src/commands/connection.rs | 68 +++++++++++++++++-- 5 files changed, 166 insertions(+), 7 deletions(-) diff --git a/agents/drivers/mongodb/src/main/java/com/dbx/agent/mongodb/MongoAgent.java b/agents/drivers/mongodb/src/main/java/com/dbx/agent/mongodb/MongoAgent.java index 93f2fae86..073b9d578 100644 --- a/agents/drivers/mongodb/src/main/java/com/dbx/agent/mongodb/MongoAgent.java +++ b/agents/drivers/mongodb/src/main/java/com/dbx/agent/mongodb/MongoAgent.java @@ -1,6 +1,7 @@ package com.dbx.agent.mongodb; import com.dbx.agent.AgentProtocol; +import com.dbx.agent.IndexInfo; import com.google.gson.Gson; import com.google.gson.JsonElement; import com.google.gson.JsonNull; @@ -297,6 +298,52 @@ public final class MongoAgent { return result; } + private static Object listIndexes(JsonObject params) { + MongoClient c = requireClient(); + String database = params.get("database").getAsString(); + String collection = params.get("table").getAsString(); + List result = new ArrayList<>(); + for (Document index : c.getDatabase(database).getCollection(collection).listIndexes()) { + result.add(indexInfoFromDocument(index)); + } + return result; + } + + static IndexInfo indexInfoFromDocument(Document index) { + Document keys = index.get("key") instanceof Document document ? document : new Document(); + String name = index.getString("name"); + if (name == null || name.isBlank()) { + List parts = new ArrayList<>(); + for (Map.Entry entry : keys.entrySet()) { + parts.add(entry.getKey() + "_" + String.valueOf(entry.getValue())); + } + name = String.join("_", parts); + } + + List columns = new ArrayList<>(keys.keySet()); + String indexType = null; + if (!keys.isEmpty()) { + List parts = new ArrayList<>(); + for (Map.Entry entry : keys.entrySet()) { + parts.add(entry.getKey() + ": " + String.valueOf(entry.getValue())); + } + indexType = String.join(", ", parts); + } + + Object unique = index.get("unique"); + Document filter = index.get("partialFilterExpression") instanceof Document document ? document : null; + return new IndexInfo( + name, + columns, + unique instanceof Boolean && (Boolean) unique, + "_id_".equals(name), + filter == null ? null : filter.toJson(), + indexType, + null, + null + ); + } + private static Object findDocuments(JsonObject params) { MongoClient c = requireClient(); String database = params.get("database").getAsString(); @@ -510,6 +557,7 @@ public final class MongoAgent { case AgentProtocol.METHOD_CONNECT -> connect(params); case AgentProtocol.MONGO_METHOD_LIST_DATABASES -> listDatabases(); case AgentProtocol.MONGO_METHOD_LIST_COLLECTIONS -> listCollections(params); + case AgentProtocol.METHOD_LIST_INDEXES -> listIndexes(params); case AgentProtocol.MONGO_METHOD_FIND_DOCUMENTS -> findDocuments(params); case AgentProtocol.MONGO_METHOD_INSERT_DOCUMENT -> insertDocument(params); case AgentProtocol.MONGO_METHOD_UPDATE_DOCUMENT -> updateDocument(params); diff --git a/agents/drivers/mongodb/src/test/java/com/dbx/agent/mongodb/MongoAgentTest.java b/agents/drivers/mongodb/src/test/java/com/dbx/agent/mongodb/MongoAgentTest.java index 7c1458827..03c32f404 100644 --- a/agents/drivers/mongodb/src/test/java/com/dbx/agent/mongodb/MongoAgentTest.java +++ b/agents/drivers/mongodb/src/test/java/com/dbx/agent/mongodb/MongoAgentTest.java @@ -1,11 +1,13 @@ package com.dbx.agent.mongodb; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import com.dbx.agent.AgentProtocol; +import com.dbx.agent.IndexInfo; import com.google.gson.JsonArray; import com.google.gson.JsonObject; import com.google.gson.JsonParser; @@ -94,6 +96,35 @@ class MongoAgentTest { assertTrue(containsCapability(result.getAsJsonArray("capabilities"), AgentProtocol.CAPABILITY_METADATA)); } + @Test + void listIndexesMethodIsRecognizedOverJsonRpc() { + String response = MongoAgent.handleRequest( + "{\"jsonrpc\":\"2.0\",\"id\":8,\"method\":\"list_indexes\"," + + "\"params\":{\"database\":\"app\",\"schema\":\"\",\"table\":\"orders\"}}"); + + JsonObject json = JsonParser.parseString(response).getAsJsonObject(); + assertEquals(8, json.get("id").getAsInt()); + assertEquals("Not connected", json.getAsJsonObject("error").get("message").getAsString()); + assertFalse(json.getAsJsonObject("error").get("message").getAsString().contains("Unknown method")); + } + + @Test + void convertsMongoIndexDocumentToIndexInfo() { + Document index = new Document("name", "idx_user_status") + .append("key", new Document("user_id", 1).append("status", -1)) + .append("unique", true) + .append("partialFilterExpression", new Document("deleted", false)); + + IndexInfo info = MongoAgent.indexInfoFromDocument(index); + + assertEquals("idx_user_status", info.getName()); + assertEquals(java.util.List.of("user_id", "status"), info.getColumns()); + assertEquals(true, info.getIs_unique()); + assertEquals(false, info.getIs_primary()); + assertEquals("user_id: 1, status: -1", info.getIndex_type()); + assertTrue(info.getFilter().contains("\"deleted\"")); + } + @Test void usesAuthSourceFromUrlParamsAsAuthenticationDatabase() { JsonObject connection = new JsonObject(); diff --git a/apps/desktop/src/components/connection/ConnectionDialog.vue b/apps/desktop/src/components/connection/ConnectionDialog.vue index efa5daaba..f114e16ba 100644 --- a/apps/desktop/src/components/connection/ConnectionDialog.vue +++ b/apps/desktop/src/components/connection/ConnectionDialog.vue @@ -392,6 +392,7 @@ const driverProfiles: Record< duckdb: { type: "duckdb", port: 0, user: "", label: "DuckDB", icon: "duckdb" }, access: { type: "access", port: 0, user: "", label: "Microsoft Access", icon: "access" }, mongodb: { type: "mongodb", port: 27017, user: "", label: "MongoDB", icon: "mongodb" }, + "mongodb-legacy": { type: "mongodb", port: 27017, user: "", label: "MongoDB (Legacy)", icon: "mongodb" }, clickhouse: { type: "clickhouse", port: 8123, diff --git a/apps/desktop/src/stores/connectionStore.ts b/apps/desktop/src/stores/connectionStore.ts index 6809b990b..21258ba82 100644 --- a/apps/desktop/src/stores/connectionStore.ts +++ b/apps/desktop/src/stores/connectionStore.ts @@ -56,6 +56,8 @@ import { completionSchemasFromTree, completionTablesFromTree } from "@/lib/compl const PINNED_TREE_NODES_STORAGE_KEY = "dbx-pinned-tree-nodes"; const ACTIVE_CONNECTION_STORAGE_KEY = "dbx-active-connection"; const CONNECTION_HEALTH_CHECK_TTL_MS = 2000; +const MONGO_LEGACY_DRIVER_PROFILE = "mongodb-legacy"; +const MONGO_LEGACY_DRIVER_LABEL = "MongoDB (Legacy)"; function sidebarObjectGroupPageSize(): number { const settingsStore = useSettingsStore(); const size = settingsStore.desktopSettings.sidebar_table_page_size; @@ -336,6 +338,7 @@ export const useConnectionStore = defineStore("connection", () => { sqlserver: "SQL Server", mongodb: "MongoDB", oracle: "Oracle", + "mongodb-legacy": MONGO_LEGACY_DRIVER_LABEL, elasticsearch: "Elasticsearch", qdrant: "Qdrant", milvus: "Milvus", @@ -790,6 +793,26 @@ export const useConnectionStore = defineStore("connection", () => { } } + async function syncMongoLegacyDriverFallback(connectionId: string, previousConfig: ConnectionConfig) { + if (!isDesktop || previousConfig.db_type !== "mongodb" || previousConfig.driver_profile === MONGO_LEGACY_DRIVER_PROFILE) { + return; + } + + const savedConnections = await api.loadConnections().catch(() => null); + const savedConfig = savedConnections?.map((connection) => normalizeConnection(connection)).find((connection) => connection.id === connectionId && connection.driver_profile === MONGO_LEGACY_DRIVER_PROFILE); + if (!savedConfig) return; + + const idx = connections.value.findIndex((connection) => connection.id === connectionId); + if (idx < 0) return; + const nextConnections = [...connections.value]; + nextConnections[idx] = { + ...savedConfig, + driver_label: savedConfig.driver_label || MONGO_LEGACY_DRIVER_LABEL, + }; + connections.value = nextConnections; + rebuildTreeNodes(); + } + async function setDefaultDatabase(connectionId: string, database: string) { const config = getConfig(connectionId); if (!config || config.database === database) return; @@ -868,6 +891,7 @@ export const useConnectionStore = defineStore("connection", () => { try { await beforeConnectHandler?.(config); const id = await withConnectionAttemptTimeout(api.connectDb(config), config); + await syncMongoLegacyDriverFallback(id, config); activeConnectionId.value = id; connectedIds.value.add(id); markConnectionHealthChecked(id); @@ -984,6 +1008,7 @@ export const useConnectionStore = defineStore("connection", () => { try { await beforeConnectHandler?.(config); await withConnectionAttemptTimeout(api.connectDb(config), config); + await syncMongoLegacyDriverFallback(connectionId, config); connectedIds.value.add(connectionId); markConnectionHealthChecked(connectionId); activeConnectionId.value = connectionId; diff --git a/src-tauri/src/commands/connection.rs b/src-tauri/src/commands/connection.rs index 68b8003e9..756512c10 100644 --- a/src-tauri/src/commands/connection.rs +++ b/src-tauri/src/commands/connection.rs @@ -18,12 +18,42 @@ use dbx_core::db::agent_driver::AgentMethod; use dbx_core::models::connection::{rewrite_jdbc_url_host, ConnectionConfig, DatabaseType}; pub use dbx_core::path_utils::expand_tilde; +const MONGO_LEGACY_DRIVER_PROFILE: &str = "mongodb-legacy"; +const MONGO_LEGACY_DRIVER_LABEL: &str = "MongoDB (Legacy)"; + fn mongo_legacy_connect_params(config: &ConnectionConfig, host: &str, port: u16) -> serde_json::Value { serde_json::json!({ "connection": agent_connect_params(config, host, port, config.effective_database().unwrap_or("")) }) } +fn mark_mongo_legacy_driver(config: &mut ConnectionConfig) -> bool { + if config.db_type != DatabaseType::MongoDb { + return false; + } + let changed = config.driver_profile.as_deref() != Some(MONGO_LEGACY_DRIVER_PROFILE) + || config.driver_label.as_deref() != Some(MONGO_LEGACY_DRIVER_LABEL); + config.driver_profile = Some(MONGO_LEGACY_DRIVER_PROFILE.to_string()); + config.driver_label = Some(MONGO_LEGACY_DRIVER_LABEL.to_string()); + changed +} + +async fn persist_mongo_legacy_driver_profile(state: &AppState, config: &ConnectionConfig) -> Result<(), String> { + if config.one_time { + return Ok(()); + } + + let mut configs: Vec = + state.storage.load_connections().await?.into_iter().map(|config| config.canonicalized()).collect(); + let Some(saved_config) = configs.iter_mut().find(|saved_config| saved_config.id == config.id) else { + return Ok(()); + }; + if !mark_mongo_legacy_driver(saved_config) { + return Ok(()); + } + save_connection_configs(state, &configs).await +} + async fn test_agent_connection( state: &Arc, config: &ConnectionConfig, @@ -168,10 +198,16 @@ async fn connect_agent_pool( #[cfg(test)] mod tests { - use super::{load_connection_configs, mongo_legacy_connect_params, save_connection_configs}; - use dbx_core::connection::{AppState, PoolKind}; + use super::{ + mark_mongo_legacy_driver, mongo_legacy_connect_params, MONGO_LEGACY_DRIVER_LABEL, MONGO_LEGACY_DRIVER_PROFILE, + }; use dbx_core::models::connection::{ConnectionConfig, DatabaseType}; - use dbx_core::storage::Storage; + #[cfg(feature = "mq-admin")] + use { + super::{load_connection_configs, save_connection_configs}, + dbx_core::connection::{AppState, PoolKind}, + dbx_core::storage::Storage, + }; fn mongodb_config() -> ConnectionConfig { ConnectionConfig { @@ -222,6 +258,7 @@ mod tests { } } + #[cfg(feature = "mq-admin")] fn mq_config(id: &str, admin_url: &str) -> ConnectionConfig { let mut config = mongodb_config(); config.id = id.to_string(); @@ -259,6 +296,16 @@ mod tests { ); } + #[test] + fn mark_mongo_legacy_driver_updates_profile_and_label() { + let mut config = mongodb_config(); + + assert!(mark_mongo_legacy_driver(&mut config)); + assert_eq!(config.driver_profile.as_deref(), Some(MONGO_LEGACY_DRIVER_PROFILE)); + assert_eq!(config.driver_label.as_deref(), Some(MONGO_LEGACY_DRIVER_LABEL)); + assert!(!mark_mongo_legacy_driver(&mut config)); + } + #[cfg(feature = "mq-admin")] #[tokio::test] async fn save_connection_configs_updates_runtime_cache_and_drops_mq_adapter() { @@ -759,6 +806,8 @@ pub async fn connect_db(state: State<'_, Arc>, config: ConnectionConfi let config = config.canonicalized(); let id = config.id.clone(); let db_config = metadata_connection_config(&config); + let mut connected_config = config.clone(); + let mut connected_db_config = db_config.clone(); state.remove_connection_pools(&id).await; state.reset_connection_transport_for_config(&id, &db_config).await; @@ -829,7 +878,8 @@ pub async fn connect_db(state: State<'_, Arc>, config: ConnectionConfi DatabaseType::DuckDb => return Err("DuckDB support not compiled (enable duckdb-bundled feature)".to_string()), DatabaseType::MongoDb => { if mongo_uses_legacy_driver(&db_config) { - let mut client = state.agent_manager.spawn(&db_config.db_type, Some("mongodb-legacy")).await?; + let mut client = + state.agent_manager.spawn(&db_config.db_type, Some(MONGO_LEGACY_DRIVER_PROFILE)).await?; client .connect(mongo_legacy_connect_params(&db_config, &host, port)) .await @@ -857,13 +907,17 @@ pub async fn connect_db(state: State<'_, Arc>, config: ConnectionConfi }; if should_retry_mongo_with_legacy_driver(&native_err) { log::info!("Native MongoDB driver failed ({native_err}), falling back to agent driver"); - let mut client = state.agent_manager.spawn(&db_config.db_type, Some("mongodb-legacy")).await?; + let mut client = + state.agent_manager.spawn(&db_config.db_type, Some(MONGO_LEGACY_DRIVER_PROFILE)).await?; client.connect(mongo_legacy_connect_params(&db_config, &host, port)).await.map_err(|err| { format!( "{native_err}\n\nFallback with MongoDB (Legacy) driver failed: {}", mongo_legacy_error_with_auth_hint(&err) ) })?; + mark_mongo_legacy_driver(&mut connected_config); + connected_db_config = metadata_connection_config(&connected_config); + persist_mongo_legacy_driver_profile(state.inner(), &connected_config).await?; PoolKind::Agent(std::sync::Arc::new(tokio::sync::Mutex::new(client))) } else { return Err(native_err); @@ -999,8 +1053,8 @@ pub async fn connect_db(state: State<'_, Arc>, config: ConnectionConfi db_type => return Err(format!("Unsupported database type: {db_type:?}")), }; - state.insert_connection_pool(id.clone(), pool, &db_config).await; - state.configs.write().await.insert(id.clone(), config); + state.insert_connection_pool(id.clone(), pool, &connected_db_config).await; + state.configs.write().await.insert(id.clone(), connected_config); Ok(id) }