From 51e2ef0d1f4bbb0b48d1a31fc96fadbd6109283f Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Fri, 12 Jun 2026 00:59:52 +0800 Subject: [PATCH] fix(postgres): support legacy routine metadata --- crates/dbx-core/src/db/postgres.rs | 141 +++++++++++++++++++++++++++-- crates/dbx-core/src/schema.rs | 30 ++++++ 2 files changed, 164 insertions(+), 7 deletions(-) diff --git a/crates/dbx-core/src/db/postgres.rs b/crates/dbx-core/src/db/postgres.rs index f8a2e9293..400e6f8f1 100644 --- a/crates/dbx-core/src/db/postgres.rs +++ b/crates/dbx-core/src/db/postgres.rs @@ -1225,15 +1225,112 @@ fn list_objects_sql(include_timestamps: bool) -> &'static str { ORDER BY sort_order, object_name" } +fn list_objects_legacy_routines_sql(include_timestamps: bool) -> &'static str { + if include_timestamps { + return "SELECT c.relname AS object_name, \ + CASE c.relkind \ + WHEN 'v' THEN 'VIEW' \ + WHEN 'm' THEN 'VIEW' \ + WHEN 'S' THEN 'SEQUENCE' \ + ELSE 'TABLE' \ + END AS object_type, \ + obj_description(c.oid) AS object_comment, \ + stat.creation::text AS created_at, \ + COALESCE( \ + CASE WHEN current_setting('track_commit_timestamp', true) = 'on' \ + THEN pg_xact_commit_timestamp(c.xmin)::text END, \ + stat.modification::text \ + ) AS updated_at, \ + CASE WHEN pc.relkind = 'p' THEN pn.nspname ELSE NULL END AS parent_schema, \ + CASE WHEN pc.relkind = 'p' THEN pc.relname ELSE NULL END AS parent_name, \ + CASE c.relkind WHEN 'v' THEN 1 WHEN 'm' THEN 1 WHEN 'S' THEN 4 ELSE 0 END AS sort_order \ + FROM pg_catalog.pg_class c \ + JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ + LEFT JOIN pg_catalog.pg_inherits i ON i.inhrelid = c.oid \ + LEFT JOIN pg_catalog.pg_class pc ON pc.oid = i.inhparent \ + LEFT JOIN pg_catalog.pg_namespace pn ON pn.oid = pc.relnamespace \ + LEFT JOIN LATERAL pg_stat_file( \ + CASE WHEN c.relkind IN ('r','m','f','p') THEN pg_relation_filepath(c.oid) END, true \ + ) stat ON true \ + WHERE n.nspname = $1 AND c.relkind IN ('r','v','m','f','p','S') \ + UNION ALL \ + SELECT p.proname AS object_name, \ + CASE WHEN EXISTS ( \ + SELECT 1 FROM information_schema.routines r \ + WHERE r.specific_schema = n.nspname \ + AND r.routine_name = p.proname \ + AND upper(r.routine_type) = 'PROCEDURE' \ + ) THEN 'PROCEDURE' ELSE 'FUNCTION' END AS object_type, \ + obj_description(p.oid) AS object_comment, \ + NULL::text AS created_at, \ + CASE WHEN current_setting('track_commit_timestamp', true) = 'on' \ + THEN pg_xact_commit_timestamp(p.xmin)::text END AS updated_at, \ + NULL::text AS parent_schema, \ + NULL::text AS parent_name, \ + CASE WHEN EXISTS ( \ + SELECT 1 FROM information_schema.routines r \ + WHERE r.specific_schema = n.nspname \ + AND r.routine_name = p.proname \ + AND upper(r.routine_type) = 'PROCEDURE' \ + ) THEN 2 ELSE 3 END AS sort_order \ + FROM pg_catalog.pg_proc p \ + JOIN pg_catalog.pg_namespace n ON n.oid = p.pronamespace \ + WHERE n.nspname = $1 AND NOT p.proisagg AND NOT p.proiswindow \ + ORDER BY sort_order, object_name"; + } + + "SELECT c.relname AS object_name, \ + CASE c.relkind \ + WHEN 'v' THEN 'VIEW' \ + WHEN 'm' THEN 'VIEW' \ + WHEN 'S' THEN 'SEQUENCE' \ + ELSE 'TABLE' \ + END AS object_type, \ + obj_description(c.oid) AS object_comment, \ + NULL::text AS created_at, \ + NULL::text AS updated_at, \ + CASE WHEN pc.relkind = 'p' THEN pn.nspname ELSE NULL END AS parent_schema, \ + CASE WHEN pc.relkind = 'p' THEN pc.relname ELSE NULL END AS parent_name, \ + CASE c.relkind WHEN 'v' THEN 1 WHEN 'm' THEN 1 WHEN 'S' THEN 4 ELSE 0 END AS sort_order \ + FROM pg_catalog.pg_class c \ + JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ + LEFT JOIN pg_catalog.pg_inherits i ON i.inhrelid = c.oid \ + LEFT JOIN pg_catalog.pg_class pc ON pc.oid = i.inhparent \ + LEFT JOIN pg_catalog.pg_namespace pn ON pn.oid = pc.relnamespace \ + WHERE n.nspname = $1 AND c.relkind IN ('r','v','m','f','p','S') \ + UNION ALL \ + SELECT p.proname AS object_name, \ + CASE WHEN EXISTS ( \ + SELECT 1 FROM information_schema.routines r \ + WHERE r.specific_schema = n.nspname \ + AND r.routine_name = p.proname \ + AND upper(r.routine_type) = 'PROCEDURE' \ + ) THEN 'PROCEDURE' ELSE 'FUNCTION' END AS object_type, \ + obj_description(p.oid) AS object_comment, \ + NULL::text AS created_at, \ + NULL::text AS updated_at, \ + NULL::text AS parent_schema, \ + NULL::text AS parent_name, \ + CASE WHEN EXISTS ( \ + SELECT 1 FROM information_schema.routines r \ + WHERE r.specific_schema = n.nspname \ + AND r.routine_name = p.proname \ + AND upper(r.routine_type) = 'PROCEDURE' \ + ) THEN 2 ELSE 3 END AS sort_order \ + FROM pg_catalog.pg_proc p \ + JOIN pg_catalog.pg_namespace n ON n.oid = p.pronamespace \ + WHERE n.nspname = $1 AND NOT p.proisagg AND NOT p.proiswindow \ + ORDER BY sort_order, object_name" +} + pub async fn list_objects(pool: &Pool, schema: &str) -> Result, String> { let client = pool.get().await.map_err(|e| e.to_string())?; - let stmt = client.prepare_cached(list_objects_sql(true)).await.map_err(|e| e.to_string())?; - let rows = match client.query(&stmt, &[&schema]).await { - Ok(rows) => rows, - Err(_) => { - let stmt = client.prepare_cached(list_objects_sql(false)).await.map_err(|e| e.to_string())?; - client.query(&stmt, &[&schema]).await.map_err(|e| e.to_string())? - } + let rows = match client.prepare_cached(list_objects_sql(true)).await { + Ok(stmt) => match client.query(&stmt, &[&schema]).await { + Ok(rows) => rows, + Err(_) => query_list_objects_fallbacks(&client, schema).await?, + }, + Err(_) => query_list_objects_fallbacks(&client, schema).await?, }; Ok(rows @@ -1251,6 +1348,22 @@ pub async fn list_objects(pool: &Pool, schema: &str) -> Result, .collect()) } +async fn query_list_objects_fallbacks( + client: &deadpool_postgres::Object, + schema: &str, +) -> Result, String> { + for sql in + [list_objects_sql(false), list_objects_legacy_routines_sql(true), list_objects_legacy_routines_sql(false)] + { + if let Ok(stmt) = client.prepare_cached(sql).await { + if let Ok(rows) = client.query(&stmt, &[&schema]).await { + return Ok(rows); + } + } + } + Err("failed to list PostgreSQL objects with compatible metadata queries".to_string()) +} + pub async fn list_schemas(pool: &Pool) -> Result, String> { let client = pool.get().await.map_err(|e| e.to_string())?; let stmt = client @@ -2153,6 +2266,20 @@ mod tests { assert!(list_objects_sql(false).contains("pg_catalog.pg_proc")); } + #[test] + fn legacy_list_objects_sql_avoids_pg11_prokind() { + let sql = list_objects_legacy_routines_sql(false); + + assert!(sql.contains("pg_catalog.pg_proc")); + assert!(sql.contains("information_schema.routines")); + assert!(sql.contains("'PROCEDURE'")); + assert!(sql.contains("'FUNCTION'")); + assert!(sql.contains("NOT p.proisagg")); + assert!(sql.contains("NOT p.proiswindow")); + assert!(!sql.contains("p.prokind")); + assert!(sql.contains("'SEQUENCE'")); + } + #[test] fn transaction_recovery_statement_detection_matches_common_postgres_commands() { assert!(is_transaction_recovery_statement("ROLLBACK")); diff --git a/crates/dbx-core/src/schema.rs b/crates/dbx-core/src/schema.rs index cc5f450af..bff4ce000 100644 --- a/crates/dbx-core/src/schema.rs +++ b/crates/dbx-core/src/schema.rs @@ -1060,6 +1060,19 @@ pub fn postgres_object_source_sql(schema: &str, name: &str, kind: &db::ObjectSou } } +pub fn postgres_routine_object_source_legacy_sql(schema: &str, name: &str) -> String { + format!( + "SELECT pg_get_functiondef(p.oid) \ + FROM pg_proc p \ + JOIN pg_namespace n ON n.oid = p.pronamespace \ + WHERE n.nspname = {} AND p.proname = {} \ + AND NOT p.proisagg AND NOT p.proiswindow \ + ORDER BY p.oid LIMIT 1", + sql_string(schema), + sql_string(name) + ) +} + pub fn oracle_object_source_sql(schema: &str, name: &str, kind: &db::ObjectSourceKind) -> String { let object_type = match kind { db::ObjectSourceKind::View => "VIEW", @@ -1313,6 +1326,13 @@ async fn postgres_object_source( .and_then(first_string_cell) .map_err(|fallback_err| format!("{primary_err}; fallback failed: {fallback_err}")) } + Err(primary_err) if matches!(object_type, db::ObjectSourceKind::Procedure | db::ObjectSourceKind::Function) => { + let fallback_sql = postgres_routine_object_source_legacy_sql(schema, name); + db::postgres::execute_query(pool, &fallback_sql) + .await + .and_then(first_string_cell) + .map_err(|fallback_err| format!("{primary_err}; fallback failed: {fallback_err}")) + } Err(err) => Err(err), } } @@ -1342,6 +1362,16 @@ mod object_source_tests { ); } + #[test] + fn builds_postgres_legacy_routine_source_sql_without_prokind() { + let sql = postgres_routine_object_source_legacy_sql("public", "recalc_score"); + + assert!(sql.contains("pg_get_functiondef(p.oid)")); + assert!(sql.contains("NOT p.proisagg")); + assert!(sql.contains("NOT p.proiswindow")); + assert!(!sql.contains("p.prokind")); + } + #[test] fn builds_postgres_object_source_sql_for_sequences() { let sql = postgres_object_source_sql("tenant's schema", "order id seq", &ObjectSourceKind::Sequence);