diff --git a/crates/dbx-core/src/db/mysql.rs b/crates/dbx-core/src/db/mysql.rs index 1cecb21b8..766b91bee 100644 --- a/crates/dbx-core/src/db/mysql.rs +++ b/crates/dbx-core/src/db/mysql.rs @@ -335,7 +335,7 @@ pub async fn list_objects(pool: &MySqlPool, database: &str) -> Result String { format!( "SELECT c.COLUMN_NAME, c.COLUMN_TYPE, c.IS_NULLABLE, c.COLUMN_DEFAULT, c.EXTRA, c.COLUMN_COMMENT, \ - c.NUMERIC_PRECISION, c.NUMERIC_SCALE, c.CHARACTER_MAXIMUM_LENGTH \ + c.COLUMN_KEY, c.NUMERIC_PRECISION, c.NUMERIC_SCALE, c.CHARACTER_MAXIMUM_LENGTH \ FROM information_schema.COLUMNS c \ WHERE c.TABLE_SCHEMA = {} AND c.TABLE_NAME = {} \ ORDER BY c.ORDINAL_POSITION", @@ -355,6 +355,10 @@ fn primary_key_columns_sql(database: &str, table: &str) -> String { ) } +fn is_primary_key_column(primary_key_columns: &HashSet, name: &str, column_key: &str) -> bool { + primary_key_columns.contains(name) || column_key.eq_ignore_ascii_case("PRI") +} + pub async fn get_columns(pool: &MySqlPool, database: &str, table: &str) -> Result, String> { let pk_sql = primary_key_columns_sql(database, table); let pk_rows: Vec = sqlx::raw_sql(&pk_sql).fetch_all(pool).await.map_err(|e| e.to_string())?; @@ -367,8 +371,9 @@ pub async fn get_columns(pool: &MySqlPool, database: &str, table: &str) -> Resul .iter() .map(|row| { let name = get_str_by_name(row, "COLUMN_NAME"); + let column_key = get_str_by_name(row, "COLUMN_KEY"); ColumnInfo { - is_primary_key: primary_key_columns.contains(&name), + is_primary_key: is_primary_key_column(&primary_key_columns, &name, &column_key), name, data_type: get_str_by_name(row, "COLUMN_TYPE"), is_nullable: get_str_by_name(row, "IS_NULLABLE") == "YES", @@ -666,6 +671,20 @@ mod tests { assert!(sql.contains("CONSTRAINT_NAME = 'PRIMARY'")); } + #[test] + fn mysql_columns_sql_selects_column_key_for_starrocks_primary_fallback() { + let sql = columns_sql("app", "users"); + + assert!(sql.contains("c.COLUMN_KEY")); + } + + #[test] + fn mysql_column_key_marks_primary_when_key_column_usage_is_unavailable() { + let primary_key_columns = HashSet::new(); + + assert!(is_primary_key_column(&primary_key_columns, "id", "PRI")); + } + #[test] fn mysql_management_show_queries_use_text_protocol() { assert!(requires_text_protocol_query("SHOW PROCESSLIST"));