fix(mongodb): support legacy indexes and driver mode
This commit is contained in:
parent
7b03af97c8
commit
489da1b7db
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in New Issue