From 24d952dc12b6c284dab95bcdf06698eacc4bf355 Mon Sep 17 00:00:00 2001 From: lexmin0412 Date: Thu, 2 Jul 2026 12:08:30 +0800 Subject: [PATCH] feat(vector): add Qdrant/Milvus collection dimension display (#2396) --- .../src/components/sidebar/TreeItem.vue | 12 +++ apps/desktop/src/lib/api.ts | 1 + apps/desktop/src/lib/http.ts | 4 + apps/desktop/src/lib/tauri.ts | 4 + crates/dbx-core/src/db/vector_driver.rs | 96 +++++++++++++++++++ crates/dbx-core/src/schema.rs | 19 ++++ crates/dbx-web/src/main.rs | 1 + crates/dbx-web/src/routes/mongo.rs | 23 +++++ src-tauri/src/commands/mongo_cmd.rs | 10 ++ src-tauri/src/lib.rs | 1 + 10 files changed, 171 insertions(+) diff --git a/apps/desktop/src/components/sidebar/TreeItem.vue b/apps/desktop/src/components/sidebar/TreeItem.vue index 317845db5..eae8df97d 100644 --- a/apps/desktop/src/components/sidebar/TreeItem.vue +++ b/apps/desktop/src/components/sidebar/TreeItem.vue @@ -578,6 +578,18 @@ async function toggle() { const collectionRef = node.id.includes("__vector_collection:") ? node.id.split("__vector_collection:").pop() || node.label : node.label; const tab = queryStore.createTab(node.connectionId, node.database || "default", node.label, "vector"); queryStore.updateSql(tab, collectionRef); + api + .vectorGetCollectionDetail(node.connectionId, node.database || "default", collectionRef) + .then((info) => { + if (info.dimension != null) { + if (node.meta) { + (node.meta as Record).dimension = info.dimension; + } else { + node.meta = { dimension: info.dimension } as any; + } + } + }) + .catch(() => {}); } else if (node.type === "database" && node.connectionId && hasTreeNodeDatabaseContext(node)) { const config = connectionStore.getConfig(node.connectionId); const effectiveDbType = effectiveDatabaseTypeForConnection(config); diff --git a/apps/desktop/src/lib/api.ts b/apps/desktop/src/lib/api.ts index 2f45a04b5..f4b5299a3 100644 --- a/apps/desktop/src/lib/api.ts +++ b/apps/desktop/src/lib/api.ts @@ -375,6 +375,7 @@ export const documentListDatabases = forward("documentListDatabases"); export const mongoListDatabases = forward("mongoListDatabases"); export const documentListCollections = forward("documentListCollections"); export const mongoListCollections = forward("mongoListCollections"); +export const vectorGetCollectionDetail = forward("vectorGetCollectionDetail"); export const mongoCreateDatabase = forward("mongoCreateDatabase"); export const mongoDropDatabase = forward("mongoDropDatabase"); export const mongoDropCollection = forward("mongoDropCollection"); diff --git a/apps/desktop/src/lib/http.ts b/apps/desktop/src/lib/http.ts index ee89281f7..eea469078 100644 --- a/apps/desktop/src/lib/http.ts +++ b/apps/desktop/src/lib/http.ts @@ -1768,6 +1768,10 @@ export async function vectorListCollections(connectionId: string, database?: str return documentListCollections(connectionId, database || "default"); } +export async function vectorGetCollectionDetail(connectionId: string, database: string, collection: string): Promise { + return post("/api/mongo/vector-collection-detail", { connectionId, database, collection }); +} + export async function mongoFindDocuments(connectionId: string, database: string, collection: string, skip: number, limit: number, filter?: string, projection?: string, sort?: string, executionId?: string): Promise { return documentFindDocuments(connectionId, database, collection, skip, limit, filter, projection, sort, executionId); } diff --git a/apps/desktop/src/lib/tauri.ts b/apps/desktop/src/lib/tauri.ts index 9f40b42e1..bfd2e4dac 100644 --- a/apps/desktop/src/lib/tauri.ts +++ b/apps/desktop/src/lib/tauri.ts @@ -1480,6 +1480,10 @@ export async function mongoListCollections(connectionId: string, database: strin return documentListCollections(connectionId, database); } +export async function vectorGetCollectionDetail(connectionId: string, database: string, collection: string): Promise { + return invoke("vector_collection_detail", { connectionId, database, collection }); +} + export async function mongoCreateDatabase(connectionId: string, database: string): Promise { return invoke("mongo_create_database", { connectionId, database }); } diff --git a/crates/dbx-core/src/db/vector_driver.rs b/crates/dbx-core/src/db/vector_driver.rs index 58946d2d4..ca19d543d 100644 --- a/crates/dbx-core/src/db/vector_driver.rs +++ b/crates/dbx-core/src/db/vector_driver.rs @@ -241,6 +241,102 @@ async fn list_chroma_collections(client: &VectorClient) -> Result Result { + match client.kind { + VectorDbKind::Qdrant => get_qdrant_collection_detail(client, collection).await, + VectorDbKind::Milvus => get_milvus_collection_detail(client, database, collection).await, + VectorDbKind::Weaviate => { + // Weaviate REST API does not expose vector dimension + Ok(CollectionInfo { name: collection.to_string(), id: collection.to_string(), dimension: None }) + } + VectorDbKind::ChromaDb => get_chroma_collection_detail(client, collection).await, + } +} + +async fn get_qdrant_collection_detail(client: &VectorClient, collection: &str) -> Result { + let body = send_json(client.get(&format!("/collections/{}", path_segment(collection))), "Qdrant").await?; + let dim = body + .pointer("/result/config/params/vectors/size") + .and_then(Value::as_u64) + .or_else(|| { + body.pointer("/result/config/params/vectors") + .and_then(Value::as_object) + .and_then(|obj| obj.values().find_map(|v| v.get("size").and_then(|s| s.as_u64()))) + }) + .map(|d| d as u32); + Ok(CollectionInfo { name: collection.to_string(), id: collection.to_string(), dimension: dim }) +} + +fn milvus_vector_dim_from_field(field: &Value) -> Option { + if let Some(dim) = field.pointer("/params/dim").and_then(Value::as_u64) { + return Some(dim as u32); + } + if let Some(params) = field.get("params").and_then(Value::as_array) { + for param in params { + if param.get("key").and_then(Value::as_str) == Some("dim") { + if let Some(v) = param.get("value").and_then(Value::as_str) { + return v.parse().ok(); + } + if let Some(v) = param.get("value").and_then(Value::as_u64) { + return Some(v as u32); + } + } + } + } + None +} + +async fn get_milvus_collection_detail( + client: &VectorClient, + database: &str, + collection: &str, +) -> Result { + let db_name = if database.is_empty() { "default" } else { database }; + let body = send_json( + client + .post("/v2/vectordb/collections/describe") + .json(&serde_json::json!({ "dbName": db_name, "collectionName": collection })), + "Milvus", + ) + .await?; + if body.get("code").and_then(Value::as_i64) != Some(0) { + let msg = body.get("message").and_then(Value::as_str).unwrap_or("unknown error"); + return Err(format!("Milvus collection detail error: {msg}")); + } + let fields = body.pointer("/data/fields").and_then(Value::as_array); + let dim = fields + .and_then(|f| { + f.iter().find(|f| { + let t = f.get("type"); + t.and_then(Value::as_str) == Some("FloatVector") + || t.and_then(Value::as_str) == Some("BinaryVector") + || t.and_then(Value::as_i64) == Some(101) + || t.and_then(Value::as_i64) == Some(102) + }) + }) + .and_then(milvus_vector_dim_from_field); + Ok(CollectionInfo { name: collection.to_string(), id: collection.to_string(), dimension: dim }) +} + +async fn get_chroma_collection_detail(client: &VectorClient, collection: &str) -> Result { + let body = send_json( + client.get(&format!( + "/api/v2/tenants/default_tenant/databases/default_database/collections/{}", + path_segment(collection) + )), + "ChromaDB", + ) + .await?; + let name = body.get("name").and_then(Value::as_str).unwrap_or(collection); + let id = body.get("id").and_then(Value::as_str).unwrap_or(collection); + let dimension = body.get("dimension").and_then(|v| v.as_u64()).map(|d| d as u32); + Ok(CollectionInfo { name: name.to_string(), id: id.to_string(), dimension }) +} + fn chroma_get_response_to_rows(body: &Value) -> Vec { let ids = body.get("ids").and_then(Value::as_array).cloned().unwrap_or_default(); let docs = body.get("documents").and_then(Value::as_array).cloned().unwrap_or_default(); diff --git a/crates/dbx-core/src/schema.rs b/crates/dbx-core/src/schema.rs index 7b87d0577..9aa835e32 100644 --- a/crates/dbx-core/src/schema.rs +++ b/crates/dbx-core/src/schema.rs @@ -844,6 +844,25 @@ pub async fn list_vector_collections_core( db::vector_driver::list_collections_with_db(&client, database).await } +/// Get detailed metadata for a single vector collection (dimension, config, etc). +pub async fn get_vector_collection_detail_core( + state: &AppState, + connection_id: &str, + database: &str, + collection: &str, +) -> Result { + let pool_key = + state.get_or_create_pool(connection_id, if database.is_empty() { None } else { Some(database) }).await?; + let client = { + let connections = state.connections.read().await; + match connections.get(&pool_key) { + Some(PoolKind::VectorDb(client)) => client.clone(), + _ => return Err("Not a vector database connection".to_string()), + } + }; + db::vector_driver::get_collection_detail(&client, database, collection).await +} + pub async fn get_table_comment_core( state: &AppState, connection_id: &str, diff --git a/crates/dbx-web/src/main.rs b/crates/dbx-web/src/main.rs index 7a4d1b541..fd255b1fa 100644 --- a/crates/dbx-web/src/main.rs +++ b/crates/dbx-web/src/main.rs @@ -442,6 +442,7 @@ async fn main() { // MongoDB .route("/mongo/list-databases", post(routes::mongo::list_databases)) .route("/mongo/list-collections", post(routes::mongo::list_collections)) + .route("/mongo/vector-collection-detail", post(routes::mongo::vector_collection_detail)) .route("/mongo/create-database", post(routes::mongo::create_database)) .route("/mongo/drop-database", post(routes::mongo::drop_database)) .route("/mongo/drop-collection", post(routes::mongo::drop_collection)) diff --git a/crates/dbx-web/src/routes/mongo.rs b/crates/dbx-web/src/routes/mongo.rs index 83a51e8c1..a1654e060 100644 --- a/crates/dbx-web/src/routes/mongo.rs +++ b/crates/dbx-web/src/routes/mongo.rs @@ -196,6 +196,29 @@ pub async fn list_collections( Ok(Json(result)) } +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct VectorCollectionDetailRequest { + pub connection_id: String, + pub database: String, + pub collection: String, +} + +pub async fn vector_collection_detail( + State(state): State>, + Json(req): Json, +) -> Result, AppError> { + let result = dbx_core::schema::get_vector_collection_detail_core( + &state.app, + &req.connection_id, + &req.database, + &req.collection, + ) + .await + .map_err(AppError)?; + Ok(Json(result)) +} + pub async fn create_database( State(state): State>, Json(req): Json, diff --git a/src-tauri/src/commands/mongo_cmd.rs b/src-tauri/src/commands/mongo_cmd.rs index 49ea86804..7d3721088 100644 --- a/src-tauri/src/commands/mongo_cmd.rs +++ b/src-tauri/src/commands/mongo_cmd.rs @@ -40,6 +40,16 @@ pub async fn mongo_list_collections( dbx_core::mongo_ops::mongo_list_collections_core(&state, &connection_id, &database).await } +#[tauri::command] +pub async fn vector_collection_detail( + state: State<'_, Arc>, + connection_id: String, + database: String, + collection: String, +) -> Result { + dbx_core::schema::get_vector_collection_detail_core(&state, &connection_id, &database, &collection).await +} + #[tauri::command] pub async fn mongo_create_database( state: State<'_, Arc>, diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index c51e6ee8a..33f3b6d55 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -848,6 +848,7 @@ pub fn run() { commands::sqlite_backup::backup_sqlite_database, commands::mongo_cmd::mongo_list_databases, commands::mongo_cmd::mongo_list_collections, + commands::mongo_cmd::vector_collection_detail, commands::mongo_cmd::mongo_create_database, commands::mongo_cmd::mongo_drop_database, commands::mongo_cmd::mongo_drop_collection,