fix(vastbase): tolerate metadata row type mismatches

This commit is contained in:
gggaiitx
2026-07-08 14:57:14 +08:00
committed by GitHub
parent 3a35f5b911
commit 651c7b8db3
3 changed files with 116 additions and 67 deletions
+114 -67
View File
@@ -1405,7 +1405,7 @@ pub async fn list_databases(pool: &Pool) -> Result<Vec<DatabaseInfo>, String> {
.await
.map_err(|e| e.to_string())?;
Ok(rows.iter().map(|row| DatabaseInfo { name: row.get::<_, String>(0) }).collect())
Ok(rows.iter().map(|row| DatabaseInfo { name: pg_row_try_string(row, 0) }).collect())
}
pub async fn list_tables(pool: &Pool, schema: &str) -> Result<Vec<TableInfo>, String> {
@@ -1438,8 +1438,8 @@ pub async fn list_tables_filtered(
Ok(rows
.iter()
.map(|row| TableInfo {
name: row.get::<_, String>(0),
table_type: row.get::<_, String>(1),
name: pg_row_try_string(row, 0),
table_type: pg_row_try_string(row, 1),
comment: row.try_get::<_, Option<String>>(2).ok().flatten().filter(|s| !s.is_empty()),
parent_schema: row.try_get::<_, Option<String>>(3).ok().flatten().filter(|s| !s.is_empty()),
parent_name: row.try_get::<_, Option<String>>(4).ok().flatten().filter(|s| !s.is_empty()),
@@ -1474,7 +1474,7 @@ pub async fn completion_assistant_search(
.await
.map_err(|e| e.to_string())?
{
let schema_name: String = row.get(0);
let schema_name: String = pg_row_try_string(&row, 0);
candidates.push(CompletionAssistantCandidate {
name: schema_name.clone(),
kind: CompletionAssistantCandidateKind::Schema,
@@ -1498,16 +1498,16 @@ pub async fn completion_assistant_search(
.await
.map_err(|e| e.to_string())?;
for row in rows {
let table_type: String = row.get(2);
let table_type: String = pg_row_try_string(&row, 2);
candidates.push(CompletionAssistantCandidate {
name: row.get(0),
name: pg_row_try_string(&row, 0),
kind: if table_type.contains("VIEW") {
CompletionAssistantCandidateKind::View
} else {
CompletionAssistantCandidateKind::Table
},
database: Some(request.database.clone()),
schema: Some(row.get(1)),
schema: Some(pg_row_try_string(&row, 1)),
parent_schema: row.try_get::<_, Option<String>>(4).ok().flatten(),
parent_name: row.try_get::<_, Option<String>>(5).ok().flatten(),
comment: row.try_get::<_, Option<String>>(3).ok().flatten(),
@@ -1526,16 +1526,16 @@ pub async fn completion_assistant_search(
.await
.map_err(|e| e.to_string())?;
for row in rows {
let routine_type: String = row.get(2);
let routine_type: String = pg_row_try_string(&row, 2);
candidates.push(CompletionAssistantCandidate {
name: row.get(0),
name: pg_row_try_string(&row, 0),
kind: if routine_type == "PROCEDURE" {
CompletionAssistantCandidateKind::Procedure
} else {
CompletionAssistantCandidateKind::Function
},
database: Some(request.database.clone()),
schema: Some(row.get(1)),
schema: Some(pg_row_try_string(&row, 1)),
parent_schema: None,
parent_name: None,
comment: row.try_get::<_, Option<String>>(3).ok().flatten(),
@@ -1556,14 +1556,14 @@ pub async fn completion_assistant_search(
.map_err(|e| e.to_string())?;
for row in rows {
candidates.push(CompletionAssistantCandidate {
name: row.get(0),
name: pg_row_try_string(&row, 0),
kind: CompletionAssistantCandidateKind::Column,
database: Some(request.database.clone()),
schema: Some(schema.to_string()),
parent_schema: Some(schema.to_string()),
parent_name: Some(table.to_string()),
comment: row.try_get::<_, Option<String>>(2).ok().flatten(),
data_type: Some(row.get(1)),
data_type: Some(pg_row_try_string(&row, 1)),
});
}
}
@@ -1911,7 +1911,7 @@ fn postgres_proc_has_prokind_sql() -> &'static str {
async fn postgres_proc_has_prokind(client: &deadpool_postgres::Client) -> Result<bool, String> {
let row =
postgres_query_one_cached(client, postgres_proc_has_prokind_sql(), &[]).await.map_err(|e| e.to_string())?;
Ok(row.get(0))
Ok(pg_row_try_bool(&row, 0).unwrap_or(false))
}
fn postgres_proc_has_prosp_sql() -> &'static str {
@@ -1926,7 +1926,7 @@ fn postgres_proc_has_prosp_sql() -> &'static str {
async fn postgres_proc_has_prosp(client: &deadpool_postgres::Client) -> Result<bool, String> {
let row = postgres_query_one_cached(client, postgres_proc_has_prosp_sql(), &[]).await.map_err(|e| e.to_string())?;
Ok(row.get(0))
Ok(pg_row_try_bool(&row, 0).unwrap_or(false))
}
async fn list_objects_rows(
@@ -1962,8 +1962,8 @@ pub async fn list_objects(pool: &Pool, schema: &str) -> Result<Vec<ObjectInfo>,
Ok(rows
.iter()
.map(|row| ObjectInfo {
name: row.get::<_, String>(0),
object_type: row.get::<_, String>(1),
name: pg_row_try_string(row, 0),
object_type: pg_row_try_string(row, 1),
schema: Some(schema.to_string()),
comment: row.try_get::<_, Option<String>>(2).ok().flatten().filter(|s| !s.is_empty()),
created_at: row.try_get::<_, Option<String>>(3).ok().flatten().filter(|s| !s.is_empty()),
@@ -1994,7 +1994,7 @@ pub async fn list_object_statistics(pool: &Pool, schema: &str) -> Result<Vec<Obj
Ok(rows
.iter()
.map(|row| ObjectStatistics {
name: row.get::<_, String>(0),
name: pg_row_try_string(row, 0),
schema: Some(schema.to_string()),
estimated_rows: row.try_get::<_, i64>(1).ok(),
total_bytes: row.try_get::<_, i64>(2).ok(),
@@ -2028,7 +2028,7 @@ pub async fn list_schema_infos(pool: &Pool) -> Result<Vec<SchemaInfo>, String> {
Ok(rows
.iter()
.map(|row| SchemaInfo {
name: row.get::<_, String>(0),
name: pg_row_try_string(row, 0),
comment: row.try_get::<_, Option<String>>(1).ok().flatten(),
})
.collect())
@@ -2124,23 +2124,64 @@ const POSTGRES_COLUMNS_INFORMATION_SCHEMA_SQL: &str = "SELECT c.column_name, \
WHERE c.table_schema = $1 AND c.table_name = $2 \
ORDER BY c.ordinal_position";
fn column_info_from_row(row: &Row) -> Result<ColumnInfo, tokio_postgres::Error> {
/// Read a boolean column from a PostgreSQL row, tolerating databases that
/// encode booleans as integers (0/1) or text ('t'/'f') instead of the standard
/// `bool` OID. Returns `None` when the column is NULL or truly unreadable.
fn pg_row_try_bool(row: &Row, idx: usize) -> Option<bool> {
if let Ok(v) = row.try_get::<_, bool>(idx) {
return Some(v);
}
if let Ok(v) = row.try_get::<_, i32>(idx) {
return Some(v != 0);
}
if let Ok(v) = row.try_get::<_, i16>(idx) {
return Some(v != 0);
}
if let Ok(Some(v)) = row.try_get::<_, Option<String>>(idx) {
match v.as_str() {
"t" | "true" | "1" | "yes" | "YES" => return Some(true),
"f" | "false" | "0" | "no" | "NO" => return Some(false),
_ => return None,
}
}
None
}
/// Read a String column from a PostgreSQL row, tolerating databases that
/// return text as other types. Falls back to i64/i32/i16/bool formatting.
fn pg_row_try_string(row: &Row, idx: usize) -> String {
if let Ok(v) = row.try_get::<_, String>(idx) {
return v;
}
if let Ok(v) = row.try_get::<_, i64>(idx) {
return v.to_string();
}
if let Ok(v) = row.try_get::<_, i32>(idx) {
return v.to_string();
}
if let Ok(v) = row.try_get::<_, i16>(idx) {
return v.to_string();
}
if let Some(v) = pg_row_try_bool(row, idx) {
return v.to_string();
}
String::new()
}
fn column_info_from_row(row: &Row) -> ColumnInfo {
let full_type = row.try_get::<_, Option<String>>(1).ok().flatten().unwrap_or_default();
Ok(ColumnInfo {
// Postgres-compatible servers may expose catalog expression columns
// with unexpected type OIDs. Use try_get so metadata decoding can
// fall back instead of panicking under release panic=abort.
name: row.try_get::<_, String>(0)?,
ColumnInfo {
name: pg_row_try_string(row, 0),
data_type: full_type,
is_nullable: row.try_get::<_, bool>(2)?,
is_nullable: pg_row_try_bool(row, 2).unwrap_or(true),
column_default: row.try_get::<_, Option<String>>(3).ok().flatten(),
is_primary_key: row.try_get::<_, bool>(4)?,
is_primary_key: pg_row_try_bool(row, 4).unwrap_or(false),
extra: row.try_get::<_, Option<String>>(6).ok().flatten(),
comment: row.try_get::<_, Option<String>>(5).ok().flatten(),
numeric_precision: row.try_get::<_, Option<i32>>(7).ok().flatten(),
numeric_scale: row.try_get::<_, Option<i32>>(8).ok().flatten(),
character_maximum_length: row.try_get::<_, Option<i32>>(9).ok().flatten(),
})
}
}
async fn get_columns_with_sql(
@@ -2151,7 +2192,7 @@ async fn get_columns_with_sql(
) -> Result<Vec<ColumnInfo>, tokio_postgres::Error> {
let rows = postgres_query_cached(client, sql, &[&schema, &table]).await?;
rows.iter().map(column_info_from_row).collect()
Ok(rows.iter().map(column_info_from_row).collect())
}
pub async fn get_columns(pool: &Pool, schema: &str, table: &str) -> Result<Vec<ColumnInfo>, String> {
@@ -2750,16 +2791,16 @@ async fn list_indexes_with_sql(
Ok(rows
.iter()
.map(|row| {
let all_cols: Vec<String> = row.get::<_, Vec<String>>(1);
let all_cols: Vec<String> = row.try_get::<_, Vec<String>>(1).unwrap_or_default();
let nkeyatts = row.try_get::<_, Option<i16>>(6).ok().flatten().unwrap_or(all_cols.len() as i16) as usize;
let split_at = nkeyatts.min(all_cols.len());
let key_cols = all_cols[..split_at].to_vec();
let included = if split_at < all_cols.len() { all_cols[split_at..].to_vec() } else { vec![] };
IndexInfo {
name: row.get::<_, String>(0),
name: pg_row_try_string(row, 0),
columns: key_cols,
is_unique: row.get::<_, bool>(2),
is_primary: row.get::<_, bool>(3),
is_unique: pg_row_try_bool(row, 2).unwrap_or(false),
is_primary: pg_row_try_bool(row, 3).unwrap_or(false),
filter: row.try_get::<_, Option<String>>(4).ok().flatten(),
index_type: row.try_get::<_, Option<String>>(5).ok().flatten(),
included_columns: if included.is_empty() { None } else { Some(included) },
@@ -2829,13 +2870,13 @@ pub async fn list_foreign_keys(pool: &Pool, schema: &str, table: &str) -> Result
Ok(rows
.iter()
.map(|row| ForeignKeyInfo {
name: row.get::<_, String>(0),
column: row.get::<_, String>(1),
ref_schema: Some(row.get::<_, String>(2)),
ref_table: row.get::<_, String>(3),
ref_column: row.get::<_, String>(4),
on_update: postgres_foreign_key_action(row.get::<_, String>(5)),
on_delete: postgres_foreign_key_action(row.get::<_, String>(6)),
name: pg_row_try_string(row, 0),
column: pg_row_try_string(row, 1),
ref_schema: Some(pg_row_try_string(row, 2)),
ref_table: pg_row_try_string(row, 3),
ref_column: pg_row_try_string(row, 4),
on_update: postgres_foreign_key_action(pg_row_try_string(row, 5)),
on_delete: postgres_foreign_key_action(pg_row_try_string(row, 6)),
})
.collect())
}
@@ -2856,9 +2897,9 @@ pub async fn list_triggers(pool: &Pool, schema: &str, table: &str) -> Result<Vec
Ok(rows
.iter()
.map(|row| TriggerInfo {
name: row.get::<_, String>(0),
event: row.get::<_, String>(1),
timing: row.get::<_, String>(2),
name: pg_row_try_string(row, 0),
event: pg_row_try_string(row, 1),
timing: pg_row_try_string(row, 2),
statement: None,
})
.collect())
@@ -2903,7 +2944,7 @@ pub async fn list_functions(pool: &Pool, schema: &str) -> Result<Vec<FunctionInf
Ok(rows
.iter()
.map(|row| {
let def: String = row.get::<_, String>(3);
let def: String = pg_row_try_string(row, 3);
// Remove schema qualification from CREATE FUNCTION statement
// to avoid false differences when comparing across schemas.
// Handle both "schema.name" and schema.name formats.
@@ -2911,11 +2952,11 @@ pub async fn list_functions(pool: &Pool, schema: &str) -> Result<Vec<FunctionInf
.replace(&format!("CREATE OR REPLACE FUNCTION \"{}\".", schema), "CREATE OR REPLACE FUNCTION ")
.replace(&format!("CREATE OR REPLACE FUNCTION {}.", schema), "CREATE OR REPLACE FUNCTION ");
FunctionInfo {
name: row.get::<_, String>(0),
function_type: row.get::<_, String>(1),
data_type: row.get::<_, String>(2),
name: pg_row_try_string(row, 0),
function_type: pg_row_try_string(row, 1),
data_type: pg_row_try_string(row, 2),
definition: normalized_def,
arguments: row.get::<_, String>(4),
arguments: pg_row_try_string(row, 4),
}
})
.collect())
@@ -2947,13 +2988,13 @@ pub async fn list_sequences(pool: &Pool, schema: &str, with_last_values: bool) -
let mut sequences: Vec<SequenceInfo> = rows
.iter()
.map(|row| SequenceInfo {
name: row.get::<_, String>(0),
data_type: row.get::<_, String>(1),
start_value: row.get::<_, String>(2),
min_value: row.get::<_, String>(3),
max_value: row.get::<_, String>(4),
increment: row.get::<_, String>(5),
cycle: row.get::<_, String>(6) == "YES",
name: pg_row_try_string(row, 0),
data_type: pg_row_try_string(row, 1),
start_value: pg_row_try_string(row, 2),
min_value: pg_row_try_string(row, 3),
max_value: pg_row_try_string(row, 4),
increment: pg_row_try_string(row, 5),
cycle: pg_row_try_string(row, 6) == "YES",
last_value: None,
})
.collect();
@@ -2966,7 +3007,7 @@ pub async fn list_sequences(pool: &Pool, schema: &str, with_last_values: bool) -
WHERE c.relkind = 'S' AND n.nspname = $1";
if let Ok(rows) = postgres_query_cached(&client, sql, &[&schema]).await {
for row in rows {
let name: String = row.get(0);
let name: String = pg_row_try_string(&row, 0);
if let Ok(val) = row.try_get::<_, i64>(1) {
if let Some(seq) = sequences.iter_mut().find(|s| s.name == name) {
seq.last_value = Some(val.to_string());
@@ -2995,9 +3036,9 @@ pub async fn list_rules(pool: &Pool, schema: &str) -> Result<Vec<RuleInfo>, Stri
Ok(rows
.iter()
.map(|row| RuleInfo {
name: row.get::<_, String>(2),
table_name: row.get::<_, String>(1),
definition: row.get::<_, String>(3),
name: pg_row_try_string(row, 2),
table_name: pg_row_try_string(row, 1),
definition: pg_row_try_string(row, 3),
})
.collect())
}
@@ -3020,8 +3061,8 @@ pub async fn list_extensions(pool: &Pool, schema: &str) -> Result<Vec<ExtensionI
Ok(rows
.iter()
.map(|row| ExtensionInfo {
name: row.get::<_, String>(0),
version: row.get::<_, String>(1),
name: pg_row_try_string(row, 0),
version: pg_row_try_string(row, 1),
comment: row.try_get::<_, Option<String>>(2).ok().flatten().filter(|s| !s.is_empty()),
schema: Some(schema.to_string()),
})
@@ -3044,8 +3085,8 @@ pub async fn list_available_extensions(pool: &Pool) -> Result<Vec<ExtensionInfo>
Ok(rows
.iter()
.map(|row| ExtensionInfo {
name: row.get::<_, String>(0),
version: row.get::<_, String>(1),
name: pg_row_try_string(row, 0),
version: pg_row_try_string(row, 1),
comment: row.try_get::<_, Option<String>>(2).ok().flatten().filter(|s| !s.is_empty()),
schema: None,
})
@@ -3059,11 +3100,11 @@ pub async fn list_owners(pool: &Pool, schema: &str) -> Result<Vec<OwnerInfo>, St
Ok(rows
.iter()
.map(|row| {
let relkind: String = row.get(2);
let relkind: String = pg_row_try_string(row, 2);
OwnerInfo {
object_name: row.get::<_, String>(1),
object_name: pg_row_try_string(row, 1),
object_type: postgres_owner_object_type(&relkind).to_string(),
owner: row.get::<_, String>(3),
owner: pg_row_try_string(row, 3),
}
})
.collect())
@@ -3635,7 +3676,7 @@ mod tests {
#[tokio::test]
#[ignore = "requires DBX_TEST_POSTGRES_URL pointing at a writable PostgreSQL database"]
async fn postgres_column_metadata_decode_type_mismatch_returns_error() {
async fn postgres_column_metadata_decode_type_mismatch_uses_fallbacks() {
let url = std::env::var("DBX_TEST_POSTGRES_URL").expect("DBX_TEST_POSTGRES_URL");
let pool = connect(&url, std::time::Duration::from_secs(5)).await.expect("connect postgres");
let client =
@@ -3658,7 +3699,13 @@ mod tests {
.await
.expect("query mismatched metadata row");
assert!(column_info_from_row(&row).is_err());
let info = column_info_from_row(&row);
// int4 column_name should be converted to string "1" instead of panicking
assert_eq!(info.name, "1");
// text 'YES' is not a standard bool, pg_row_try_bool falls back to string match
assert!(info.is_nullable);
// int4 1 should be interpreted as true for is_primary_key
assert!(info.is_primary_key);
}
#[test]
+1
View File
@@ -303,6 +303,7 @@ mod tests {
external_config: None,
jdbc_driver_class: None,
jdbc_driver_paths: Vec::new(),
agent_java_options: Vec::new(),
one_time: false,
read_only: false,
}
+1
View File
@@ -201,6 +201,7 @@ mod tests {
external_config: None,
jdbc_driver_class: None,
jdbc_driver_paths: Vec::new(),
agent_java_options: Vec::new(),
one_time: false,
read_only: false,
}