From e4796853120740f974729f4037eaafb9059bfefd Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Fri, 24 Jul 2026 15:21:23 +0800 Subject: [PATCH] feat(kingbase): extract schema metadata helpers --- crates/dbx-core/src/schema.rs | 236 ++----------------------- crates/dbx-core/src/schema/kingbase.rs | 219 +++++++++++++++++++++++ 2 files changed, 233 insertions(+), 222 deletions(-) create mode 100644 crates/dbx-core/src/schema/kingbase.rs diff --git a/crates/dbx-core/src/schema.rs b/crates/dbx-core/src/schema.rs index 0900163836..fb33e8d1d0 100644 --- a/crates/dbx-core/src/schema.rs +++ b/crates/dbx-core/src/schema.rs @@ -7,6 +7,8 @@ use std::future::Future; use std::sync::Arc; use std::time::{Duration, Instant}; +mod kingbase; + macro_rules! extract_pool { ($connections:expr, $key:expr, $variant:ident) => { $connections.get($key).and_then(|v| match v { @@ -1716,19 +1718,6 @@ fn dameng_object_statistics_rows_only_sql(schema: &str) -> String { ) } -fn kingbase_object_statistics_sql(schema: &str) -> String { - format!( - "SELECT c.relname, n.nspname, \ - CAST(CASE WHEN c.reltuples < 0 THEN 0 ELSE c.reltuples END AS BIGINT) AS estimated_rows, \ - CAST(sys_total_relation_size(c.oid) AS BIGINT) AS total_bytes \ - FROM sys_catalog.sys_class c \ - JOIN sys_catalog.sys_namespace n ON n.oid = c.relnamespace \ - WHERE n.nspname = {} AND c.relkind IN ('r','m','f','p') \ - ORDER BY c.relname", - sql_string(schema), - ) -} - fn gbase8a_object_statistics_sql(database: &str) -> String { format!( "SELECT TABLE_NAME, TABLE_SCHEMA, TABLE_ROWS, \ @@ -2675,14 +2664,13 @@ mod tests { dameng_object_statistics_rows_only_sql, dameng_object_statistics_user_segments_sql, deduplicate_column_infos, filter_mysql_system_databases_for_config, filter_object_infos, filter_table_infos, filter_visible_schema_names, gbase8a_object_statistics_sql, is_agent_postgres_metadata_fallback_config, is_retryable_metadata_error, - kingbase_object_statistics_sql, mysql_object_source_sql, mysql_table_metadata_catalog, - normalize_information_schema_table_type, oracle_columns_from_query_result, oracle_columns_sql, - oracle_object_statistics_dba_segments_sql, oracle_object_statistics_from_query_result, - oracle_object_statistics_rows_only_sql, oracle_object_statistics_sql, - oracle_object_statistics_user_segments_sql, oracle_table_comment_from_query_result, oracle_table_comment_sql, - oracle_table_comments_from_query_result, oracle_table_comments_sql, presto_like_columns_from_query_result, - presto_like_information_schema_columns_sql, presto_like_information_schema_tables_sql, - presto_like_tables_from_query_result, visible_schema_filter, + mysql_object_source_sql, mysql_table_metadata_catalog, normalize_information_schema_table_type, + oracle_columns_from_query_result, oracle_columns_sql, oracle_object_statistics_dba_segments_sql, + oracle_object_statistics_from_query_result, oracle_object_statistics_rows_only_sql, + oracle_object_statistics_sql, oracle_object_statistics_user_segments_sql, + oracle_table_comment_from_query_result, oracle_table_comment_sql, oracle_table_comments_from_query_result, + oracle_table_comments_sql, presto_like_columns_from_query_result, presto_like_information_schema_columns_sql, + presto_like_information_schema_tables_sql, presto_like_tables_from_query_result, visible_schema_filter, }; #[cfg(feature = "duckdb-bundled")] use super::{ @@ -3444,57 +3432,6 @@ mod tests { assert!(!is_agent_postgres_metadata_fallback_config(&test_connection_config(DatabaseType::Mysql))); } - #[test] - fn kingbase_extension_sql_uses_sys_catalog_and_escapes_schema() { - let sql = super::kingbase_list_extensions_sql(Some("app's"), super::KingbaseExtensionCatalog::Sys); - - assert!(sql.contains("FROM sys_catalog.sys_extension e")); - assert!(sql.contains("JOIN sys_catalog.sys_namespace n")); - assert!(sql.contains("LEFT JOIN sys_catalog.sys_description d")); - assert!(sql.contains("WHERE n.nspname = 'app''s'")); - } - - #[test] - fn kingbase_available_extension_sql_supports_pg_catalog_fallback() { - let sql = super::kingbase_list_available_extensions_sql(super::KingbaseExtensionCatalog::Pg); - - assert!(sql.contains("FROM pg_catalog.pg_available_extensions")); - assert!(sql.contains("WHERE installed_version IS NULL")); - } - - #[test] - fn extension_infos_from_query_result_maps_installed_extensions() { - let result = db::QueryResult { - columns: vec![ - "extname".to_string(), - "extversion".to_string(), - "description".to_string(), - "nspname".to_string(), - ], - column_types: Vec::new(), - column_sortables: Vec::new(), - rows: vec![vec![ - serde_json::json!("kdb_utils"), - serde_json::json!("1.0"), - serde_json::json!("KingBase utilities"), - serde_json::json!("public"), - ]], - affected_rows: 0, - execution_time_ms: 0, - truncated: false, - session_id: None, - has_more: false, - }; - - let extensions = super::extension_infos_from_query_result(result, true); - - assert_eq!(extensions.len(), 1); - assert_eq!(extensions[0].name, "kdb_utils"); - assert_eq!(extensions[0].version, "1.0"); - assert_eq!(extensions[0].comment.as_deref(), Some("KingBase utilities")); - assert_eq!(extensions[0].schema.as_deref(), Some("public")); - } - #[test] fn agent_metadata_timeout_defaults_to_sixty_seconds_and_honors_longer_config() { assert_eq!(super::agent_metadata_timeout(None), Some(std::time::Duration::from_secs(60))); @@ -3706,13 +3643,7 @@ mod tests { } #[test] - fn agent_database_object_statistics_sql_uses_native_catalogs() { - let kingbase_sql = kingbase_object_statistics_sql("core's"); - assert!(kingbase_sql.contains("sys_catalog.sys_class")); - assert!(kingbase_sql.contains("sys_catalog.sys_namespace")); - assert!(kingbase_sql.contains("sys_total_relation_size")); - assert!(kingbase_sql.contains("n.nspname = 'core''s'")); - + fn gbase8a_object_statistics_sql_uses_information_schema() { let gbase_sql = gbase8a_object_statistics_sql("shop's"); assert!(gbase_sql.contains("information_schema.TABLES")); assert!(gbase_sql.contains("DATA_LENGTH")); @@ -4201,7 +4132,7 @@ async fn list_object_statistics_once( .await; } if db_config.as_ref().is_some_and(|config| config.db_type == DatabaseType::Kingbase) { - let sql = kingbase_object_statistics_sql(schema); + let sql = kingbase::object_statistics_sql(schema); drop(connections); return agent_list_object_statistics( client, @@ -5186,13 +5117,8 @@ pub async fn list_extensions_core( let connections = state.connections.read().await; if let Some(client) = extract_pool!(&connections, &pool_key, Agent) { drop(connections); - return kingbase_agent_list_extensions( - client, - database, - schema, - agent_metadata_timeout(db_config.as_ref()), - ) - .await; + return kingbase::list_extensions(client, database, schema, agent_metadata_timeout(db_config.as_ref())) + .await; } } @@ -5219,7 +5145,7 @@ pub async fn list_available_extensions_core( let connections = state.connections.read().await; if let Some(client) = extract_pool!(&connections, &pool_key, Agent) { drop(connections); - return kingbase_agent_list_available_extensions( + return kingbase::list_available_extensions( client, database, agent_metadata_timeout(db_config.as_ref()), @@ -5239,140 +5165,6 @@ pub async fn list_available_extensions_core( .await } -#[derive(Clone, Copy)] -enum KingbaseExtensionCatalog { - Sys, - Pg, -} - -impl KingbaseExtensionCatalog { - fn catalog_name(self) -> &'static str { - match self { - Self::Sys => "sys_catalog", - Self::Pg => "pg_catalog", - } - } - - fn prefix(self) -> &'static str { - match self { - Self::Sys => "sys", - Self::Pg => "pg", - } - } -} - -fn kingbase_list_extensions_sql(schema: Option<&str>, catalog: KingbaseExtensionCatalog) -> String { - let catalog_name = catalog.catalog_name(); - let prefix = catalog.prefix(); - let schema_filter = schema - .map(str::trim) - .filter(|value| !value.is_empty()) - .map(|schema| format!("WHERE n.nspname = {}", sql_string(schema))) - .unwrap_or_default(); - let order_by = if schema_filter.is_empty() { "n.nspname, e.extname" } else { "e.extname" }; - format!( - "SELECT e.extname, COALESCE(e.extversion, '') AS extversion, d.description, n.nspname \ - FROM {catalog_name}.{prefix}_extension e \ - JOIN {catalog_name}.{prefix}_namespace n ON n.oid = e.extnamespace \ - LEFT JOIN {catalog_name}.{prefix}_description d ON d.objoid = e.oid AND d.objsubid = 0 \ - {schema_filter} \ - ORDER BY {order_by}" - ) -} - -fn kingbase_list_available_extensions_sql(catalog: KingbaseExtensionCatalog) -> String { - let catalog_name = catalog.catalog_name(); - let prefix = catalog.prefix(); - format!( - "SELECT name, default_version, comment \ - FROM {catalog_name}.{prefix}_available_extensions \ - WHERE installed_version IS NULL \ - ORDER BY name" - ) -} - -async fn kingbase_agent_query_result( - client: Arc>, - database: &str, - sql: &str, - max_rows: usize, - timeout_duration: Option, -) -> Result { - let params = agent_execute_query_params( - sql, - if database.is_empty() { None } else { Some(database) }, - None, - QueryExecutionOptions { max_rows: Some(max_rows), ..Default::default() }, - ); - let mut client = client.lock().await; - client.execute_query_with_timeout(params, timeout_duration).await -} - -async fn kingbase_agent_query_result_with_catalog_fallback( - client: Arc>, - database: &str, - sys_sql: String, - pg_sql: String, - max_rows: usize, - timeout_duration: Option, -) -> Result { - match kingbase_agent_query_result(client.clone(), database, &sys_sql, max_rows, timeout_duration).await { - Ok(result) => Ok(result), - Err(sys_error) => kingbase_agent_query_result(client, database, &pg_sql, max_rows, timeout_duration) - .await - .map_err(|pg_error| format!("{sys_error}; pg_catalog fallback failed: {pg_error}")), - } -} - -async fn kingbase_agent_list_extensions( - client: Arc>, - database: &str, - schema: Option<&str>, - timeout_duration: Option, -) -> Result, String> { - let result = kingbase_agent_query_result_with_catalog_fallback( - client, - database, - kingbase_list_extensions_sql(schema, KingbaseExtensionCatalog::Sys), - kingbase_list_extensions_sql(schema, KingbaseExtensionCatalog::Pg), - 10_000, - timeout_duration, - ) - .await?; - Ok(extension_infos_from_query_result(result, true)) -} - -async fn kingbase_agent_list_available_extensions( - client: Arc>, - database: &str, - timeout_duration: Option, -) -> Result, String> { - let result = kingbase_agent_query_result_with_catalog_fallback( - client, - database, - kingbase_list_available_extensions_sql(KingbaseExtensionCatalog::Sys), - kingbase_list_available_extensions_sql(KingbaseExtensionCatalog::Pg), - 10_000, - timeout_duration, - ) - .await?; - Ok(extension_infos_from_query_result(result, false)) -} - -fn extension_infos_from_query_result(result: db::QueryResult, include_schema: bool) -> Vec { - result - .rows - .into_iter() - .filter_map(|row| { - let name = query_result_cell_string(&row, 0)?; - let version = query_result_cell_string(&row, 1).unwrap_or_default(); - let comment = query_result_cell_string(&row, 2).filter(|value| !value.trim().is_empty()); - let schema = include_schema.then(|| query_result_cell_string(&row, 3)).flatten(); - Some(db::ExtensionInfo { name, version, comment, schema }) - }) - .collect() -} - pub async fn list_owners_core( state: &AppState, connection_id: &str, diff --git a/crates/dbx-core/src/schema/kingbase.rs b/crates/dbx-core/src/schema/kingbase.rs new file mode 100644 index 0000000000..9433fc1d31 --- /dev/null +++ b/crates/dbx-core/src/schema/kingbase.rs @@ -0,0 +1,219 @@ +use crate::db; +use crate::query::{agent_execute_query_params, QueryExecutionOptions}; +use std::sync::Arc; +use std::time::Duration; + +use super::{query_result_cell_string, sql_string}; + +#[derive(Clone, Copy)] +enum ExtensionCatalog { + Sys, + Pg, +} + +impl ExtensionCatalog { + fn catalog_name(self) -> &'static str { + match self { + Self::Sys => "sys_catalog", + Self::Pg => "pg_catalog", + } + } + + fn prefix(self) -> &'static str { + match self { + Self::Sys => "sys", + Self::Pg => "pg", + } + } +} + +pub(super) fn object_statistics_sql(schema: &str) -> String { + format!( + "SELECT c.relname, n.nspname, \ + CAST(CASE WHEN c.reltuples < 0 THEN 0 ELSE c.reltuples END AS BIGINT) AS estimated_rows, \ + CAST(sys_total_relation_size(c.oid) AS BIGINT) AS total_bytes \ + FROM sys_catalog.sys_class c \ + JOIN sys_catalog.sys_namespace n ON n.oid = c.relnamespace \ + WHERE n.nspname = {} AND c.relkind IN ('r','m','f','p') \ + ORDER BY c.relname", + sql_string(schema), + ) +} + +fn list_extensions_sql(schema: Option<&str>, catalog: ExtensionCatalog) -> String { + let catalog_name = catalog.catalog_name(); + let prefix = catalog.prefix(); + let schema_filter = schema + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(|schema| format!("WHERE n.nspname = {}", sql_string(schema))) + .unwrap_or_default(); + let order_by = if schema_filter.is_empty() { "n.nspname, e.extname" } else { "e.extname" }; + format!( + "SELECT e.extname, COALESCE(e.extversion, '') AS extversion, d.description, n.nspname \ + FROM {catalog_name}.{prefix}_extension e \ + JOIN {catalog_name}.{prefix}_namespace n ON n.oid = e.extnamespace \ + LEFT JOIN {catalog_name}.{prefix}_description d ON d.objoid = e.oid AND d.objsubid = 0 \ + {schema_filter} \ + ORDER BY {order_by}" + ) +} + +fn list_available_extensions_sql(catalog: ExtensionCatalog) -> String { + let catalog_name = catalog.catalog_name(); + let prefix = catalog.prefix(); + format!( + "SELECT name, default_version, comment \ + FROM {catalog_name}.{prefix}_available_extensions \ + WHERE installed_version IS NULL \ + ORDER BY name" + ) +} + +async fn query_result( + client: Arc>, + database: &str, + sql: &str, + max_rows: usize, + timeout_duration: Option, +) -> Result { + let params = agent_execute_query_params( + sql, + if database.is_empty() { None } else { Some(database) }, + None, + QueryExecutionOptions { max_rows: Some(max_rows), ..Default::default() }, + ); + let mut client = client.lock().await; + client.execute_query_with_timeout(params, timeout_duration).await +} + +async fn query_result_with_catalog_fallback( + client: Arc>, + database: &str, + sys_sql: String, + pg_sql: String, + max_rows: usize, + timeout_duration: Option, +) -> Result { + match query_result(client.clone(), database, &sys_sql, max_rows, timeout_duration).await { + Ok(result) => Ok(result), + Err(sys_error) => query_result(client, database, &pg_sql, max_rows, timeout_duration) + .await + .map_err(|pg_error| format!("{sys_error}; pg_catalog fallback failed: {pg_error}")), + } +} + +pub(super) async fn list_extensions( + client: Arc>, + database: &str, + schema: Option<&str>, + timeout_duration: Option, +) -> Result, String> { + let result = query_result_with_catalog_fallback( + client, + database, + list_extensions_sql(schema, ExtensionCatalog::Sys), + list_extensions_sql(schema, ExtensionCatalog::Pg), + 10_000, + timeout_duration, + ) + .await?; + Ok(extension_infos_from_query_result(result, true)) +} + +pub(super) async fn list_available_extensions( + client: Arc>, + database: &str, + timeout_duration: Option, +) -> Result, String> { + let result = query_result_with_catalog_fallback( + client, + database, + list_available_extensions_sql(ExtensionCatalog::Sys), + list_available_extensions_sql(ExtensionCatalog::Pg), + 10_000, + timeout_duration, + ) + .await?; + Ok(extension_infos_from_query_result(result, false)) +} + +fn extension_infos_from_query_result(result: db::QueryResult, include_schema: bool) -> Vec { + result + .rows + .into_iter() + .filter_map(|row| { + let name = query_result_cell_string(&row, 0)?; + let version = query_result_cell_string(&row, 1).unwrap_or_default(); + let comment = query_result_cell_string(&row, 2).filter(|value| !value.trim().is_empty()); + let schema = include_schema.then(|| query_result_cell_string(&row, 3)).flatten(); + Some(db::ExtensionInfo { name, version, comment, schema }) + }) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn object_statistics_sql_uses_native_catalogs() { + let sql = object_statistics_sql("core's"); + + assert!(sql.contains("sys_catalog.sys_class")); + assert!(sql.contains("sys_catalog.sys_namespace")); + assert!(sql.contains("sys_total_relation_size")); + assert!(sql.contains("n.nspname = 'core''s'")); + } + + #[test] + fn extension_sql_uses_sys_catalog_and_escapes_schema() { + let sql = list_extensions_sql(Some("app's"), ExtensionCatalog::Sys); + + assert!(sql.contains("FROM sys_catalog.sys_extension e")); + assert!(sql.contains("JOIN sys_catalog.sys_namespace n")); + assert!(sql.contains("LEFT JOIN sys_catalog.sys_description d")); + assert!(sql.contains("WHERE n.nspname = 'app''s'")); + } + + #[test] + fn available_extension_sql_supports_pg_catalog_fallback() { + let sql = list_available_extensions_sql(ExtensionCatalog::Pg); + + assert!(sql.contains("FROM pg_catalog.pg_available_extensions")); + assert!(sql.contains("WHERE installed_version IS NULL")); + } + + #[test] + fn extension_infos_map_installed_extensions() { + let result = db::QueryResult { + columns: vec![ + "extname".to_string(), + "extversion".to_string(), + "description".to_string(), + "nspname".to_string(), + ], + column_types: Vec::new(), + column_sortables: Vec::new(), + rows: vec![vec![ + serde_json::json!("kdb_utils"), + serde_json::json!("1.0"), + serde_json::json!("KingBase utilities"), + serde_json::json!("public"), + ]], + affected_rows: 0, + execution_time_ms: 0, + truncated: false, + session_id: None, + has_more: false, + }; + + let extensions = extension_infos_from_query_result(result, true); + + assert_eq!(extensions.len(), 1); + assert_eq!(extensions[0].name, "kdb_utils"); + assert_eq!(extensions[0].version, "1.0"); + assert_eq!(extensions[0].comment.as_deref(), Some("KingBase utilities")); + assert_eq!(extensions[0].schema.as_deref(), Some("public")); + } +}