github/dbx
4
0
Fork 0
mirror of https://github.com/t8y2/dbx.git synced 2026-10-23 18:20:01 +00:00
Code Issues Packages Projects Releases Wiki Activity

perf(doris): speed up metadata browsing

Browse Source
This commit is contained in:
t8y2
2026-06-04 00:05:14 +08:00
parent f540770bce
commit 69ecde8fdf
7 changed files with 248 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/dbx-core/assets/database-drivers.manifest.json
+2 -2
View File
@@ -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
},
crates/dbx-core/src/connection.rs
+51 -1
View File
@@ -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<db::mysql::MySqlPool, String> {
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?)
crates/dbx-core/src/database_capabilities.rs
+1 -1
View File
@@ -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 {
crates/dbx-core/src/db/mysql.rs
+166
View File
@@ -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<String> {
.or_else(|| row_get::<Vec<u8>, _>(row, name).map(|b| String::from_utf8_lossy(&b).to_string()))
}
fn get_opt_metadata_string(row: &mysql_async::Row, name: &str) -> Option<String> {
get_opt_str(row, name)
.or_else(|| row_get::<NaiveDateTime, _>(row, name).map(|value| value.to_string()))
.or_else(|| row_get::<NaiveDate, _>(row, name).map(|value| value.to_string()))
.or_else(|| row_get::<NaiveTime, _>(row, name).map(|value| value.to_string()))
}
fn numeric_metadata_u64_to_i32(value: Option<u64>) -> Option<i32> {
value.and_then(|v| i32::try_from(v).ok())
}
@@ -719,6 +727,15 @@ pub async fn list_databases(pool: &MySqlPool) -> Result<Vec<DatabaseInfo>, Strin
Ok(rows.iter().map(|row| DatabaseInfo { name: get_str(row, 0) }).collect())
}
pub async fn list_databases_show(pool: &MySqlPool) -> Result<Vec<DatabaseInfo>, 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<mysql_async::Row> = result.collect_and_drop().await.map_err(|e| e.to_string())?;
let mut databases: Vec<DatabaseInfo> = 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<Vec<TableInfo>, 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<Vec<TableIn
.collect())
}
#[derive(Clone, Debug, Default)]
struct TableStatusMeta {
comment: Option<String>,
created_at: Option<String>,
updated_at: Option<String>,
}
async fn list_table_status_show(pool: &MySqlPool, database: &str) -> Result<HashMap<String, TableStatusMeta>, 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<mysql_async::Row> = 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<Vec<TableInfo>, 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<mysql_async::Row> = 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<TableInfo> = 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<TableInfo>, HashMap<String, TableStatusMeta>), 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<Vec<TableInfo>, 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<Vec<Object
Ok(objects)
}
pub async fn list_table_objects_show(pool: &MySqlPool, database: &str) -> Result<Vec<ObjectInfo>, 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<ObjectInfo> = 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<Vec<ObjectInfo>, 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<mysql_async::Row> = 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<Vec<ObjectInfo>, 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<Vec<ColumnInfo>, 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<mysql_async::Row> = 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>) -> usize {
max_rows.unwrap_or(crate::query::MAX_ROWS).max(1)
}
crates/dbx-core/src/schema.rs
+19
View File
@@ -289,10 +289,14 @@ async fn list_databases_once(state: &AppState, connection_id: &str) -> Result<Ve
}
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;
let pool = connections.get(connection_id).ok_or("Connection not found")?;
match pool {
PoolKind::Mysql(p, _) if db_config.as_ref().is_some_and(is_doris_family_config) => {
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<Vec<db::TableInfo>, 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<Vec<db::ColumnInfo>, 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('\'', "''"))
}
crates/dbx-core/tests/database_capabilities.rs
+2 -2
View File
@@ -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));
}
src-tauri/src/commands/connection.rs
+7 -5
View File
@@ -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<AppState>>, 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?)
}