diff --git a/crates/dbx-core/src/db/postgres.rs b/crates/dbx-core/src/db/postgres.rs index 8bf4dd5dd..78fcfed0a 100644 --- a/crates/dbx-core/src/db/postgres.rs +++ b/crates/dbx-core/src/db/postgres.rs @@ -1513,7 +1513,7 @@ fn like_fuzzy_pattern(value: &str) -> String { }) } -fn list_objects_sql(include_timestamps: bool) -> &'static str { +fn list_object_relations_sql(include_timestamps: bool) -> &'static str { if include_timestamps { return "SELECT c.relname AS object_name, \ CASE c.relkind \ @@ -1540,21 +1540,7 @@ fn list_objects_sql(include_timestamps: bool) -> &'static str { 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 p.prokind WHEN 'p' 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 p.prokind WHEN 'p' 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 p.prokind IN ('p','f') \ - ORDER BY sort_order, object_name"; + WHERE n.nspname = $1 AND c.relkind IN ('r','v','m','f','p','S')"; } "SELECT c.relname AS object_name, \ @@ -1575,9 +1561,27 @@ fn list_objects_sql(include_timestamps: bool) -> &'static str { 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, \ + WHERE n.nspname = $1 AND c.relkind IN ('r','v','m','f','p','S')" +} + +fn list_object_routines_sql(include_timestamps: bool, has_proc_prokind: bool) -> &'static str { + if has_proc_prokind { + if include_timestamps { + return "SELECT p.proname AS object_name, \ + CASE p.prokind WHEN 'p' 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 p.prokind WHEN 'p' 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 p.prokind IN ('p','f')"; + } + + return "SELECT p.proname AS object_name, \ CASE p.prokind WHEN 'p' THEN 'PROCEDURE' ELSE 'FUNCTION' END AS object_type, \ obj_description(p.oid) AS object_comment, \ NULL::text AS created_at, \ @@ -1587,17 +1591,71 @@ fn list_objects_sql(include_timestamps: bool) -> &'static str { CASE p.prokind WHEN 'p' 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 p.prokind IN ('p','f') \ - ORDER BY sort_order, object_name" + WHERE n.nspname = $1 AND p.prokind IN ('p','f')"; + } + + if include_timestamps { + return "SELECT p.proname AS object_name, \ + 'FUNCTION' 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, \ + 3 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"; + } + + "SELECT p.proname AS object_name, \ + 'FUNCTION' 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, \ + 3 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" +} + +fn list_objects_sql(include_timestamps: bool, has_proc_prokind: bool) -> String { + format!( + "{} UNION ALL {} ORDER BY sort_order, object_name", + list_object_relations_sql(include_timestamps), + list_object_routines_sql(include_timestamps, has_proc_prokind) + ) +} + +fn postgres_proc_has_prokind_sql() -> &'static str { + "SELECT EXISTS ( \ + SELECT 1 \ + FROM pg_catalog.pg_attribute \ + WHERE attrelid = 'pg_catalog.pg_proc'::regclass \ + AND attname = 'prokind' \ + AND NOT attisdropped \ + )" +} + +async fn postgres_proc_has_prokind(client: &deadpool_postgres::Client) -> Result { + let stmt = client.prepare_cached(postgres_proc_has_prokind_sql()).await.map_err(|e| e.to_string())?; + let row = client.query_one(&stmt, &[]).await.map_err(|e| e.to_string())?; + Ok(row.get(0)) } pub async fn list_objects(pool: &Pool, schema: &str) -> Result, String> { let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; - let stmt = client.prepare_cached(list_objects_sql(true)).await.map_err(|e| e.to_string())?; + let has_proc_prokind = postgres_proc_has_prokind(&client).await?; + let sql = list_objects_sql(true, has_proc_prokind); + let stmt = client.prepare_cached(&sql).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())?; + let fallback_sql = list_objects_sql(false, has_proc_prokind); + let stmt = client.prepare_cached(&fallback_sql).await.map_err(|e| e.to_string())?; client.query(&stmt, &[&schema]).await.map_err(|e| e.to_string())? } }; @@ -3154,7 +3212,7 @@ mod tests { #[test] fn list_objects_sql_includes_routines() { - let sql = list_objects_sql(true); + let sql = list_objects_sql(true, true); assert!(sql.contains("pg_catalog.pg_class")); assert!(sql.contains("pg_catalog.pg_proc")); assert!(sql.contains("pg_catalog.pg_inherits")); @@ -3169,7 +3227,7 @@ mod tests { #[test] fn list_objects_sql_without_timestamps_omits_stat_file() { - let sql = list_objects_sql(false); + let sql = list_objects_sql(false, true); assert!(!sql.contains("pg_stat_file")); assert!(sql.contains("NULL::text AS created_at")); assert!(sql.contains("NULL::text AS updated_at")); @@ -3177,14 +3235,36 @@ mod tests { #[test] fn both_list_objects_sql_variants_use_parameter() { - assert!(list_objects_sql(true).contains("$1")); - assert!(list_objects_sql(false).contains("$1")); + assert!(list_objects_sql(true, true).contains("$1")); + assert!(list_objects_sql(false, true).contains("$1")); + assert!(list_objects_sql(true, false).contains("$1")); + assert!(list_objects_sql(false, false).contains("$1")); } #[test] fn both_list_objects_sql_variants_include_pg_proc() { - assert!(list_objects_sql(true).contains("pg_catalog.pg_proc")); - assert!(list_objects_sql(false).contains("pg_catalog.pg_proc")); + assert!(list_objects_sql(true, true).contains("pg_catalog.pg_proc")); + assert!(list_objects_sql(false, true).contains("pg_catalog.pg_proc")); + assert!(list_objects_sql(true, false).contains("pg_catalog.pg_proc")); + assert!(list_objects_sql(false, false).contains("pg_catalog.pg_proc")); + } + + #[test] + fn legacy_list_objects_sql_avoids_pg11_proc_kind_column() { + let sql = list_objects_sql(true, false); + assert!(!sql.contains("p.prokind")); + assert!(sql.contains("NOT p.proisagg")); + assert!(sql.contains("NOT p.proiswindow")); + assert!(sql.contains("'FUNCTION' AS object_type")); + assert!(!sql.contains("'PROCEDURE'")); + } + + #[test] + fn postgres_proc_has_prokind_sql_checks_catalog_attribute() { + let sql = postgres_proc_has_prokind_sql(); + assert!(sql.contains("pg_catalog.pg_attribute")); + assert!(sql.contains("'pg_catalog.pg_proc'::regclass")); + assert!(sql.contains("attname = 'prokind'")); } #[test] diff --git a/crates/dbx-core/src/schema.rs b/crates/dbx-core/src/schema.rs index 374ec04ee..ba898834b 100644 --- a/crates/dbx-core/src/schema.rs +++ b/crates/dbx-core/src/schema.rs @@ -3062,6 +3062,18 @@ fn postgres_object_source_sql_without_relispopulated(schema: &str, name: &str, k postgres_object_source_sql_inner(schema, name, kind, false) } +fn postgres_function_object_source_sql_without_prokind(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) + ) +} + fn postgres_object_source_sql_inner( schema: &str, name: &str, @@ -3427,6 +3439,16 @@ async fn postgres_object_source( .and_then(first_string_cell) .map_err(|fallback_err| format!("{primary_err}; relispopulated fallback failed: {fallback_err}")) } + Err(primary_err) + if postgres_missing_prokind_error(&primary_err) + && matches!(object_type, db::ObjectSourceKind::Function) => + { + let fallback_sql = postgres_function_object_source_sql_without_prokind(schema, name); + db::postgres::execute_query(pool, &fallback_sql) + .await + .and_then(first_string_cell) + .map_err(|fallback_err| format!("{primary_err}; prokind fallback failed: {fallback_err}")) + } Err(primary_err) if matches!(object_type, db::ObjectSourceKind::View) => { let fallback_sql = postgres_view_source_fallback_sql(schema, name); db::postgres::execute_query(pool, &fallback_sql) @@ -3438,6 +3460,14 @@ async fn postgres_object_source( } } +fn postgres_missing_prokind_error(err: &str) -> bool { + let lower = err.to_ascii_lowercase(); + lower.contains("does not exist") + && (lower.contains("column p.prokind") + || lower.contains("column \"p\".\"prokind\"") + || lower.contains("column \"prokind\"")) +} + fn postgres_missing_relispopulated_error(err: &str) -> bool { let lower = err.to_ascii_lowercase(); lower.contains("does not exist") @@ -3488,6 +3518,16 @@ mod object_source_tests { assert!(!sql.contains("relispopulated")); } + #[test] + fn builds_postgres_function_source_sql_without_prokind_for_legacy_catalogs() { + let sql = postgres_function_object_source_sql_without_prokind("public", "recalc_score"); + + assert_eq!( + sql, + "SELECT pg_get_functiondef(p.oid) FROM pg_proc p JOIN pg_namespace n ON n.oid = p.pronamespace WHERE n.nspname = 'public' AND p.proname = 'recalc_score' AND NOT p.proisagg AND NOT p.proiswindow ORDER BY p.oid LIMIT 1" + ); + } + #[test] fn keeps_legacy_materialized_viewdef_when_it_already_contains_create_statement() { let sql = postgres_object_source_sql("public", "active_users", &ObjectSourceKind::MaterializedView); @@ -3508,6 +3548,13 @@ mod object_source_tests { assert!(!postgres_missing_relispopulated_error("ERROR: relation public.relispopulated does not exist")); } + #[test] + fn detects_legacy_postgres_prokind_errors() { + assert!(postgres_missing_prokind_error("ERROR: column p.prokind does not exist")); + assert!(postgres_missing_prokind_error("ERROR: column \"p\".\"prokind\" does not exist")); + assert!(!postgres_missing_prokind_error("ERROR: relation public.prokind does not exist")); + } + #[test] fn builds_postgres_view_source_sql_without_regclass_cast() { let sql = postgres_object_source_sql("tenant's schema", "active users", &ObjectSourceKind::View);