From 69ecde8fdffc1d935b3c0c45a1e3589edc3836a6 Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Thu, 4 Jun 2026 00:04:32 +0800 Subject: [PATCH] perf(doris): speed up metadata browsing --- .../assets/database-drivers.manifest.json | 4 +- crates/dbx-core/src/connection.rs | 52 +++++- crates/dbx-core/src/database_capabilities.rs | 2 +- crates/dbx-core/src/db/mysql.rs | 166 ++++++++++++++++++ crates/dbx-core/src/schema.rs | 19 ++ .../dbx-core/tests/database_capabilities.rs | 4 +- src-tauri/src/commands/connection.rs | 12 +- 7 files changed, 248 insertions(+), 11 deletions(-) diff --git a/crates/dbx-core/assets/database-drivers.manifest.json b/crates/dbx-core/assets/database-drivers.manifest.json index 592b7725e7..10c1089af6 100644 --- a/crates/dbx-core/assets/database-drivers.manifest.json +++ b/crates/dbx-core/assets/database-drivers.manifest.json @@ -111,7 +111,7 @@ "runtimeMode": "native", "mcpMode": "direct", "singleConnectionPool": false, - "metadataConnectionScoped": true, + "metadataConnectionScoped": false, "skipTcpProbe": false, "defaultPort": 9030 }, @@ -121,7 +121,7 @@ "runtimeMode": "native", "mcpMode": "direct", "singleConnectionPool": false, - "metadataConnectionScoped": true, + "metadataConnectionScoped": false, "skipTcpProbe": false, "defaultPort": 9030 }, diff --git a/crates/dbx-core/src/connection.rs b/crates/dbx-core/src/connection.rs index 81ed2a22af..5272eadcc8 100644 --- a/crates/dbx-core/src/connection.rs +++ b/crates/dbx-core/src/connection.rs @@ -137,6 +137,51 @@ pub async fn connect_mysql_metadata_pool( } } +pub async fn connect_bare_metadata_pool( + db_config: &ConnectionConfig, + host: &str, + port: u16, + connect_timeout: std::time::Duration, +) -> Result { + let url = connection_url_for_endpoint(db_config, host, port); + if db_config.effective_database().is_none() { + return db::mysql::connect_bare(&url, connect_timeout).await; + } + + let mut unscoped_config = db_config.clone(); + unscoped_config.database = None; + let unscoped_url = connection_url_for_endpoint(&unscoped_config, host, port); + if unscoped_url == url { + return db::mysql::connect_bare(&url, connect_timeout).await; + } + + let preferred = db::mysql::connect_bare(&url, connect_timeout); + let unscoped = db::mysql::connect_bare(&unscoped_url, connect_timeout); + tokio::pin!(preferred); + tokio::pin!(unscoped); + + tokio::select! { + result = &mut preferred => match result { + Ok(pool) => Ok(pool), + Err(preferred_err) => match (&mut unscoped).await { + Ok(pool) => Ok(pool), + Err(unscoped_err) => Err(format!( + "Connection with the configured database failed: {preferred_err}\n\nConnection without a default database also failed: {unscoped_err}" + )), + }, + }, + result = &mut unscoped => match result { + Ok(pool) => Ok(pool), + Err(unscoped_err) => match (&mut preferred).await { + Ok(pool) => Ok(pool), + Err(preferred_err) => Err(format!( + "Connection with the configured database failed: {preferred_err}\n\nConnection without a default database also failed: {unscoped_err}" + )), + }, + }, + } +} + fn mysql_metadata_fallback_url( config: &ConnectionConfig, db_config: &ConnectionConfig, @@ -276,7 +321,12 @@ impl AppState { PoolKind::Mysql(pool, mode) } DatabaseType::Doris | DatabaseType::StarRocks => { - PoolKind::Mysql(db::mysql::connect_bare(&url, connect_timeout).await?, MysqlMode::Bare) + let pool = if database.is_none() { + connect_bare_metadata_pool(&db_config, &host, port, connect_timeout).await? + } else { + db::mysql::connect_bare(&url, connect_timeout).await? + }; + PoolKind::Mysql(pool, MysqlMode::Bare) } DatabaseType::Postgres | DatabaseType::Redshift | DatabaseType::Gaussdb | DatabaseType::OpenGauss => { PoolKind::Postgres(db::postgres::connect(&url, connect_timeout).await?) diff --git a/crates/dbx-core/src/database_capabilities.rs b/crates/dbx-core/src/database_capabilities.rs index d982647d96..3066e5bd4e 100644 --- a/crates/dbx-core/src/database_capabilities.rs +++ b/crates/dbx-core/src/database_capabilities.rs @@ -31,7 +31,7 @@ pub fn is_single_connection_pool(db_type: &DatabaseType) -> bool { } pub fn is_metadata_connection_scoped(db_type: &DatabaseType) -> bool { - matches!(db_type, DatabaseType::Mysql | DatabaseType::Doris | DatabaseType::StarRocks) + matches!(db_type, DatabaseType::Mysql) } pub fn skips_tcp_probe(db_type: &DatabaseType) -> bool { diff --git a/crates/dbx-core/src/db/mysql.rs b/crates/dbx-core/src/db/mysql.rs index 404e05dca6..f5faa2d2cc 100644 --- a/crates/dbx-core/src/db/mysql.rs +++ b/crates/dbx-core/src/db/mysql.rs @@ -5,6 +5,7 @@ use mysql_async::prelude::*; use percent_encoding::percent_decode_str; use rust_decimal::Decimal; use std::borrow::Cow; +use std::collections::HashMap; use std::path::PathBuf; use std::time::Duration; use std::time::Instant; @@ -67,6 +68,13 @@ fn get_opt_str(row: &mysql_async::Row, name: &str) -> Option { .or_else(|| row_get::, _>(row, name).map(|b| String::from_utf8_lossy(&b).to_string())) } +fn get_opt_metadata_string(row: &mysql_async::Row, name: &str) -> Option { + get_opt_str(row, name) + .or_else(|| row_get::(row, name).map(|value| value.to_string())) + .or_else(|| row_get::(row, name).map(|value| value.to_string())) + .or_else(|| row_get::(row, name).map(|value| value.to_string())) +} + fn numeric_metadata_u64_to_i32(value: Option) -> Option { value.and_then(|v| i32::try_from(v).ok()) } @@ -719,6 +727,15 @@ pub async fn list_databases(pool: &MySqlPool) -> Result, Strin Ok(rows.iter().map(|row| DatabaseInfo { name: get_str(row, 0) }).collect()) } +pub async fn list_databases_show(pool: &MySqlPool) -> Result, String> { + let mut conn = pool.get_conn().await.map_err(|e| e.to_string())?; + let result = conn.query_iter("SHOW DATABASES").await.map_err(|e| e.to_string())?; + let rows: Vec = result.collect_and_drop().await.map_err(|e| e.to_string())?; + let mut databases: Vec = rows.iter().map(|row| DatabaseInfo { name: get_str(row, 0) }).collect(); + databases.sort_by(|a, b| a.name.cmp(&b.name)); + Ok(databases) +} + pub async fn list_tables(pool: &MySqlPool, database: &str) -> Result, String> { let sql = format!( "SELECT TABLE_NAME, TABLE_TYPE, TABLE_COMMENT FROM information_schema.TABLES WHERE TABLE_SCHEMA = {} ORDER BY TABLE_NAME", @@ -740,6 +757,87 @@ pub async fn list_tables(pool: &MySqlPool, database: &str) -> Result, + created_at: Option, + updated_at: Option, +} + +async fn list_table_status_show(pool: &MySqlPool, database: &str) -> Result, String> { + let sql = format!("SHOW TABLE STATUS FROM {}", quote_identifier(database)); + let mut conn = pool.get_conn().await.map_err(|e| e.to_string())?; + let result = conn.query_iter(&sql).await.map_err(|e| e.to_string())?; + let rows: Vec = result.collect_and_drop().await.map_err(|e| e.to_string())?; + Ok(rows + .iter() + .map(|row| { + ( + get_str_by_name(row, "Name"), + TableStatusMeta { + comment: get_opt_metadata_string(row, "Comment").filter(|s| !s.is_empty()), + created_at: get_opt_metadata_string(row, "Create_time"), + updated_at: get_opt_metadata_string(row, "Update_time"), + }, + ) + }) + .filter(|(name, _)| !name.is_empty()) + .collect()) +} + +async fn list_table_names_show(pool: &MySqlPool, database: &str) -> Result, String> { + let sql = format!("SHOW FULL TABLES FROM {}", quote_identifier(database)); + let mut conn = pool.get_conn().await.map_err(|e| e.to_string())?; + let rows: Vec = match conn.query_iter(&sql).await { + Ok(result) => result.collect_and_drop().await.map_err(|e| e.to_string())?, + Err(_) => { + let sql = format!("SHOW TABLES FROM {}", quote_identifier(database)); + let result = conn.query_iter(&sql).await.map_err(|e| e.to_string())?; + result.collect_and_drop().await.map_err(|e| e.to_string())? + } + }; + let mut tables: Vec = rows + .iter() + .map(|row| { + let table_type = get_str(row, 1); + TableInfo { + name: get_str(row, 0), + table_type: if table_type.is_empty() { "TABLE".to_string() } else { table_type }, + comment: None, + parent_schema: None, + parent_name: None, + } + }) + .collect(); + tables.sort_by(|a, b| a.name.cmp(&b.name)); + Ok(tables) +} + +async fn list_tables_show_with_status( + pool: &MySqlPool, + database: &str, +) -> Result<(Vec, HashMap), String> { + let (tables, status) = tokio::join!(list_table_names_show(pool, database), list_table_status_show(pool, database)); + let mut tables = tables?; + let status = match status { + Ok(status) => status, + Err(err) => { + log::warn!("Skipping table status for database `{}`: {}", database, err); + HashMap::new() + } + }; + for table in &mut tables { + if let Some(meta) = status.get(&table.name) { + table.comment = meta.comment.clone(); + } + } + Ok((tables, status)) +} + +pub async fn list_tables_show(pool: &MySqlPool, database: &str) -> Result, String> { + list_tables_show_with_status(pool, database).await.map(|(tables, _)| tables) +} + fn list_tables_objects_sql(database: &str) -> String { format!( "SELECT TABLE_NAME AS object_name, \ @@ -824,6 +922,43 @@ pub async fn list_objects(pool: &MySqlPool, database: &str) -> Result Result, String> { + let (tables, routines) = + tokio::join!(list_tables_show_with_status(pool, database), list_routine_objects(pool, database)); + let (tables, status) = tables?; + let mut objects: Vec = tables + .into_iter() + .map(|table| { + let meta = status.get(&table.name); + ObjectInfo { + name: table.name, + object_type: if table.table_type.eq_ignore_ascii_case("VIEW") { "VIEW" } else { "TABLE" }.to_string(), + schema: Some(database.to_string()), + comment: table.comment, + created_at: meta.and_then(|meta| meta.created_at.clone()), + updated_at: meta.and_then(|meta| meta.updated_at.clone()), + parent_schema: table.parent_schema, + parent_name: table.parent_name, + } + }) + .collect(); + + match routines { + Ok(routines) => objects.extend(routines), + Err(err) => log::warn!("Skipping routines for database `{}` in object browser: {}", database, err), + } + + Ok(objects) +} + +async fn list_routine_objects(pool: &MySqlPool, database: &str) -> Result, String> { + let mut conn = pool.get_conn().await.map_err(|e| e.to_string())?; + let routines_sql = list_routines_sql(database); + let result = conn.query_iter(&routines_sql).await.map_err(|e| e.to_string())?; + let rows: Vec = result.collect_and_drop().await.map_err(|e| e.to_string())?; + Ok(rows.iter().map(|row| row_to_object(row, database)).collect()) +} + pub async fn list_completion_objects(pool: &MySqlPool, database: &str) -> Result, String> { let mut conn = pool.get_conn().await.map_err(|e| e.to_string())?; let mut objects = Vec::new(); @@ -895,6 +1030,37 @@ pub async fn get_columns(pool: &MySqlPool, database: &str, table: &str) -> Resul .collect()) } +pub async fn get_columns_show(pool: &MySqlPool, database: &str, table: &str) -> Result, String> { + let sql = format!("SHOW FULL COLUMNS FROM {}.{}", quote_identifier(database), quote_identifier(table)); + let mut conn = pool.get_conn().await.map_err(|e| e.to_string())?; + let rows: Vec = match conn.query_iter(&sql).await { + Ok(result) => result.collect_and_drop().await.map_err(|e| e.to_string())?, + Err(_) => { + let sql = format!("SHOW COLUMNS FROM {}.{}", quote_identifier(database), quote_identifier(table)); + let result = conn.query_iter(&sql).await.map_err(|e| e.to_string())?; + result.collect_and_drop().await.map_err(|e| e.to_string())? + } + }; + Ok(rows + .iter() + .map(|row| { + let key = get_str_by_name(row, "Key"); + ColumnInfo { + name: get_str_by_name(row, "Field"), + data_type: get_str_by_name(row, "Type"), + is_nullable: get_str_by_name(row, "Null").eq_ignore_ascii_case("YES"), + column_default: get_opt_str(row, "Default"), + is_primary_key: key.eq_ignore_ascii_case("PRI"), + extra: get_opt_str(row, "Extra"), + comment: get_opt_str(row, "Comment").filter(|s| !s.is_empty()), + numeric_precision: None, + numeric_scale: None, + character_maximum_length: None, + } + }) + .collect()) +} + fn query_result_row_limit(max_rows: Option) -> usize { max_rows.unwrap_or(crate::query::MAX_ROWS).max(1) } diff --git a/crates/dbx-core/src/schema.rs b/crates/dbx-core/src/schema.rs index db6be893f2..56773f320b 100644 --- a/crates/dbx-core/src/schema.rs +++ b/crates/dbx-core/src/schema.rs @@ -289,10 +289,14 @@ async fn list_databases_once(state: &AppState, connection_id: &str) -> Result { + db::mysql::list_databases_show(p).await + } PoolKind::Mysql(p, mode) => dispatch_mysql!(p, mode, db::mysql::list_databases, db::ob_oracle::list_databases), PoolKind::Postgres(p) => db::postgres::list_databases(p).await, PoolKind::Sqlite(p) => db::sqlite::list_databases(p).await, @@ -369,6 +373,7 @@ async fn list_tables_once( ) -> Result, String> { let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?; let duckdb_attached_names = duckdb_attached_database_names(state, connection_id).await; + let db_config = connection_config(state, connection_id).await; { let connections = state.connections.read().await; @@ -410,6 +415,9 @@ async fn list_tables_once( let pool = connections.get(&pool_key).ok_or("Pool not found")?; match pool { + PoolKind::Mysql(p, _) if db_config.as_ref().is_some_and(is_doris_family_config) => { + db::mysql::list_tables_show(p, database).await.map(|tables| filter_table_infos(tables, filter, limit)) + } PoolKind::Mysql(p, mode) => { dispatch_mysql!(p, mode, db::mysql::list_tables, db::ob_oracle::list_tables, schema) .map(|tables| filter_table_infos(tables, filter, limit)) @@ -628,6 +636,8 @@ async fn list_objects_once( // Note: mysql and ob_oracle take different second args (database vs schema) if *mode == MysqlMode::OceanBaseOracle { db::ob_oracle::list_objects(p, schema).await + } else if db_config.as_ref().is_some_and(is_doris_family_config) { + db::mysql::list_table_objects_show(p, database).await } else { db::mysql::list_objects(p, database).await } @@ -744,6 +754,7 @@ pub async fn get_columns_core( ) -> Result, String> { let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?; let duckdb_attached_names = duckdb_attached_database_names(state, connection_id).await; + let db_config = connection_config(state, connection_id).await; { let connections = state.connections.read().await; @@ -805,6 +816,9 @@ pub async fn get_columns_core( let pool = connections.get(&pool_key).ok_or("Pool not found")?; match pool { + PoolKind::Mysql(p, _) if db_config.as_ref().is_some_and(is_doris_family_config) => { + db::mysql::get_columns_show(p, database, table).await.map(deduplicate_column_infos) + } PoolKind::Mysql(p, mode) => { dispatch_mysql!(p, mode, db::mysql::get_columns, db::ob_oracle::get_columns, database, table) .map(deduplicate_column_infos) @@ -1018,6 +1032,11 @@ fn is_opengauss_family_config(config: &ConnectionConfig) -> bool { || matches!(config.driver_profile.as_deref(), Some("opengauss" | "gaussdb")) } +fn is_doris_family_config(config: &ConnectionConfig) -> bool { + matches!(config.db_type, DatabaseType::Doris | DatabaseType::StarRocks) + || matches!(config.driver_profile.as_deref(), Some("doris" | "selectdb" | "starrocks")) +} + fn sql_string(value: &str) -> String { format!("'{}'", value.replace('\'', "''")) } diff --git a/crates/dbx-core/tests/database_capabilities.rs b/crates/dbx-core/tests/database_capabilities.rs index 321d0ccb40..d26af273e1 100644 --- a/crates/dbx-core/tests/database_capabilities.rs +++ b/crates/dbx-core/tests/database_capabilities.rs @@ -107,8 +107,8 @@ fn identifies_single_connection_pool_types() { #[test] fn identifies_metadata_connections_that_drop_database_scope() { assert!(is_metadata_connection_scoped(&DatabaseType::Mysql)); - assert!(is_metadata_connection_scoped(&DatabaseType::Doris)); - assert!(is_metadata_connection_scoped(&DatabaseType::StarRocks)); + assert!(!is_metadata_connection_scoped(&DatabaseType::Doris)); + assert!(!is_metadata_connection_scoped(&DatabaseType::StarRocks)); assert!(!is_metadata_connection_scoped(&DatabaseType::Postgres)); assert!(!is_metadata_connection_scoped(&DatabaseType::Oracle)); } diff --git a/src-tauri/src/commands/connection.rs b/src-tauri/src/commands/connection.rs index 8a33e2ee5d..5be1349edc 100644 --- a/src-tauri/src/commands/connection.rs +++ b/src-tauri/src/commands/connection.rs @@ -6,8 +6,9 @@ pub use dbx_core::agent_connection::{ oracle_auth_fallback_profiles, should_retry_oracle_with_10g_driver, }; pub use dbx_core::connection::{ - connect_mysql_metadata_pool, connection_url_for_endpoint, expand_tilde, metadata_connection_config, - probe_connection_endpoint, redacted_connection_url_for_endpoint, AppState, MysqlMode, PoolKind, + connect_bare_metadata_pool, connect_mysql_metadata_pool, connection_url_for_endpoint, expand_tilde, + metadata_connection_config, probe_connection_endpoint, redacted_connection_url_for_endpoint, AppState, MysqlMode, + PoolKind, }; use dbx_core::database_capabilities; use dbx_core::db; @@ -439,9 +440,10 @@ pub async fn connect_db(state: State<'_, Arc>, config: ConnectionConfi let (pool, mode) = connect_mysql_metadata_pool(&config, &db_config, &host, port, connect_timeout).await?; PoolKind::Mysql(pool, mode) } - DatabaseType::Doris | DatabaseType::StarRocks => { - PoolKind::Mysql(db::mysql::connect_bare(&url, connect_timeout).await?, MysqlMode::Bare) - } + DatabaseType::Doris | DatabaseType::StarRocks => PoolKind::Mysql( + connect_bare_metadata_pool(&db_config, &host, port, connect_timeout).await?, + MysqlMode::Bare, + ), DatabaseType::Postgres | DatabaseType::Redshift | DatabaseType::Gaussdb | DatabaseType::OpenGauss => { PoolKind::Postgres(db::postgres::connect(&url, connect_timeout).await?) }