From 03bc87c77a4915b85807884fcd610c685c165e8e Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Mon, 13 Jul 2026 10:29:03 +0800 Subject: [PATCH] fix(postgres): fall back for routine signatures --- crates/dbx-core/src/db/postgres.rs | 119 ++++++++++++++++++++++------- 1 file changed, 92 insertions(+), 27 deletions(-) diff --git a/crates/dbx-core/src/db/postgres.rs b/crates/dbx-core/src/db/postgres.rs index 6c019eb36..a95603ab2 100644 --- a/crates/dbx-core/src/db/postgres.rs +++ b/crates/dbx-core/src/db/postgres.rs @@ -1921,12 +1921,42 @@ fn list_object_routines_sql(include_timestamps: bool, has_proc_prokind: bool, ha WHERE n.nspname = $1 AND NOT p.proisagg AND NOT p.proiswindow" } -fn list_objects_sql(include_timestamps: bool, has_proc_prokind: bool, has_proc_prosp: bool) -> String { - format!( +fn list_objects_sql( + include_timestamps: bool, + has_proc_prokind: bool, + has_proc_prosp: bool, + has_function_identity_arguments: bool, +) -> String { + let sql = format!( "{} UNION ALL {} ORDER BY sort_order, object_name", list_object_relations_sql(include_timestamps), list_object_routines_sql(include_timestamps, has_proc_prokind, has_proc_prosp) - ) + ); + if has_function_identity_arguments { + sql + } else { + // Redshift and older PostgreSQL-compatible servers may only expose the + // older formatter. It includes argument names but still distinguishes + // overloads instead of making the whole schema browser unavailable. + sql.replace("pg_get_function_identity_arguments(p.oid)", "pg_get_function_arguments(p.oid)") + } +} + +fn postgres_has_function_identity_arguments_sql() -> &'static str { + "SELECT EXISTS ( \ + SELECT 1 \ + FROM pg_catalog.pg_proc p \ + JOIN pg_catalog.pg_namespace n ON n.oid = p.pronamespace \ + WHERE n.nspname = 'pg_catalog' \ + AND p.proname = 'pg_get_function_identity_arguments' \ + )" +} + +async fn postgres_has_function_identity_arguments(client: &deadpool_postgres::Client) -> Result { + let row = postgres_query_one_cached(client, postgres_has_function_identity_arguments_sql(), &[]) + .await + .map_err(|e| e.to_string())?; + Ok(pg_row_try_bool(&row, 0).unwrap_or(false)) } fn postgres_proc_has_prokind_sql() -> &'static str { @@ -1966,8 +1996,9 @@ async fn list_objects_rows( include_timestamps: bool, has_proc_prokind: bool, has_proc_prosp: bool, + has_function_identity_arguments: bool, ) -> Result, String> { - let sql = list_objects_sql(include_timestamps, has_proc_prokind, has_proc_prosp); + let sql = list_objects_sql(include_timestamps, has_proc_prokind, has_proc_prosp, has_function_identity_arguments); postgres_query_cached(client, &sql, &[&schema]).await.map_err(|e| e.to_string()) } @@ -1977,11 +2008,30 @@ pub async fn list_objects(pool: &Pool, schema: &str) -> Result, // Some GaussDB-compatible catalogs expose prosp alongside, or instead of, // PostgreSQL 11's prokind. Treat prosp as an extra procedure signal. let has_proc_prosp = postgres_proc_has_prosp(&client).await?; - let rows = match list_objects_rows(&client, schema, true, has_proc_prokind, has_proc_prosp).await { + let has_function_identity_arguments = postgres_has_function_identity_arguments(&client).await?; + let rows = match list_objects_rows( + &client, + schema, + true, + has_proc_prokind, + has_proc_prosp, + has_function_identity_arguments, + ) + .await + { Ok(rows) => rows, Err(primary_error) => { log::debug!("[postgres][list_objects:timestamp-fallback] primary_error={}", primary_error); - match list_objects_rows(&client, schema, false, has_proc_prokind, has_proc_prosp).await { + match list_objects_rows( + &client, + schema, + false, + has_proc_prokind, + has_proc_prosp, + has_function_identity_arguments, + ) + .await + { Ok(rows) => rows, Err(fallback_error) => { return Err(format!("{primary_error}; timestamp fallback failed: {fallback_error}")); @@ -3892,7 +3942,7 @@ mod tests { #[test] fn list_objects_sql_includes_routines() { - let sql = list_objects_sql(true, true, false); + let sql = list_objects_sql(true, true, false, true); assert!(sql.contains("pg_catalog.pg_class")); assert!(sql.contains("pg_catalog.pg_proc")); assert!(sql.contains("pg_catalog.pg_inherits")); @@ -3909,39 +3959,54 @@ mod tests { #[test] fn list_objects_sql_without_timestamps_omits_stat_file() { - let sql = list_objects_sql(false, true, false); + let sql = list_objects_sql(false, true, false, true); assert!(!sql.contains("pg_stat_file")); assert!(sql.contains("NULL::text AS created_at")); assert!(sql.contains("NULL::text AS updated_at")); } + #[test] + fn redshift_compatible_list_objects_sql_uses_legacy_argument_formatter() { + let sql = list_objects_sql(false, false, false, false); + assert!(sql.contains("pg_get_function_arguments(p.oid) AS signature")); + assert!(!sql.contains("pg_get_function_identity_arguments")); + } + + #[test] + fn function_identity_arguments_probe_uses_pg_proc() { + let sql = postgres_has_function_identity_arguments_sql(); + assert!(sql.contains("pg_catalog.pg_proc")); + assert!(sql.contains("n.nspname = 'pg_catalog'")); + assert!(sql.contains("p.proname = 'pg_get_function_identity_arguments'")); + } + #[test] fn both_list_objects_sql_variants_use_parameter() { - assert!(list_objects_sql(true, true, true).contains("$1")); - assert!(list_objects_sql(false, true, true).contains("$1")); - assert!(list_objects_sql(true, true, false).contains("$1")); - assert!(list_objects_sql(false, true, false).contains("$1")); - assert!(list_objects_sql(true, false, true).contains("$1")); - assert!(list_objects_sql(false, false, true).contains("$1")); - assert!(list_objects_sql(true, false, false).contains("$1")); - assert!(list_objects_sql(false, false, false).contains("$1")); + assert!(list_objects_sql(true, true, true, true).contains("$1")); + assert!(list_objects_sql(false, true, true, true).contains("$1")); + assert!(list_objects_sql(true, true, false, true).contains("$1")); + assert!(list_objects_sql(false, true, false, true).contains("$1")); + assert!(list_objects_sql(true, false, true, true).contains("$1")); + assert!(list_objects_sql(false, false, true, true).contains("$1")); + assert!(list_objects_sql(true, false, false, true).contains("$1")); + assert!(list_objects_sql(false, false, false, true).contains("$1")); } #[test] fn both_list_objects_sql_variants_include_pg_proc() { - assert!(list_objects_sql(true, true, true).contains("pg_catalog.pg_proc")); - assert!(list_objects_sql(false, true, true).contains("pg_catalog.pg_proc")); - assert!(list_objects_sql(true, true, false).contains("pg_catalog.pg_proc")); - assert!(list_objects_sql(false, true, false).contains("pg_catalog.pg_proc")); - assert!(list_objects_sql(true, false, true).contains("pg_catalog.pg_proc")); - assert!(list_objects_sql(false, false, true).contains("pg_catalog.pg_proc")); - assert!(list_objects_sql(true, false, false).contains("pg_catalog.pg_proc")); - assert!(list_objects_sql(false, false, false).contains("pg_catalog.pg_proc")); + assert!(list_objects_sql(true, true, true, true).contains("pg_catalog.pg_proc")); + assert!(list_objects_sql(false, true, true, true).contains("pg_catalog.pg_proc")); + assert!(list_objects_sql(true, true, false, true).contains("pg_catalog.pg_proc")); + assert!(list_objects_sql(false, true, false, true).contains("pg_catalog.pg_proc")); + assert!(list_objects_sql(true, false, true, true).contains("pg_catalog.pg_proc")); + assert!(list_objects_sql(false, false, true, true).contains("pg_catalog.pg_proc")); + assert!(list_objects_sql(true, false, false, true).contains("pg_catalog.pg_proc")); + assert!(list_objects_sql(false, false, false, true).contains("pg_catalog.pg_proc")); } #[test] fn legacy_list_objects_sql_avoids_pg11_proc_kind_column() { - let sql = list_objects_sql(true, false, false); + let sql = list_objects_sql(true, false, false, true); assert!(!sql.contains("p.prokind")); assert!(!sql.contains("p.prosp")); assert!(sql.contains("NOT p.proisagg")); @@ -3953,7 +4018,7 @@ mod tests { #[test] fn gaussdb_compatible_list_objects_sql_uses_prosp_when_prokind_is_missing() { - let sql = list_objects_sql(true, false, true); + let sql = list_objects_sql(true, false, true, true); assert!(!sql.contains("p.prokind")); assert!(sql.contains("CASE WHEN p.prosp THEN 'PROCEDURE' ELSE 'FUNCTION' END AS object_type")); assert!(sql.contains("CASE WHEN p.prosp THEN 2 ELSE 3 END AS sort_order")); @@ -3964,7 +4029,7 @@ mod tests { #[test] fn gaussdb_compatible_list_objects_sql_uses_prosp_with_prokind_when_available() { - let sql = list_objects_sql(true, true, true); + let sql = list_objects_sql(true, true, true, true); assert!( sql.contains("CASE WHEN p.prokind = 'p' OR p.prosp THEN 'PROCEDURE' ELSE 'FUNCTION' END AS object_type") );