github/dbx
4
0 Fork 0
mirror of https://github.com/t8y2/dbx.git synced 2026-11-15 22:53:23 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(transfer): complete PostgreSQL 9.x table transfer

Browse Source 
Closes #6070
This commit is contained in:
t8y2 committed 2026-08-13 16:13:07 +08:00
1 parent a0354726b1
commit 8548ecacc0
1 file changed
+161 -85
Download Patch File Download Diff File Expand all files Collapse all files
crates/dbx-core/src/transfer.rs
+161 -85
View File
@@ -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<String> {
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<bool, String> {
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<bool> {
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<PostgresTransferCatalogCapabilities, String> {
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<Vec<String>, 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::<Vec<_>>().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\""