From 8548ecacc07a63815549124dea69b92f4e6a2076 Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Thu, 13 Aug 2026 16:13:07 +0800 Subject: [PATCH] fix(transfer): complete PostgreSQL 9.x table transfer Closes #6070 --- crates/dbx-core/src/transfer.rs | 246 +++++++++++++++++++++----------- 1 file changed, 161 insertions(+), 85 deletions(-) diff --git a/crates/dbx-core/src/transfer.rs b/crates/dbx-core/src/transfer.rs index de6f1a37e8..709987e8c9 100644 --- a/crates/dbx-core/src/transfer.rs +++ b/crates/dbx-core/src/transfer.rs @@ -1235,6 +1235,14 @@ fn is_postgres_transfer_dialect(db_type: &DatabaseType) -> bool { matches!(db_type, DatabaseType::Postgres | DatabaseType::Kingbase) } +fn transfer_table_needs_inline_postgres_schema_ensure( + source_db_type: &DatabaseType, + target_db_type: &DatabaseType, +) -> bool { + is_postgres_transfer_dialect(target_db_type) + && !matches!((source_db_type, target_db_type), (DatabaseType::Postgres, DatabaseType::Postgres)) +} + fn postgres_integer_bounds(data_type: &str) -> Option<(i128, i128)> { let normalized = data_type.trim().to_ascii_lowercase(); match normalized.split(['(', ' ']).next().unwrap_or("") { @@ -4580,27 +4588,7 @@ async fn get_postgres_owned_sequences_for_transfer( } }; let client = pool.get().await.map_err(|e| e.to_string())?; - let rows = client - .query( - "SELECT c.relname, \ - t.relname, \ - a.attname \ - FROM pg_class c \ - JOIN pg_namespace n ON n.oid = c.relnamespace \ - LEFT JOIN pg_sequence s ON s.seqrelid = c.oid \ - JOIN pg_depend d ON d.classid = 'pg_class'::regclass \ - AND d.objid = c.oid \ - AND d.refclassid = 'pg_class'::regclass \ - AND d.deptype IN ('a', 'i') \ - JOIN pg_class t ON t.oid = d.refobjid \ - JOIN pg_namespace tn ON tn.oid = t.relnamespace AND tn.nspname = n.nspname \ - JOIN pg_attribute a ON a.attrelid = t.oid AND a.attnum = d.refobjsubid \ - WHERE c.relkind = 'S' AND n.nspname = $1 \ - ORDER BY t.relname, c.relname", - &[&schema], - ) - .await - .map_err(|e| e.to_string())?; + let rows = client.query(POSTGRES_OWNED_SEQUENCES_SQL, &[&schema]).await.map_err(|e| e.to_string())?; let selected: HashSet<&str> = tables.iter().map(String::as_str).collect(); Ok(rows @@ -4619,6 +4607,21 @@ async fn get_postgres_owned_sequences_for_transfer( .collect()) } +const POSTGRES_OWNED_SEQUENCES_SQL: &str = "SELECT c.relname, \ + t.relname, \ + a.attname \ + FROM pg_class c \ + JOIN pg_namespace n ON n.oid = c.relnamespace \ + JOIN pg_depend d ON d.classid = 'pg_class'::regclass \ + AND d.objid = c.oid \ + AND d.refclassid = 'pg_class'::regclass \ + AND d.deptype IN ('a', 'i') \ + JOIN pg_class t ON t.oid = d.refobjid \ + JOIN pg_namespace tn ON tn.oid = t.relnamespace AND tn.nspname = n.nspname \ + JOIN pg_attribute a ON a.attrelid = t.oid AND a.attnum = d.refobjsubid \ + WHERE c.relkind = 'S' AND n.nspname = $1 \ + ORDER BY t.relname, c.relname"; + async fn get_postgres_sequence_snapshots_for_transfer( state: &AppState, pool_key: &str, @@ -4632,26 +4635,7 @@ async fn get_postgres_sequence_snapshots_for_transfer( } }; let client = pool.get().await.map_err(|e| e.to_string())?; - let rows = client - .query( - "SELECT c.relname, \ - t.relname, \ - a.attname \ - FROM pg_class c \ - JOIN pg_namespace n ON n.oid = c.relnamespace \ - LEFT JOIN pg_sequence s ON s.seqrelid = c.oid \ - LEFT JOIN pg_depend d ON d.classid = 'pg_class'::regclass \ - AND d.objid = c.oid \ - AND d.refclassid = 'pg_class'::regclass \ - AND d.deptype IN ('a', 'i') \ - LEFT JOIN pg_class t ON t.oid = d.refobjid \ - LEFT JOIN pg_attribute a ON a.attrelid = t.oid AND a.attnum = d.refobjsubid \ - WHERE c.relkind = 'S' AND n.nspname = $1 \ - ORDER BY c.relname", - &[&schema], - ) - .await - .map_err(|e| e.to_string())?; + let rows = client.query(POSTGRES_SEQUENCE_SNAPSHOTS_SQL, &[&schema]).await.map_err(|e| e.to_string())?; Ok(rows .iter() @@ -4663,6 +4647,20 @@ async fn get_postgres_sequence_snapshots_for_transfer( .collect()) } +const POSTGRES_SEQUENCE_SNAPSHOTS_SQL: &str = "SELECT c.relname, \ + t.relname, \ + a.attname \ + FROM pg_class c \ + JOIN pg_namespace n ON n.oid = c.relnamespace \ + LEFT JOIN pg_depend d ON d.classid = 'pg_class'::regclass \ + AND d.objid = c.oid \ + AND d.refclassid = 'pg_class'::regclass \ + AND d.deptype IN ('a', 'i') \ + LEFT JOIN pg_class t ON t.oid = d.refobjid \ + LEFT JOIN pg_attribute a ON a.attrelid = t.oid AND a.attnum = d.refobjsubid \ + WHERE c.relkind = 'S' AND n.nspname = $1 \ + ORDER BY c.relname"; + fn postgres_selected_sequences_sql(schema: &str, names: &[String]) -> Option { if names.is_empty() { return None; @@ -4779,10 +4777,8 @@ async fn prepare_postgres_owned_sequences_for_transfer( &request.target_schema, )?; if should_create { - let create_sql = format!( - "CREATE SEQUENCE IF NOT EXISTS {}", - postgres_sequence_qualified_name(&request.target_schema, &sequence.name) - ); + let create_sql = + format!("CREATE SEQUENCE {}", postgres_sequence_qualified_name(&request.target_schema, &sequence.name)); execute_on_pool(state, target_pool_key, &create_sql) .await .map_err(|e| format!("Failed to create PostgreSQL sequence for {target_table}: {e}"))?; @@ -5332,31 +5328,62 @@ where } Ok(outcome) } -fn postgres_transfer_routine_catalog_capability_sql() -> &'static str { +fn postgres_transfer_catalog_capabilities_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 \ + ), EXISTS ( \ + SELECT 1 \ + FROM pg_catalog.pg_class c \ + JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ + WHERE n.nspname = 'pg_catalog' AND c.relname = 'pg_policy' \ + ), EXISTS ( \ + SELECT 1 \ + FROM pg_catalog.pg_attribute a \ + JOIN pg_catalog.pg_class c ON c.oid = a.attrelid \ + JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ + WHERE n.nspname = 'pg_catalog' \ + AND c.relname = 'pg_policy' \ + AND a.attname = 'polpermissive' \ + AND NOT a.attisdropped \ )" } -async fn postgres_transfer_has_prokind(state: &AppState, pool_key: &str) -> Result { - let result = execute_read_on_pool(state, pool_key, postgres_transfer_routine_catalog_capability_sql()).await?; - result - .rows - .first() - .and_then(|row| row.first()) - .and_then(|value| value.as_bool().or_else(|| value.as_str().and_then(|value| value.parse().ok()))) - .ok_or_else(|| "Failed to inspect PostgreSQL routine catalog capabilities".to_string()) +struct PostgresTransferCatalogCapabilities { + has_prokind: bool, + has_pg_policy: bool, + supports_policy_permissiveness: bool, +} + +fn postgres_capability_cell(row: &[serde_json::Value], index: usize) -> Option { + row.get(index).and_then(|value| value.as_bool().or_else(|| value.as_str().and_then(|value| value.parse().ok()))) +} + +async fn postgres_transfer_catalog_capabilities( + state: &AppState, + pool_key: &str, +) -> Result { + let result = execute_read_on_pool(state, pool_key, postgres_transfer_catalog_capabilities_sql()).await?; + let row = + result.rows.first().ok_or_else(|| "Failed to inspect PostgreSQL transfer catalog capabilities".to_string())?; + Ok(PostgresTransferCatalogCapabilities { + has_prokind: postgres_capability_cell(row, 0) + .ok_or_else(|| "Failed to inspect PostgreSQL routine catalog capabilities".to_string())?, + has_pg_policy: postgres_capability_cell(row, 1) + .ok_or_else(|| "Failed to inspect PostgreSQL policy catalog capabilities".to_string())?, + supports_policy_permissiveness: postgres_capability_cell(row, 2) + .ok_or_else(|| "Failed to inspect PostgreSQL policy catalog capabilities".to_string())?, + }) } fn postgres_transfer_routine_catalog_sql(has_prokind: bool) -> (&'static str, &'static str) { if has_prokind { ("CASE p.prokind WHEN 'p' THEN 'PROCEDURE' ELSE 'FUNCTION' END", "p.prokind IN ('p','f')") } else { - ("'FUNCTION'", "NOT p.proisagg AND NOT p.proiswindow") + ("'FUNCTION'::text", "NOT p.proisagg AND NOT p.proiswindow") } } @@ -5574,10 +5601,17 @@ async fn get_postgres_policy_statements_for_transfer( source_schema: &str, target_schema: &str, tables: &[String], + has_pg_policy: bool, + supports_policy_permissiveness: bool, ) -> Result, String> { - if tables.is_empty() { + if tables.is_empty() || !has_pg_policy { return Ok(Vec::new()); } + let (policy_permissiveness_select, policy_permissiveness_clause) = if supports_policy_permissiveness { + ("p.polpermissive", "CASE WHEN polpermissive THEN ' AS PERMISSIVE' ELSE ' AS RESTRICTIVE' END") + } else { + ("true", "''") + }; let table_list = tables.iter().map(|table| quote_string_literal(table)).collect::>().join(", "); let sql = format!( "WITH selected_tables AS ( \ @@ -5587,7 +5621,8 @@ async fn get_postgres_policy_statements_for_transfer( WHERE n.nspname = {source_schema} AND c.relkind IN ('r','p') AND c.relname IN ({table_list}) \ ), \ policy_rows AS ( \ - SELECT t.relname, t.relrowsecurity, t.relforcerowsecurity, p.polname, p.polpermissive, p.polcmd, \ + SELECT t.relname, t.relrowsecurity, t.relforcerowsecurity, p.polname, \ + {policy_permissiveness_select} AS polpermissive, p.polcmd, \ COALESCE((SELECT string_agg(CASE WHEN role_oid = 0 THEN 'PUBLIC' ELSE quote_ident(r.rolname) END, ', ' ORDER BY CASE WHEN role_oid = 0 THEN '' ELSE r.rolname END) \ FROM unnest(p.polroles) AS role_oid LEFT JOIN pg_roles r ON r.oid = role_oid), '') AS role_list, \ pg_get_expr(p.polqual, p.polrelid) AS using_expr, \ @@ -5606,9 +5641,9 @@ async fn get_postgres_policy_statements_for_transfer( FROM policy_rows \ UNION ALL \ SELECT format( \ - 'CREATE POLICY %I ON %I.%I AS %s FOR %s%s%s%s', \ + 'CREATE POLICY %I ON %I.%I%s FOR %s%s%s%s', \ polname, {target_schema}, relname, \ - CASE WHEN polpermissive THEN 'PERMISSIVE' ELSE 'RESTRICTIVE' END, \ + {policy_permissiveness_clause}, \ CASE polcmd WHEN 'r' THEN 'SELECT' WHEN 'a' THEN 'INSERT' WHEN 'w' THEN 'UPDATE' WHEN 'd' THEN 'DELETE' ELSE 'ALL' END, \ CASE WHEN role_list <> '' THEN ' TO ' || role_list ELSE '' END, \ CASE WHEN using_expr IS NOT NULL THEN ' USING (' || using_expr || ')' ELSE '' END, \ @@ -5619,6 +5654,8 @@ async fn get_postgres_policy_statements_for_transfer( ORDER BY relname, sort_order, stmt", source_schema = quote_string_literal(source_schema), target_schema = quote_string_literal(target_schema), + policy_permissiveness_select = policy_permissiveness_select, + policy_permissiveness_clause = policy_permissiveness_clause, ); Ok(result_rows_to_string_statements(execute_on_pool(state, pool_key, &sql).await?.rows)) } @@ -5735,7 +5772,7 @@ pub async fn preview_transfer_ownership( return Ok(TransferOwnershipPreview { missing_owners: Vec::new(), target_owner: String::new() }); } - let has_prokind = postgres_transfer_has_prokind(state, source_pool_key).await?; + let has_prokind = postgres_transfer_catalog_capabilities(state, source_pool_key).await?.has_prokind; let relation_names = postgres_transfer_relation_names(request); let statements = get_postgres_ownership_statements_for_transfer( state, @@ -5776,10 +5813,12 @@ fn postgres_transfer_grant_statements_sql( CASE WHEN a.grantee = 0 THEN 'PUBLIC' ELSE quote_ident(grantee.rolname) END, \ CASE WHEN bool_or(a.is_grantable) THEN ' WITH GRANT OPTION' ELSE '' END \ ) AS stmt \ - FROM pg_catalog.pg_namespace n \ - JOIN LATERAL aclexplode(n.nspacl) a ON true \ + FROM ( \ + SELECT n.nspname, (aclexplode(n.nspacl)).* \ + FROM pg_catalog.pg_namespace n \ + WHERE n.nspname = {source_schema} \ + ) a \ LEFT JOIN pg_roles grantee ON grantee.oid = a.grantee \ - WHERE n.nspname = {source_schema} \ GROUP BY a.grantee, grantee.rolname \ ), \ relation_grants AS ( \ @@ -5792,12 +5831,14 @@ fn postgres_transfer_grant_statements_sql( CASE WHEN bool_or(is_grantable) THEN ' WITH GRANT OPTION' ELSE '' END \ ) AS stmt \ FROM ( \ - SELECT c.relname, c.relkind, a.grantee, a.privilege_type, a.is_grantable, grantee.rolname \ - FROM pg_catalog.pg_class c \ - JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ - JOIN LATERAL aclexplode(c.relacl) a ON true \ + SELECT a.relname, a.relkind, a.grantee, a.privilege_type, a.is_grantable, grantee.rolname \ + FROM ( \ + SELECT c.relname, c.relkind, (aclexplode(c.relacl)).* \ + FROM pg_catalog.pg_class c \ + JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ + WHERE n.nspname = {source_schema} AND (c.relkind IN ('v','m') OR ({table_filter} AND c.relkind IN ('r','p','f','S'))) \ + ) a \ LEFT JOIN pg_roles grantee ON grantee.oid = a.grantee \ - WHERE n.nspname = {source_schema} AND (c.relkind IN ('v','m') OR ({table_filter} AND c.relkind IN ('r','p','f','S'))) \ ) rels \ GROUP BY relname, relkind, grantee, rolname \ ), \ @@ -5811,14 +5852,16 @@ fn postgres_transfer_grant_statements_sql( CASE WHEN bool_or(is_grantable) THEN ' WITH GRANT OPTION' ELSE '' END \ ) AS stmt \ FROM ( \ - SELECT p.proname, {routine_kind} AS routine_kind, \ - pg_get_function_identity_arguments(p.oid) AS identity_args, a.grantee, a.privilege_type, \ + SELECT a.proname, a.routine_kind, a.identity_args, a.grantee, a.privilege_type, \ a.is_grantable, grantee.rolname \ - FROM pg_catalog.pg_proc p \ - JOIN pg_catalog.pg_namespace n ON n.oid = p.pronamespace \ - JOIN LATERAL aclexplode(p.proacl) a ON true \ + FROM ( \ + SELECT p.proname, {routine_kind} AS routine_kind, \ + pg_get_function_identity_arguments(p.oid) AS identity_args, (aclexplode(p.proacl)).* \ + FROM pg_catalog.pg_proc p \ + JOIN pg_catalog.pg_namespace n ON n.oid = p.pronamespace \ + WHERE n.nspname = {source_schema} AND {routine_filter} \ + ) a \ LEFT JOIN pg_roles grantee ON grantee.oid = a.grantee \ - WHERE n.nspname = {source_schema} AND {routine_filter} \ ) routines \ GROUP BY proname, routine_kind, identity_args, grantee, rolname \ ) \ @@ -6422,7 +6465,9 @@ where // Create table on target if requested if request.create_table { - if is_postgres_transfer_dialect(target_db_type) && !request.target_schema.trim().is_empty() { + if transfer_table_needs_inline_postgres_schema_ensure(source_db_type, target_db_type) + && !request.target_schema.trim().is_empty() + { let create_schema_sql = format!("CREATE SCHEMA IF NOT EXISTS {}", quote_identifier(&request.target_schema, target_db_type)); execute_on_pool(state, target_pool_key, &create_schema_sql) @@ -6951,7 +6996,8 @@ where return Ok(TransferObjectOutcome::default()); } - let has_prokind = postgres_transfer_has_prokind(state, source_pool_key).await?; + let catalog_capabilities = postgres_transfer_catalog_capabilities(state, source_pool_key).await?; + let has_prokind = catalog_capabilities.has_prokind; let mut outcome = TransferObjectOutcome::default(); let object_sources = filter_object_sources_by_selection( get_postgres_schema_object_sources_for_transfer(state, source_pool_key, &request.source_schema, has_prokind) @@ -6984,6 +7030,8 @@ where &request.source_schema, &request.target_schema, &request.tables, + catalog_capabilities.has_pg_policy, + catalog_capabilities.supports_policy_permissiveness, ) .await?; let relation_names = postgres_transfer_relation_names(request); @@ -8135,12 +8183,27 @@ mod tests { use super::*; #[test] - fn postgres_transfer_routine_catalog_probe_checks_prokind() { - let sql = postgres_transfer_routine_catalog_capability_sql(); + fn postgres_transfer_catalog_probe_checks_legacy_boundaries() { + let sql = postgres_transfer_catalog_capabilities_sql(); assert!(sql.contains("attrelid = 'pg_catalog.pg_proc'::regclass")); assert!(sql.contains("attname = 'prokind'")); assert!(sql.contains("NOT attisdropped")); + assert!(sql.contains("c.relname = 'pg_policy'")); + assert!(sql.contains("attname = 'polpermissive'")); + assert!(!sql.contains("'pg_catalog.pg_policy'::regclass")); + } + + #[test] + fn postgres_table_transfer_reuses_the_batch_schema_preflight() { + assert!(!transfer_table_needs_inline_postgres_schema_ensure( + &DatabaseType::Postgres, + &DatabaseType::Postgres + )); + assert!(transfer_table_needs_inline_postgres_schema_ensure( + &DatabaseType::Postgres, + &DatabaseType::Kingbase + )); } #[test] @@ -8178,9 +8241,13 @@ mod tests { let legacy = postgres_transfer_grant_statements_sql("public", "archive", &["items".into()], false); assert!(!legacy.contains("prokind")); - assert!(legacy.contains("'FUNCTION' AS routine_kind")); + assert!(legacy.contains("'FUNCTION'::text AS routine_kind")); assert!(legacy.contains("NOT p.proisagg")); assert!(legacy.contains("NOT p.proiswindow")); + assert!(!legacy.contains("LATERAL")); + assert!(legacy.contains("(aclexplode(n.nspacl)).*")); + assert!(legacy.contains("(aclexplode(c.relacl)).*")); + assert!(legacy.contains("(aclexplode(p.proacl)).*")); } #[test] @@ -8228,6 +8295,17 @@ mod tests { assert!(sql.contains("pg_sequence_last_value(c.oid)::text")); } + #[test] + fn postgres_owned_sequence_queries_support_pre_ten_catalogs() { + assert!(!POSTGRES_OWNED_SEQUENCES_SQL.contains("pg_sequence")); + assert!(!POSTGRES_SEQUENCE_SNAPSHOTS_SQL.contains("pg_sequence")); + for sql in [POSTGRES_OWNED_SEQUENCES_SQL, POSTGRES_SEQUENCE_SNAPSHOTS_SQL] { + assert!(sql.contains("c.relkind = 'S'")); + assert!(sql.contains("pg_depend")); + assert!(sql.contains("d.deptype IN ('a', 'i')")); + } + } + #[test] fn postgres_selected_sequence_ddl_preserves_definition_and_value() { let sequence = PostgresTransferSequence { @@ -9508,8 +9586,7 @@ mod tests { owner_table: "it_quick_entry".to_string(), owner_column: "id".to_string(), }; - let create_sql = - format!("CREATE SEQUENCE IF NOT EXISTS {}", postgres_sequence_qualified_name("public", &sequence.name)); + let create_sql = format!("CREATE SEQUENCE {}", postgres_sequence_qualified_name("public", &sequence.name)); let owner_sql = format!( "ALTER SEQUENCE {} OWNED BY {}.{}", postgres_sequence_qualified_name("public", &sequence.name), @@ -9517,7 +9594,7 @@ mod tests { quote_identifier(&sequence.owner_column, &DatabaseType::Postgres) ); - assert_eq!(create_sql, "CREATE SEQUENCE IF NOT EXISTS \"public\".\"it_quick_entry_id_seq\"".to_string()); + assert_eq!(create_sql, "CREATE SEQUENCE \"public\".\"it_quick_entry_id_seq\"".to_string()); assert_eq!( owner_sql, "ALTER SEQUENCE \"public\".\"it_quick_entry_id_seq\" OWNED BY \"public\".\"it_quick_entry\".\"id\"" @@ -9590,8 +9667,7 @@ mod tests { owner_table: "it_quick_entry".to_string(), owner_column: "id".to_string(), }; - let create_sql = - format!("CREATE SEQUENCE IF NOT EXISTS {}", postgres_sequence_qualified_name("archive", &sequence.name)); + let create_sql = format!("CREATE SEQUENCE {}", postgres_sequence_qualified_name("archive", &sequence.name)); let owner_sql = format!( "ALTER SEQUENCE {} OWNED BY {}.{}", postgres_sequence_qualified_name("archive", &sequence.name), @@ -9604,7 +9680,7 @@ mod tests { assert!(!source_ddl.contains("CREATE SEQUENCE")); assert!(rewritten.contains("CREATE TABLE \"archive\".\"it_quick_entry\"")); assert!(rewritten.contains("nextval('\"archive\".it_quick_entry_id_seq'::regclass)")); - assert_eq!(create_sql, "CREATE SEQUENCE IF NOT EXISTS \"archive\".\"it_quick_entry_id_seq\"".to_string()); + assert_eq!(create_sql, "CREATE SEQUENCE \"archive\".\"it_quick_entry_id_seq\"".to_string()); assert_eq!( owner_sql, "ALTER SEQUENCE \"archive\".\"it_quick_entry_id_seq\" OWNED BY \"archive\".\"it_quick_entry\".\"id\""