From 3ba617fc74963fa2ce9920bb990175f7d6e08e0d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=A4=A7=E5=B3=B0?= <67090734+tianyifeng-druid@users.noreply.github.com> Date: Sat, 4 Jul 2026 23:06:21 +0800 Subject: [PATCH] fix(db2): include comments in table ddl --- crates/dbx-core/src/schema.rs | 143 ++++++++++++++++++++++++++++++++++ 1 file changed, 143 insertions(+) diff --git a/crates/dbx-core/src/schema.rs b/crates/dbx-core/src/schema.rs index ed7dc1051a..4473d22b8e 100644 --- a/crates/dbx-core/src/schema.rs +++ b/crates/dbx-core/src/schema.rs @@ -3985,6 +3985,16 @@ pub async fn get_table_ddl_core( ) .await; } + if db_config.as_ref().is_some_and(|config| config.db_type == DatabaseType::Db2) { + return db2_agent_table_ddl( + client, + database, + schema, + table, + agent_metadata_timeout(db_config.as_ref()), + ) + .await; + } let mut client = client.lock().await; return client.get_table_ddl(database, schema, table, agent_metadata_timeout(db_config.as_ref())).await; } @@ -4583,6 +4593,139 @@ fn append_oracle_comments_to_ddl( result } +async fn db2_agent_table_ddl( + client: Arc>, + database: &str, + schema: &str, + table: &str, + timeout_duration: Option, +) -> Result { + let mut client = client.lock().await; + let ddl = client.get_table_ddl::(database, schema, table, timeout_duration).await?; + match append_db2_comments_to_ddl(&mut client, database, schema, table, &ddl, timeout_duration).await { + Ok(ddl) => Ok(ddl), + Err(error) => { + log::debug!( + "[schema][db2:get_table_ddl:comments-fallback-failed] schema={} table={} error={}", + schema, + table, + error + ); + Ok(ddl) + } + } +} + +async fn append_db2_comments_to_ddl( + client: &mut db::agent_driver::AgentDriverClient, + database: &str, + schema: &str, + table: &str, + ddl: &str, + timeout_duration: Option, +) -> Result { + let table_comment = db2_table_comment(client, database, schema, table, timeout_duration).await; + let column_comments = db2_column_comments(client, database, schema, table, timeout_duration).await; + let mut columns = + client.get_columns::>(database, schema, table, timeout_duration).await.unwrap_or_default(); + if !column_comments.is_empty() { + for column in &mut columns { + if column.comment.as_deref().is_none_or(|c| c.trim().is_empty() || c.trim().eq_ignore_ascii_case("null")) { + if let Some(remark) = column_comments.get(&column.name.to_uppercase()) { + column.comment = Some(remark.clone()); + } + } + } + } + Ok(append_oracle_comments_to_ddl(ddl, schema, table, table_comment.as_deref(), &columns)) +} + +async fn db2_table_comment( + client: &mut db::agent_driver::AgentDriverClient, + database: &str, + schema: &str, + table: &str, + timeout_duration: Option, +) -> Option { + // 优先使用原始值查询，支持 quoted/mixed-case 对象；如果查不到再 fallback 到大写 + for (schema_name, table_name) in + [(schema.trim(), table.trim()), (&schema.trim().to_uppercase(), &table.trim().to_uppercase())] + { + let schema_filter = if schema_name.is_empty() { "CURRENT SCHEMA".to_string() } else { sql_string(schema_name) }; + let sql = format!( + "SELECT REMARKS FROM SYSCAT.TABLES WHERE TABSCHEMA = {} AND TABNAME = {} AND REMARKS IS NOT NULL", + schema_filter, + sql_string(table_name), + ); + if let Ok(result) = client + .execute_query_with_timeout::( + agent_execute_query_params( + &sql, + if database.is_empty() { None } else { Some(database) }, + if schema.is_empty() { None } else { Some(schema) }, + QueryExecutionOptions { max_rows: Some(1), ..Default::default() }, + ), + timeout_duration, + ) + .await + { + if let Some(comment) = + result.rows.first().and_then(|row| row.first()).and_then(|v| v.as_str()).map(|s| s.to_string()) + { + return Some(comment); + } + } + } + None +} + +async fn db2_column_comments( + client: &mut db::agent_driver::AgentDriverClient, + database: &str, + schema: &str, + table: &str, + timeout_duration: Option, +) -> HashMap { + // 优先使用原始值查询，支持 quoted/mixed-case 对象；如果查不到再 fallback 到大写 + let mut comments = HashMap::new(); + for (schema_name, table_name) in + [(schema.trim(), table.trim()), (&schema.trim().to_uppercase(), &table.trim().to_uppercase())] + { + let schema_filter = if schema_name.is_empty() { "CURRENT SCHEMA".to_string() } else { sql_string(schema_name) }; + let sql = format!( + "SELECT COLNAME, REMARKS FROM SYSCAT.COLUMNS WHERE TABSCHEMA = {} AND TABNAME = {} AND REMARKS IS NOT NULL", + schema_filter, + sql_string(table_name), + ); + let result = match client + .execute_query_with_timeout::( + agent_execute_query_params( + &sql, + if database.is_empty() { None } else { Some(database) }, + if schema.is_empty() { None } else { Some(schema) }, + QueryExecutionOptions { ..Default::default() }, + ), + timeout_duration, + ) + .await + { + Ok(result) => result, + Err(_) => continue, + }; + for row in &result.rows { + let col_name = row.first().and_then(|v| v.as_str()).unwrap_or("").trim(); + let remark = row.get(1).and_then(|v| v.as_str()).unwrap_or("").trim(); + if !col_name.is_empty() && !remark.is_empty() { + comments.entry(col_name.to_uppercase()).or_insert_with(|| remark.to_string()); + } + } + if !comments.is_empty() { + break; + } + } + comments +} + async fn postgres_object_source( pool: &deadpool_postgres::Pool, schema: &str,