fix(mongodb): support legacy indexes and driver mode

This commit is contained in:
t8y2 2026-06-23 19:47:10 +08:00
parent 7b03af97c8
commit 489da1b7db
5 changed files with 166 additions and 7 deletions

View File

@ -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<IndexInfo> 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<String> parts = new ArrayList<>();
for (Map.Entry<String, Object> entry : keys.entrySet()) {
parts.add(entry.getKey() + "_" + String.valueOf(entry.getValue()));
}
name = String.join("_", parts);
}
List<String> columns = new ArrayList<>(keys.keySet());
String indexType = null;
if (!keys.isEmpty()) {
List<String> parts = new ArrayList<>();
for (Map.Entry<String, Object> 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);

View File

@ -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();

View File

@ -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,

View File

@ -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;

View File

@ -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<ConnectionConfig> =
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<AppState>,
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<AppState>>, 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<AppState>>, 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<AppState>>, 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<AppState>>, 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)
}