From 2e431d6e330406ff96861756e764a1ca306140b2 Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Fri, 29 May 2026 00:49:29 +0800 Subject: [PATCH] feat(transfer): improve postgres cross-version migration --- crates/dbx-core/src/connection.rs | 5 +- crates/dbx-core/src/table_structure_sql.rs | 24 +- crates/dbx-core/src/transfer.rs | 1440 ++++++++++++++++- .../dbx-core/tests/live_postgres_transfer.rs | 353 ++++ src-tauri/src/commands/transfer.rs | 98 ++ 5 files changed, 1887 insertions(+), 33 deletions(-) create mode 100644 crates/dbx-core/tests/live_postgres_transfer.rs diff --git a/crates/dbx-core/src/connection.rs b/crates/dbx-core/src/connection.rs index a788af5a8..101c8ae5d 100644 --- a/crates/dbx-core/src/connection.rs +++ b/crates/dbx-core/src/connection.rs @@ -1492,10 +1492,7 @@ mod tests { config.port = 0; state.configs.write().await.insert(config.id.clone(), config.clone()); - let pool_key = state - .get_or_create_pool_for_session("duckdb-conn", Some("main"), Some("tab-1")) - .await - .unwrap(); + let pool_key = state.get_or_create_pool_for_session("duckdb-conn", Some("main"), Some("tab-1")).await.unwrap(); assert_eq!(pool_key, "duckdb-conn:session:tab-1"); assert!(state.close_client_session_pool("duckdb-conn", Some("main"), "tab-1").await.unwrap()); diff --git a/crates/dbx-core/src/table_structure_sql.rs b/crates/dbx-core/src/table_structure_sql.rs index 15786fc5f..e57de5b08 100644 --- a/crates/dbx-core/src/table_structure_sql.rs +++ b/crates/dbx-core/src/table_structure_sql.rs @@ -2227,11 +2227,7 @@ mod tests { col.data_type = "int".to_string(); col.is_nullable = false; col.is_primary_key = true; - col.extra = Some(ColumnExtra { - auto_increment: Some(true), - on_update_current_timestamp: None, - identity: None, - }); + col.extra = Some(ColumnExtra { auto_increment: Some(true), on_update_current_timestamp: None, identity: None }); let result = build_create_table_sql(TableStructureSqlOptions { database_type: Some(DatabaseType::Mysql), @@ -2254,11 +2250,7 @@ mod tests { col.data_type = "timestamp".to_string(); col.is_nullable = false; col.default_value = "CURRENT_TIMESTAMP".to_string(); - col.extra = Some(ColumnExtra { - auto_increment: None, - on_update_current_timestamp: Some(true), - identity: None, - }); + col.extra = Some(ColumnExtra { auto_increment: None, on_update_current_timestamp: Some(true), identity: None }); let result = build_create_table_sql(TableStructureSqlOptions { database_type: Some(DatabaseType::Mysql), @@ -2282,11 +2274,7 @@ mod tests { col.extra = Some(ColumnExtra { auto_increment: None, on_update_current_timestamp: None, - identity: Some(ColumnIdentity { - generation: Some("BY DEFAULT".to_string()), - seed: None, - increment: None, - }), + identity: Some(ColumnIdentity { generation: Some("BY DEFAULT".to_string()), seed: None, increment: None }), }); let result = build_create_table_sql(TableStructureSqlOptions { @@ -2311,11 +2299,7 @@ mod tests { col.extra = Some(ColumnExtra { auto_increment: Some(true), on_update_current_timestamp: None, - identity: Some(ColumnIdentity { - generation: None, - seed: Some(100), - increment: Some(5), - }), + identity: Some(ColumnIdentity { generation: None, seed: Some(100), increment: Some(5) }), }); let result = build_create_table_sql(TableStructureSqlOptions { diff --git a/crates/dbx-core/src/transfer.rs b/crates/dbx-core/src/transfer.rs index ae12e6d75..70fdf0017 100644 --- a/crates/dbx-core/src/transfer.rs +++ b/crates/dbx-core/src/transfer.rs @@ -1,10 +1,12 @@ +use regex::Regex; use serde::{Deserialize, Serialize}; -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use tokio::sync::RwLock; use crate::connection::{AppState, PoolKind}; use crate::db; use crate::models::connection::DatabaseType; +use crate::object_source_sql::{build_executable_object_source_statements, EditableObjectSourceSqlInput}; use crate::query::{agent_execute_query_params, QueryExecutionOptions}; use crate::sql::starts_with_executable_sql_keyword; @@ -79,6 +81,371 @@ pub fn qualified_table(table: &str, schema: &str, db_type: &DatabaseType) -> Str } } +fn quote_string_literal(value: &str) -> String { + format!("'{}'", value.replace('\'', "''")) +} + +fn is_simple_identifier(value: &str) -> bool { + let mut chars = value.chars(); + let Some(first) = chars.next() else { + return false; + }; + if !(first == '_' || first.is_ascii_alphabetic()) { + return false; + } + chars.all(|ch| ch == '_' || ch.is_ascii_alphanumeric()) +} + +fn is_postgres_compat_transfer(source_db: &DatabaseType, target_db: &DatabaseType) -> bool { + matches!(source_db, DatabaseType::Postgres) && matches!(target_db, DatabaseType::Postgres) +} + +fn is_postgres_integer_like_type(data_type: &str) -> bool { + let normalized = data_type.trim().to_ascii_lowercase(); + matches!( + normalized.split(['(', ' ']).next().unwrap_or(""), + "smallint" | "integer" | "bigint" | "int2" | "int4" | "int8" + ) +} + +fn is_postgres_sequence_default(default_value: Option<&str>) -> bool { + default_value.is_some_and(|value| value.to_ascii_lowercase().contains("nextval(")) +} + +fn rewrite_postgres_schema_qualified_references(input: &str, source_schema: &str, target_schema: &str) -> String { + if source_schema.trim().is_empty() || source_schema == target_schema { + return input.to_string(); + } + + let quoted_source = format!("{}.", quote_identifier(source_schema, &DatabaseType::Postgres)); + let quoted_target = format!("{}.", quote_identifier(target_schema, &DatabaseType::Postgres)); + let rewritten = input.replace("ed_source, "ed_target); + let unquoted_pattern = + Regex::new(&format!(r#"(^|[^"\w]){}\."#, regex::escape(source_schema))).expect("valid postgres schema regex"); + unquoted_pattern + .replace_all(&rewritten, |captures: ®ex::Captures| format!("{}{}", &captures[1], quoted_target)) + .into_owned() +} + +fn postgres_column_type_sql( + column: &db::ColumnInfo, + source_schema: &str, + target_schema: &str, + source_db: &DatabaseType, + target_db: &DatabaseType, +) -> String { + if is_postgres_compat_transfer(source_db, target_db) { + let trimmed = column.data_type.trim(); + if !trimmed.is_empty() { + return rewrite_postgres_schema_qualified_references(trimmed, source_schema, target_schema); + } + } + map_column_type(&column.data_type, source_db, target_db) +} + +fn postgres_default_clause( + column: &db::ColumnInfo, + source_schema: &str, + target_schema: &str, + source_db: &DatabaseType, + target_db: &DatabaseType, +) -> Option { + if !is_postgres_compat_transfer(source_db, target_db) { + return None; + } + let default_value = column.column_default.as_deref()?.trim(); + if default_value.is_empty() { + return None; + } + if is_postgres_sequence_default(Some(default_value)) && is_postgres_integer_like_type(&column.data_type) { + return Some("GENERATED BY DEFAULT AS IDENTITY".to_string()); + } + Some(format!( + "DEFAULT {}", + rewrite_postgres_schema_qualified_references(default_value, source_schema, target_schema) + )) +} + +fn postgres_order_by_expression(columns: &[String], db_type: &DatabaseType) -> Option { + if columns.is_empty() { + return None; + } + Some(columns.iter().map(|column| quote_identifier(column, db_type)).collect::>().join(", ")) +} + +fn postgres_index_column_sql(column: &str) -> String { + if is_simple_identifier(column) { + quote_identifier(column, &DatabaseType::Postgres) + } else { + column.to_string() + } +} + +fn generate_postgres_index_ddl(indexes: &[db::IndexInfo], table: &str, schema: &str) -> Vec { + let full_table = qualified_table(table, schema, &DatabaseType::Postgres); + let mut statements = Vec::new(); + for index in indexes.iter().filter(|index| !index.is_primary) { + if index.name.trim().is_empty() || index.columns.is_empty() { + continue; + } + let unique = if index.is_unique { "UNIQUE " } else { "" }; + let using_clause = index + .index_type + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(|value| format!(" USING {value}")) + .unwrap_or_default(); + let columns = + index.columns.iter().map(|column| postgres_index_column_sql(column)).collect::>().join(", "); + let include_clause = index + .included_columns + .as_ref() + .filter(|columns| !columns.is_empty()) + .map(|columns| { + format!( + " INCLUDE ({})", + columns + .iter() + .map(|column| quote_identifier(column, &DatabaseType::Postgres)) + .collect::>() + .join(", ") + ) + }) + .unwrap_or_default(); + let filter_clause = index + .filter + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(|value| format!(" WHERE {value}")) + .unwrap_or_default(); + statements.push(format!( + "CREATE {unique}INDEX IF NOT EXISTS {} ON {full_table}{using_clause} ({columns}){include_clause}{filter_clause}", + quote_identifier(&index.name, &DatabaseType::Postgres) + )); + if let Some(comment) = index.comment.as_deref().map(str::trim).filter(|value| !value.is_empty()) { + let qualified_index = if schema.is_empty() { + quote_identifier(&index.name, &DatabaseType::Postgres) + } else { + format!( + "{}.{}", + quote_identifier(schema, &DatabaseType::Postgres), + quote_identifier(&index.name, &DatabaseType::Postgres) + ) + }; + statements.push(format!("COMMENT ON INDEX {qualified_index} IS {}", quote_string_literal(comment))); + } + } + statements +} + +fn generate_postgres_foreign_key_ddl(foreign_keys: &[db::ForeignKeyInfo], table: &str, schema: &str) -> Vec { + let full_table = qualified_table(table, schema, &DatabaseType::Postgres); + let mut grouped: HashMap<&str, Vec<&db::ForeignKeyInfo>> = HashMap::new(); + let mut order = Vec::new(); + + for foreign_key in foreign_keys { + if !grouped.contains_key(foreign_key.name.as_str()) { + order.push(foreign_key.name.as_str()); + } + grouped.entry(foreign_key.name.as_str()).or_default().push(foreign_key); + } + + let mut statements = Vec::new(); + for name in order { + let Some(group) = grouped.get(name) else { + continue; + }; + let columns = group + .iter() + .map(|foreign_key| quote_identifier(&foreign_key.column, &DatabaseType::Postgres)) + .collect::>() + .join(", "); + let ref_columns = group + .iter() + .map(|foreign_key| quote_identifier(&foreign_key.ref_column, &DatabaseType::Postgres)) + .collect::>() + .join(", "); + let referenced_table = qualified_table(&group[0].ref_table, schema, &DatabaseType::Postgres); + statements.push(format!( + "ALTER TABLE {full_table} ADD CONSTRAINT {} FOREIGN KEY ({columns}) REFERENCES {referenced_table} ({ref_columns})", + quote_identifier(name, &DatabaseType::Postgres) + )); + } + + statements +} + +fn generate_postgres_sequence_sync_sql(columns: &[db::ColumnInfo], table: &str, schema: &str) -> Vec { + let full_table = qualified_table(table, schema, &DatabaseType::Postgres); + columns + .iter() + .filter(|column| is_postgres_sequence_default(column.column_default.as_deref())) + .map(|column| { + let quoted_column = quote_identifier(&column.name, &DatabaseType::Postgres); + format!( + "SELECT setval(pg_get_serial_sequence({}, {}), GREATEST(COALESCE(MAX({quoted_column}), 0), 1), MAX({quoted_column}) IS NOT NULL) FROM {full_table}", + quote_string_literal(&full_table), + quote_string_literal(&column.name) + ) + }) + .collect() +} + +#[derive(Debug, Clone)] +struct PostgresTriggerSource { + table_name: String, + trigger_name: String, + source: String, +} + +#[derive(Debug, Clone)] +struct PostgresExtensionSource { + extension_name: String, +} + +#[derive(Debug, Clone)] +struct PostgresEnumSource { + type_name: String, + labels: Vec, +} + +#[derive(Debug, Clone)] +struct PostgresDomainSource { + domain_name: String, + base_type: String, + default_value: Option, + not_null: bool, + checks: Vec, +} + +#[derive(Debug, Clone)] +struct PostgresMaterializedViewSource { + view_name: String, + source: String, +} + +fn json_string_cell(row: &[serde_json::Value], index: usize) -> Option { + row.get(index).and_then(|value| value.as_str().map(str::to_string)) +} + +fn result_rows_to_string_statements(rows: Vec>) -> Vec { + rows.into_iter().filter_map(|row| json_string_cell(&row, 0)).filter(|stmt| !stmt.trim().is_empty()).collect() +} + +fn ensure_sql_statement_terminated(sql: &str) -> String { + let trimmed = sql.trim(); + if trimmed.ends_with(';') { + trimmed.to_string() + } else { + format!("{trimmed};") + } +} + +fn generate_postgres_extension_ddl(extension: &PostgresExtensionSource, target_schema: &str) -> String { + format!( + "CREATE EXTENSION IF NOT EXISTS {} WITH SCHEMA {}", + quote_identifier(&extension.extension_name, &DatabaseType::Postgres), + quote_identifier(target_schema, &DatabaseType::Postgres) + ) +} + +fn generate_postgres_enum_ddl(enum_type: &PostgresEnumSource, target_schema: &str) -> String { + let labels = enum_type.labels.iter().map(|label| quote_string_literal(label)).collect::>().join(", "); + let create_sql = format!( + "CREATE TYPE {}.{} AS ENUM ({labels})", + quote_identifier(target_schema, &DatabaseType::Postgres), + quote_identifier(&enum_type.type_name, &DatabaseType::Postgres) + ); + format!( + "DO $$ BEGIN IF NOT EXISTS (SELECT 1 FROM pg_type t JOIN pg_namespace n ON n.oid = t.typnamespace WHERE n.nspname = {} AND t.typname = {}) THEN {create_sql}; END IF; END $$", + quote_string_literal(target_schema), + quote_string_literal(&enum_type.type_name) + ) +} + +fn generate_postgres_domain_ddl(domain: &PostgresDomainSource, target_schema: &str) -> String { + let mut create_sql = format!( + "CREATE DOMAIN {}.{} AS {}", + quote_identifier(target_schema, &DatabaseType::Postgres), + quote_identifier(&domain.domain_name, &DatabaseType::Postgres), + domain.base_type + ); + if let Some(default_value) = domain.default_value.as_deref().map(str::trim).filter(|value| !value.is_empty()) { + create_sql.push_str(&format!(" DEFAULT {default_value}")); + } + if domain.not_null { + create_sql.push_str(" NOT NULL"); + } + for check in &domain.checks { + create_sql.push(' '); + create_sql.push_str(check); + } + format!( + "DO $$ BEGIN IF NOT EXISTS (SELECT 1 FROM pg_type t JOIN pg_namespace n ON n.oid = t.typnamespace WHERE n.nspname = {} AND t.typname = {}) THEN {}; END IF; END $$", + quote_string_literal(target_schema), + quote_string_literal(&domain.domain_name), + create_sql + ) +} + +fn generate_postgres_materialized_view_ddls(view: &PostgresMaterializedViewSource, target_schema: &str) -> Vec { + let qualified_name = qualified_table(&view.view_name, target_schema, &DatabaseType::Postgres); + vec![ + format!("DROP MATERIALIZED VIEW IF EXISTS {qualified_name}"), + format!("CREATE MATERIALIZED VIEW {qualified_name} AS\n{}", ensure_sql_statement_terminated(&view.source)), + ] +} + +fn rewrite_postgres_routine_schema(source: &str, target_schema: &str) -> Option { + let re = Regex::new( + r#"(?is)^(\s*CREATE\s+(?:OR\s+REPLACE\s+)?(?:(?:NON)?EDITIONABLE\s+)?(?:FUNCTION|PROCEDURE)\s+)((?:"(?:""|[^"])+"|[A-Za-z_][\w$]*)(?:\s*\.\s*(?:"(?:""|[^"])+"|[A-Za-z_][\w$]*))?)"#, + ) + .ok()?; + let captures = re.captures(source)?; + let full = captures.get(0)?; + let prefix = captures.get(1)?.as_str(); + let existing_name = captures.get(2)?.as_str(); + let name_re = Regex::new(r#""(?:""|[^"])+"|[A-Za-z_][\w$]*"#).ok()?; + let parts = name_re + .find_iter(existing_name) + .map(|part| part.as_str().trim().trim_matches('"').replace("\"\"", "\"")) + .collect::>(); + let name = parts.last()?; + let replacement = format!( + "{}.{}", + quote_identifier(target_schema, &DatabaseType::Postgres), + quote_identifier(name, &DatabaseType::Postgres) + ); + Some(format!("{}{}{}{}", &source[..full.start()], prefix, replacement, &source[full.end()..])) +} + +fn rewrite_postgres_trigger_table_schema( + source: &str, + source_schema: &str, + table_name: &str, + target_schema: &str, +) -> String { + let qualified_target_table = qualified_table(table_name, target_schema, &DatabaseType::Postgres); + let candidate_patterns = [ + format!( + " ON {}.{} ", + quote_identifier(source_schema, &DatabaseType::Postgres), + quote_identifier(table_name, &DatabaseType::Postgres) + ), + format!(" ON {source_schema}.{table_name} "), + format!(" ON {} ", quote_identifier(table_name, &DatabaseType::Postgres)), + format!(" ON {table_name} "), + ]; + for pattern in candidate_patterns { + if source.contains(&pattern) { + return source.replacen(&pattern, &format!(" ON {qualified_target_table} "), 1); + } + } + source.to_string() +} + pub fn escape_value(val: &serde_json::Value, db_type: &DatabaseType) -> String { escape_value_typed(val, db_type, None) } @@ -397,6 +764,7 @@ fn mysql_type_needs_key_prefix(mapped_type: &str) -> bool { pub fn generate_create_table_ddl( columns: &[db::ColumnInfo], table: &str, + source_schema: &str, schema: &str, target_db: &DatabaseType, source_db: &DatabaseType, @@ -416,8 +784,12 @@ pub fn generate_create_table_ddl( let mut col_lines = Vec::with_capacity(columns.len()); for c in columns { col_lines.push({ - let mapped_type = map_column_type(&c.data_type, source_db, target_db); + let mapped_type = postgres_column_type_sql(c, source_schema, schema, source_db, target_db); let mut line = format!(" {} {}", quote_identifier(&c.name, target_db), mapped_type); + if let Some(default_clause) = postgres_default_clause(c, source_schema, schema, source_db, target_db) { + line.push(' '); + line.push_str(&default_clause); + } if !c.is_nullable { line.push_str(" NOT NULL"); } @@ -763,6 +1135,33 @@ pub fn pagination_sql( } } +pub fn pagination_sql_with_order( + columns: &[String], + table: &str, + schema: &str, + db_type: &DatabaseType, + offset: u64, + limit: usize, + order_by_columns: &[String], +) -> String { + let full_table = qualified_table(table, schema, db_type); + let col_list = columns.iter().map(|c| quote_identifier(c, db_type)).collect::>().join(", "); + let order_expression = postgres_order_by_expression(order_by_columns, db_type); + + match db_type { + DatabaseType::SqlServer | DatabaseType::Oracle => { + let order_by = order_expression.unwrap_or_else(|| "(SELECT NULL)".to_string()); + format!( + "SELECT {col_list} FROM {full_table} ORDER BY {order_by} OFFSET {offset} ROWS FETCH NEXT {limit} ROWS ONLY" + ) + } + _ => { + let order_by = order_expression.map(|value| format!(" ORDER BY {value}")).unwrap_or_default(); + format!("SELECT {col_list} FROM {full_table}{order_by} LIMIT {limit} OFFSET {offset}") + } + } +} + pub fn count_sql(table: &str, schema: &str, db_type: &DatabaseType) -> String { let full_table = qualified_table(table, schema, db_type); format!("SELECT COUNT(*) FROM {full_table}") @@ -991,6 +1390,413 @@ pub async fn get_columns_for_transfer( } } +async fn get_postgres_indexes_for_transfer( + state: &AppState, + pool_key: &str, + schema: &str, + table: &str, +) -> Result, String> { + let connections = state.connections.read().await; + let Some(PoolKind::Postgres(pool)) = connections.get(pool_key) else { + return Err("PostgreSQL pool not found".to_string()); + }; + let pool = pool.clone(); + drop(connections); + db::postgres::list_indexes(&pool, schema, table).await +} + +async fn get_postgres_foreign_keys_for_transfer( + state: &AppState, + pool_key: &str, + schema: &str, + table: &str, +) -> Result, String> { + let connections = state.connections.read().await; + let Some(PoolKind::Postgres(pool)) = connections.get(pool_key) else { + return Err("PostgreSQL pool not found".to_string()); + }; + let pool = pool.clone(); + drop(connections); + db::postgres::list_foreign_keys(&pool, schema, table).await +} + +async fn get_postgres_schema_object_sources_for_transfer( + state: &AppState, + pool_key: &str, + schema: &str, +) -> Result, String> { + let views_sql = format!( + "SELECT c.relname, pg_get_viewdef(c.oid, true) \ + FROM pg_catalog.pg_class c \ + JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ + WHERE n.nspname = {} AND c.relkind = 'v' \ + ORDER BY c.relname", + quote_string_literal(schema) + ); + let routines_sql = format!( + "SELECT p.proname, CASE p.prokind WHEN 'p' THEN 'PROCEDURE' ELSE 'FUNCTION' END, pg_get_functiondef(p.oid) \ + FROM pg_catalog.pg_proc p \ + JOIN pg_catalog.pg_namespace n ON n.oid = p.pronamespace \ + WHERE n.nspname = {} AND p.prokind IN ('p', 'f') \ + ORDER BY CASE p.prokind WHEN 'p' THEN 0 ELSE 1 END, p.proname, p.oid", + quote_string_literal(schema) + ); + + let mut sources = Vec::new(); + for row in execute_on_pool(state, pool_key, &views_sql).await?.rows { + let Some(name) = json_string_cell(&row, 0) else { + continue; + }; + let Some(source) = json_string_cell(&row, 1) else { + continue; + }; + sources.push(db::ObjectSource { + name, + object_type: db::ObjectSourceKind::View, + schema: Some(schema.to_string()), + source, + }); + } + for row in execute_on_pool(state, pool_key, &routines_sql).await?.rows { + let Some(name) = json_string_cell(&row, 0) else { + continue; + }; + let kind = match json_string_cell(&row, 1).as_deref() { + Some("PROCEDURE") => db::ObjectSourceKind::Procedure, + _ => db::ObjectSourceKind::Function, + }; + let Some(source) = json_string_cell(&row, 2) else { + continue; + }; + sources.push(db::ObjectSource { name, object_type: kind, schema: Some(schema.to_string()), source }); + } + + Ok(sources) +} + +async fn get_postgres_materialized_view_sources_for_transfer( + state: &AppState, + pool_key: &str, + schema: &str, +) -> Result, String> { + let sql = format!( + "SELECT c.relname, pg_get_viewdef(c.oid, true) \ + FROM pg_catalog.pg_class c \ + JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ + WHERE n.nspname = {} AND c.relkind = 'm' \ + ORDER BY c.relname", + quote_string_literal(schema) + ); + let rows = execute_on_pool(state, pool_key, &sql).await?.rows; + Ok(rows + .into_iter() + .filter_map(|row| { + Some(PostgresMaterializedViewSource { + view_name: json_string_cell(&row, 0)?, + source: json_string_cell(&row, 1)?, + }) + }) + .collect()) +} + +async fn get_postgres_trigger_sources_for_transfer( + state: &AppState, + pool_key: &str, + schema: &str, + tables: &[String], +) -> Result, String> { + if tables.is_empty() { + return Ok(Vec::new()); + } + let table_list = tables.iter().map(|table| quote_string_literal(table)).collect::>().join(", "); + let sql = format!( + "SELECT c.relname, t.tgname, pg_get_triggerdef(t.oid, true) \ + FROM pg_catalog.pg_trigger t \ + JOIN pg_catalog.pg_class c ON c.oid = t.tgrelid \ + JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ + WHERE n.nspname = {} AND NOT t.tgisinternal AND c.relname IN ({table_list}) \ + ORDER BY c.relname, t.tgname", + quote_string_literal(schema) + ); + let rows = execute_on_pool(state, pool_key, &sql).await?.rows; + Ok(rows + .into_iter() + .filter_map(|row| { + Some(PostgresTriggerSource { + table_name: json_string_cell(&row, 0)?, + trigger_name: json_string_cell(&row, 1)?, + source: json_string_cell(&row, 2)?, + }) + }) + .collect()) +} + +async fn get_postgres_extension_sources_for_transfer( + state: &AppState, + pool_key: &str, + schema: &str, +) -> Result, String> { + let sql = format!( + "SELECT e.extname \ + FROM pg_extension e \ + JOIN pg_namespace n ON n.oid = e.extnamespace \ + WHERE n.nspname = {} \ + ORDER BY e.extname", + quote_string_literal(schema) + ); + let rows = execute_on_pool(state, pool_key, &sql).await?.rows; + Ok(rows + .into_iter() + .filter_map(|row| json_string_cell(&row, 0).map(|extension_name| PostgresExtensionSource { extension_name })) + .collect()) +} + +async fn get_postgres_enum_sources_for_transfer( + state: &AppState, + pool_key: &str, + schema: &str, +) -> Result, String> { + let sql = format!( + "SELECT t.typname, COALESCE(array_to_json(array_agg(e.enumlabel ORDER BY e.enumsortorder))::text, '[]') \ + FROM pg_type t \ + JOIN pg_namespace n ON n.oid = t.typnamespace \ + LEFT JOIN pg_enum e ON e.enumtypid = t.oid \ + WHERE n.nspname = {} AND t.typtype = 'e' \ + GROUP BY t.typname \ + ORDER BY t.typname", + quote_string_literal(schema) + ); + let rows = execute_on_pool(state, pool_key, &sql).await?.rows; + Ok(rows + .into_iter() + .filter_map(|row| { + let type_name = json_string_cell(&row, 0)?; + let labels_json = json_string_cell(&row, 1)?; + let labels = serde_json::from_str::>(&labels_json).ok()?; + Some(PostgresEnumSource { type_name, labels }) + }) + .collect()) +} + +async fn get_postgres_domain_sources_for_transfer( + state: &AppState, + pool_key: &str, + schema: &str, +) -> Result, String> { + let sql = format!( + "SELECT t.typname, \ + pg_catalog.format_type(t.typbasetype, t.typtypmod), \ + NULLIF(t.typdefault, ''), \ + t.typnotnull, \ + COALESCE(( \ + SELECT array_to_json(array_agg(pg_get_constraintdef(c.oid, true) ORDER BY c.conname))::text \ + FROM pg_constraint c \ + WHERE c.contypid = t.oid AND c.contype = 'c' \ + ), '[]') \ + FROM pg_type t \ + JOIN pg_namespace n ON n.oid = t.typnamespace \ + WHERE n.nspname = {} AND t.typtype = 'd' \ + ORDER BY t.typname", + quote_string_literal(schema) + ); + let rows = execute_on_pool(state, pool_key, &sql).await?.rows; + Ok(rows + .into_iter() + .filter_map(|row| { + let domain_name = json_string_cell(&row, 0)?; + let base_type = json_string_cell(&row, 1)?; + let default_value = json_string_cell(&row, 2); + let not_null = row.get(3).and_then(|value| value.as_bool()).unwrap_or(false); + let checks = json_string_cell(&row, 4) + .and_then(|json| serde_json::from_str::>(&json).ok()) + .unwrap_or_default(); + Some(PostgresDomainSource { domain_name, base_type, default_value, not_null, checks }) + }) + .collect()) +} + +async fn get_postgres_policy_statements_for_transfer( + state: &AppState, + pool_key: &str, + source_schema: &str, + target_schema: &str, + tables: &[String], +) -> Result, String> { + if tables.is_empty() { + return Ok(Vec::new()); + } + let table_list = tables.iter().map(|table| quote_string_literal(table)).collect::>().join(", "); + let sql = format!( + "WITH selected_tables AS ( \ + SELECT c.oid, c.relname, c.relrowsecurity, c.relforcerowsecurity \ + 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 ('r','p') AND c.relname IN ({table_list}) \ + ), \ + policy_rows AS ( \ + SELECT t.relname, t.relrowsecurity, t.relforcerowsecurity, p.polname, p.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, \ + pg_get_expr(p.polwithcheck, p.polrelid) AS with_check_expr \ + FROM selected_tables t \ + JOIN pg_catalog.pg_policy p ON p.polrelid = t.oid \ + ) \ + SELECT stmt FROM ( \ + SELECT format('ALTER TABLE %I.%I ENABLE ROW LEVEL SECURITY', {target_schema}, relname) AS stmt, relname, 0 AS sort_order \ + FROM selected_tables WHERE relrowsecurity \ + UNION ALL \ + SELECT format('ALTER TABLE %I.%I FORCE ROW LEVEL SECURITY', {target_schema}, relname) AS stmt, relname, 1 AS sort_order \ + FROM selected_tables WHERE relforcerowsecurity \ + UNION ALL \ + SELECT format('DROP POLICY IF EXISTS %I ON %I.%I', polname, {target_schema}, relname) AS stmt, relname, 2 AS sort_order \ + FROM policy_rows \ + UNION ALL \ + SELECT format( \ + 'CREATE POLICY %I ON %I.%I AS %s FOR %s%s%s%s', \ + polname, {target_schema}, relname, \ + CASE WHEN polpermissive THEN 'PERMISSIVE' ELSE 'RESTRICTIVE' END, \ + 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, \ + CASE WHEN with_check_expr IS NOT NULL THEN ' WITH CHECK (' || with_check_expr || ')' ELSE '' END \ + ) AS stmt, relname, 3 AS sort_order \ + FROM policy_rows \ + ) statements \ + ORDER BY relname, sort_order, stmt", + source_schema = quote_string_literal(source_schema), + target_schema = quote_string_literal(target_schema), + ); + Ok(result_rows_to_string_statements(execute_on_pool(state, pool_key, &sql).await?.rows)) +} + +async fn get_postgres_ownership_statements_for_transfer( + state: &AppState, + pool_key: &str, + source_schema: &str, + target_schema: &str, + tables: &[String], +) -> Result, String> { + let table_list = tables.iter().map(|table| quote_string_literal(table)).collect::>().join(", "); + let table_filter = if tables.is_empty() { "FALSE".to_string() } else { format!("c.relname IN ({table_list})") }; + let sql = format!( + "WITH relation_owners AS ( \ + SELECT CASE c.relkind \ + WHEN 'm' THEN format('ALTER MATERIALIZED VIEW %I.%I OWNER TO %I', {target_schema}, c.relname, pg_get_userbyid(c.relowner)) \ + WHEN 'v' THEN format('ALTER VIEW %I.%I OWNER TO %I', {target_schema}, c.relname, pg_get_userbyid(c.relowner)) \ + WHEN 'f' THEN format('ALTER FOREIGN TABLE %I.%I OWNER TO %I', {target_schema}, c.relname, pg_get_userbyid(c.relowner)) \ + WHEN 'S' THEN format('ALTER SEQUENCE %I.%I OWNER TO %I', {target_schema}, c.relname, pg_get_userbyid(c.relowner)) \ + ELSE format('ALTER TABLE %I.%I OWNER TO %I', {target_schema}, c.relname, pg_get_userbyid(c.relowner)) \ + END AS stmt \ + 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'))) \ + ), \ + routine_owners AS ( \ + SELECT format('ALTER %s %I.%I(%s) OWNER TO %I', \ + CASE p.prokind WHEN 'p' THEN 'PROCEDURE' ELSE 'FUNCTION' END, \ + {target_schema}, p.proname, pg_get_function_identity_arguments(p.oid), pg_get_userbyid(p.proowner)) AS stmt \ + FROM pg_catalog.pg_proc p \ + JOIN pg_catalog.pg_namespace n ON n.oid = p.pronamespace \ + WHERE n.nspname = {source_schema} AND p.prokind IN ('p','f') \ + ), \ + type_owners AS ( \ + SELECT format('ALTER %s %I.%I OWNER TO %I', \ + CASE t.typtype WHEN 'd' THEN 'DOMAIN' ELSE 'TYPE' END, \ + {target_schema}, t.typname, pg_get_userbyid(t.typowner)) AS stmt \ + FROM pg_catalog.pg_type t \ + JOIN pg_catalog.pg_namespace n ON n.oid = t.typnamespace \ + WHERE n.nspname = {source_schema} AND t.typtype IN ('e','d') \ + ) \ + SELECT stmt FROM ( \ + SELECT format('ALTER SCHEMA %I OWNER TO %I', {target_schema}, pg_get_userbyid(n.nspowner)) AS stmt \ + FROM pg_catalog.pg_namespace n WHERE n.nspname = {source_schema} \ + UNION ALL SELECT stmt FROM relation_owners \ + UNION ALL SELECT stmt FROM routine_owners \ + UNION ALL SELECT stmt FROM type_owners \ + ) statements", + source_schema = quote_string_literal(source_schema), + target_schema = quote_string_literal(target_schema), + table_filter = table_filter, + ); + Ok(result_rows_to_string_statements(execute_on_pool(state, pool_key, &sql).await?.rows)) +} + +async fn get_postgres_grant_statements_for_transfer( + state: &AppState, + pool_key: &str, + source_schema: &str, + target_schema: &str, + tables: &[String], +) -> Result, String> { + let table_list = tables.iter().map(|table| quote_string_literal(table)).collect::>().join(", "); + let table_filter = if tables.is_empty() { "FALSE".to_string() } else { format!("c.relname IN ({table_list})") }; + let sql = format!( + "WITH schema_grants AS ( \ + SELECT format( \ + 'GRANT %s ON SCHEMA %I TO %s%s', \ + string_agg(a.privilege_type, ', ' ORDER BY a.privilege_type), \ + {target_schema}, \ + 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 \ + LEFT JOIN pg_roles grantee ON grantee.oid = a.grantee \ + WHERE n.nspname = {source_schema} \ + GROUP BY a.grantee, grantee.rolname \ + ), \ + relation_grants AS ( \ + SELECT format( \ + 'GRANT %s ON %s %I.%I TO %s%s', \ + string_agg(a.privilege_type, ', ' ORDER BY a.privilege_type), \ + CASE WHEN relkind = 'S' THEN 'SEQUENCE' ELSE 'TABLE' END, \ + {target_schema}, relname, \ + 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 ( \ + 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 \ + 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 \ + ), \ + routine_grants AS ( \ + SELECT format( \ + 'GRANT %s ON %s %I.%I(%s) TO %s%s', \ + string_agg(a.privilege_type, ', ' ORDER BY a.privilege_type), \ + CASE WHEN prokind = 'p' THEN 'PROCEDURE' ELSE 'FUNCTION' END, \ + {target_schema}, proname, identity_args, \ + 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 ( \ + SELECT p.proname, p.prokind, pg_get_function_identity_arguments(p.oid) AS 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 \ + LEFT JOIN pg_roles grantee ON grantee.oid = a.grantee \ + WHERE n.nspname = {source_schema} AND p.prokind IN ('p','f') \ + ) routines \ + GROUP BY proname, prokind, identity_args, grantee, rolname \ + ) \ + SELECT stmt FROM ( \ + SELECT stmt FROM schema_grants \ + UNION ALL SELECT stmt FROM relation_grants \ + UNION ALL SELECT stmt FROM routine_grants \ + ) statements \ + WHERE stmt IS NOT NULL", + source_schema = quote_string_literal(source_schema), + target_schema = quote_string_literal(target_schema), + table_filter = table_filter, + ); + Ok(result_rows_to_string_statements(execute_on_pool(state, pool_key, &sql).await?.rows)) +} + pub async fn is_cancelled(transfer_id: &str) -> bool { CANCELLED.read().await.contains(transfer_id) } @@ -1020,6 +1826,7 @@ where F: FnMut(TransferProgress), { let total_tables = request.tables.len(); + let pg_compat_transfer = is_postgres_compat_transfer(source_db_type, target_db_type); // Get source columns (deduplicate by name) let columns = { @@ -1042,6 +1849,8 @@ where let col_names: Vec = columns.iter().map(|c| c.name.clone()).collect(); let col_types: Vec> = columns.iter().map(|c| Some(c.data_type.clone())).collect(); + let primary_key_columns: Vec = + columns.iter().filter(|c| c.is_primary_key).map(|c| c.name.clone()).collect(); log::info!("[transfer] {} has {} columns, counting rows...", table, columns.len()); // Fetch source table comment @@ -1059,6 +1868,29 @@ where .next() .and_then(|t| t.comment); + let target_table_preexisting = crate::schema::list_tables_core( + state, + &request.target_connection_id, + &request.target_database, + &request.target_schema, + Some(table), + Some(1), + ) + .await + .map(|tables| !tables.is_empty()) + .unwrap_or(false); + + let source_indexes = if request.create_table && pg_compat_transfer && !target_table_preexisting { + get_postgres_indexes_for_transfer(state, source_pool_key, &request.source_schema, table).await? + } else { + Vec::new() + }; + let source_foreign_keys = if request.create_table && pg_compat_transfer && !target_table_preexisting { + get_postgres_foreign_keys_for_transfer(state, source_pool_key, &request.source_schema, table).await? + } else { + Vec::new() + }; + // Count source rows let total_rows = { let sql = count_sql(table, &request.source_schema, source_db_type); @@ -1078,9 +1910,17 @@ where // Create table on target if requested if request.create_table { + if matches!(target_db_type, DatabaseType::Postgres) && !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) + .await + .map_err(|e| format!("Failed to ensure schema exists: {e}"))?; + } let ddl = generate_create_table_ddl( &columns, table, + &request.source_schema, &request.target_schema, target_db_type, source_db_type, @@ -1157,7 +1997,15 @@ where return Err("Cancelled".to_string()); } - let sql = pagination_sql(&col_names, table, &request.source_schema, source_db_type, offset, batch_size); + let sql = pagination_sql_with_order( + &col_names, + table, + &request.source_schema, + source_db_type, + offset, + batch_size, + &primary_key_columns, + ); let result = execute_on_pool(state, source_pool_key, &sql).await?; let row_count = result.rows.len(); @@ -1210,9 +2058,355 @@ where } } + if pg_compat_transfer { + for statement in generate_postgres_sequence_sync_sql(&columns, table, &request.target_schema) { + execute_on_pool(state, target_pool_key, &statement) + .await + .map_err(|e| format!("Failed to sync PostgreSQL sequence for {table}: {e}"))?; + } + } + + if request.create_table && pg_compat_transfer && !target_table_preexisting { + for statement in generate_postgres_index_ddl(&source_indexes, table, &request.target_schema) { + execute_on_pool(state, target_pool_key, &statement) + .await + .map_err(|e| format!("Failed to create PostgreSQL index for {table}: {e}"))?; + } + for statement in generate_postgres_foreign_key_ddl(&source_foreign_keys, table, &request.target_schema) { + execute_on_pool(state, target_pool_key, &statement) + .await + .map_err(|e| format!("Failed to create PostgreSQL foreign key for {table}: {e}"))?; + } + } + Ok(total_transferred) } +pub async fn transfer_postgres_schema_dependencies( + state: &AppState, + request: &TransferRequest, + source_pool_key: &str, + target_pool_key: &str, + mut progress_callback: F, +) -> Result<(), String> +where + F: FnMut(TransferProgress), +{ + let source_db_type = get_db_type(state, &request.source_connection_id).await?; + let target_db_type = get_db_type(state, &request.target_connection_id).await?; + if !request.create_table || !is_postgres_compat_transfer(&source_db_type, &target_db_type) { + return Ok(()); + } + + if !request.target_schema.trim().is_empty() { + let create_schema_sql = format!( + "CREATE SCHEMA IF NOT EXISTS {}", + quote_identifier(&request.target_schema, &DatabaseType::Postgres) + ); + execute_on_pool(state, target_pool_key, &create_schema_sql) + .await + .map_err(|e| format!("Failed to ensure PostgreSQL target schema exists: {e}"))?; + } + + let extensions = + get_postgres_extension_sources_for_transfer(state, source_pool_key, &request.source_schema).await?; + let enum_types = get_postgres_enum_sources_for_transfer(state, source_pool_key, &request.source_schema).await?; + let domains = get_postgres_domain_sources_for_transfer(state, source_pool_key, &request.source_schema).await?; + let total_steps = extensions.len() + enum_types.len() + domains.len(); + let table_index = 0; + let mut completed_steps = 0_u64; + + for extension in extensions { + if is_cancelled(&request.transfer_id).await { + return Err("Cancelled".to_string()); + } + completed_steps += 1; + progress_callback(TransferProgress { + transfer_id: request.transfer_id.clone(), + table: format!("extension: {}", extension.extension_name), + table_index, + total_tables: request.tables.len(), + rows_transferred: completed_steps, + total_rows: Some(total_steps as u64), + status: TransferStatus::Running, + error: None, + }); + execute_on_pool(state, target_pool_key, &generate_postgres_extension_ddl(&extension, &request.target_schema)) + .await + .map_err(|e| format!("Failed to create PostgreSQL extension {}: {e}", extension.extension_name))?; + } + + for enum_type in enum_types { + if is_cancelled(&request.transfer_id).await { + return Err("Cancelled".to_string()); + } + completed_steps += 1; + progress_callback(TransferProgress { + transfer_id: request.transfer_id.clone(), + table: format!("enum: {}", enum_type.type_name), + table_index, + total_tables: request.tables.len(), + rows_transferred: completed_steps, + total_rows: Some(total_steps as u64), + status: TransferStatus::Running, + error: None, + }); + execute_on_pool(state, target_pool_key, &generate_postgres_enum_ddl(&enum_type, &request.target_schema)) + .await + .map_err(|e| format!("Failed to create PostgreSQL enum {}: {e}", enum_type.type_name))?; + } + + for domain in domains { + if is_cancelled(&request.transfer_id).await { + return Err("Cancelled".to_string()); + } + completed_steps += 1; + progress_callback(TransferProgress { + transfer_id: request.transfer_id.clone(), + table: format!("domain: {}", domain.domain_name), + table_index, + total_tables: request.tables.len(), + rows_transferred: completed_steps, + total_rows: Some(total_steps as u64), + status: TransferStatus::Running, + error: None, + }); + execute_on_pool(state, target_pool_key, &generate_postgres_domain_ddl(&domain, &request.target_schema)) + .await + .map_err(|e| format!("Failed to create PostgreSQL domain {}: {e}", domain.domain_name))?; + } + + Ok(()) +} + +pub async fn transfer_postgres_schema_objects( + state: &AppState, + request: &TransferRequest, + source_pool_key: &str, + target_pool_key: &str, + mut progress_callback: F, +) -> Result<(), String> +where + F: FnMut(TransferProgress), +{ + let source_db_type = get_db_type(state, &request.source_connection_id).await?; + let target_db_type = get_db_type(state, &request.target_connection_id).await?; + if !request.create_table || !is_postgres_compat_transfer(&source_db_type, &target_db_type) { + return Ok(()); + } + + let object_sources = + get_postgres_schema_object_sources_for_transfer(state, source_pool_key, &request.source_schema).await?; + let materialized_views = + get_postgres_materialized_view_sources_for_transfer(state, source_pool_key, &request.source_schema).await?; + let trigger_sources = + get_postgres_trigger_sources_for_transfer(state, source_pool_key, &request.source_schema, &request.tables) + .await?; + let policy_statements = get_postgres_policy_statements_for_transfer( + state, + source_pool_key, + &request.source_schema, + &request.target_schema, + &request.tables, + ) + .await?; + let ownership_statements = get_postgres_ownership_statements_for_transfer( + state, + source_pool_key, + &request.source_schema, + &request.target_schema, + &request.tables, + ) + .await?; + let grant_statements = get_postgres_grant_statements_for_transfer( + state, + source_pool_key, + &request.source_schema, + &request.target_schema, + &request.tables, + ) + .await?; + let materialized_view_step_count = materialized_views + .iter() + .map(|view| generate_postgres_materialized_view_ddls(view, &request.target_schema).len()) + .sum::(); + let trigger_step_count = trigger_sources.len() * 2; + let total_steps = object_sources.len() + + materialized_view_step_count + + trigger_step_count + + policy_statements.len() + + ownership_statements.len() + + grant_statements.len(); + let table_index = request.tables.len(); + let mut completed_steps = 0_u64; + + for object in object_sources { + if is_cancelled(&request.transfer_id).await { + return Err("Cancelled".to_string()); + } + completed_steps += 1; + progress_callback(TransferProgress { + transfer_id: request.transfer_id.clone(), + table: format!("schema object: {}", object.name), + table_index, + total_tables: request.tables.len(), + rows_transferred: completed_steps, + total_rows: Some(total_steps as u64), + status: TransferStatus::Running, + error: None, + }); + + let rewritten_source = match object.object_type { + db::ObjectSourceKind::View => object.source.clone(), + db::ObjectSourceKind::Procedure | db::ObjectSourceKind::Function => { + rewrite_postgres_routine_schema(&object.source, &request.target_schema) + .unwrap_or_else(|| object.source.clone()) + } + }; + let statements = build_executable_object_source_statements(EditableObjectSourceSqlInput { + database_type: DatabaseType::Postgres, + object_type: object.object_type.clone(), + schema: Some(request.target_schema.clone()), + name: object.name.clone(), + source: rewritten_source, + })?; + for statement in statements { + execute_on_pool(state, target_pool_key, &statement) + .await + .map_err(|e| format!("Failed to create PostgreSQL {:?} {}: {e}", object.object_type, object.name))?; + } + } + + for view in materialized_views { + for statement in generate_postgres_materialized_view_ddls(&view, &request.target_schema) { + if is_cancelled(&request.transfer_id).await { + return Err("Cancelled".to_string()); + } + completed_steps += 1; + progress_callback(TransferProgress { + transfer_id: request.transfer_id.clone(), + table: format!("materialized view: {}", view.view_name), + table_index, + total_tables: request.tables.len(), + rows_transferred: completed_steps, + total_rows: Some(total_steps as u64), + status: TransferStatus::Running, + error: None, + }); + execute_on_pool(state, target_pool_key, &statement) + .await + .map_err(|e| format!("Failed to create PostgreSQL materialized view {}: {e}", view.view_name))?; + } + } + + for trigger in trigger_sources { + if is_cancelled(&request.transfer_id).await { + return Err("Cancelled".to_string()); + } + completed_steps += 1; + progress_callback(TransferProgress { + transfer_id: request.transfer_id.clone(), + table: format!("trigger: {}", trigger.trigger_name), + table_index, + total_tables: request.tables.len(), + rows_transferred: completed_steps, + total_rows: Some(total_steps as u64), + status: TransferStatus::Running, + error: None, + }); + let full_table = qualified_table(&trigger.table_name, &request.target_schema, &DatabaseType::Postgres); + let drop_sql = format!( + "DROP TRIGGER IF EXISTS {} ON {full_table}", + quote_identifier(&trigger.trigger_name, &DatabaseType::Postgres) + ); + execute_on_pool(state, target_pool_key, &drop_sql) + .await + .map_err(|e| format!("Failed to drop PostgreSQL trigger {}: {e}", trigger.trigger_name))?; + completed_steps += 1; + progress_callback(TransferProgress { + transfer_id: request.transfer_id.clone(), + table: format!("trigger: {}", trigger.trigger_name), + table_index, + total_tables: request.tables.len(), + rows_transferred: completed_steps, + total_rows: Some(total_steps as u64), + status: TransferStatus::Running, + error: None, + }); + let create_sql = rewrite_postgres_trigger_table_schema( + &ensure_sql_statement_terminated(&trigger.source), + &request.source_schema, + &trigger.table_name, + &request.target_schema, + ); + execute_on_pool(state, target_pool_key, &create_sql) + .await + .map_err(|e| format!("Failed to create PostgreSQL trigger {}: {e}", trigger.trigger_name))?; + } + + for statement in policy_statements { + if is_cancelled(&request.transfer_id).await { + return Err("Cancelled".to_string()); + } + completed_steps += 1; + progress_callback(TransferProgress { + transfer_id: request.transfer_id.clone(), + table: "row security policies".to_string(), + table_index, + total_tables: request.tables.len(), + rows_transferred: completed_steps, + total_rows: Some(total_steps as u64), + status: TransferStatus::Running, + error: None, + }); + execute_on_pool(state, target_pool_key, &statement) + .await + .map_err(|e| format!("Failed to apply PostgreSQL row security statement: {e}"))?; + } + + for statement in ownership_statements { + if is_cancelled(&request.transfer_id).await { + return Err("Cancelled".to_string()); + } + completed_steps += 1; + progress_callback(TransferProgress { + transfer_id: request.transfer_id.clone(), + table: "ownership".to_string(), + table_index, + total_tables: request.tables.len(), + rows_transferred: completed_steps, + total_rows: Some(total_steps as u64), + status: TransferStatus::Running, + error: None, + }); + execute_on_pool(state, target_pool_key, &statement) + .await + .map_err(|e| format!("Failed to apply PostgreSQL ownership statement: {e}"))?; + } + + for statement in grant_statements { + if is_cancelled(&request.transfer_id).await { + return Err("Cancelled".to_string()); + } + completed_steps += 1; + progress_callback(TransferProgress { + transfer_id: request.transfer_id.clone(), + table: "grants".to_string(), + table_index, + total_tables: request.tables.len(), + rows_transferred: completed_steps, + total_rows: Some(total_steps as u64), + status: TransferStatus::Running, + error: None, + }); + execute_on_pool(state, target_pool_key, &statement) + .await + .map_err(|e| format!("Failed to apply PostgreSQL grant statement: {e}"))?; + } + + Ok(()) +} + #[cfg(test)] mod tests { use super::*; @@ -1300,7 +2494,7 @@ mod tests { db::ColumnInfo { comment: None, ..test_column("age", "int") }, ]; - let ddl = generate_create_table_ddl(&cols, "users", "", &DatabaseType::Mysql, &DatabaseType::Mysql, None); + let ddl = generate_create_table_ddl(&cols, "users", "", "", &DatabaseType::Mysql, &DatabaseType::Mysql, None); assert!(ddl.contains("COMMENT '用户ID'")); assert!(ddl.contains("COMMENT '用户姓名'")); @@ -1309,12 +2503,82 @@ mod tests { assert!(ddl.contains("PRIMARY KEY (`id`)")); } + #[test] + fn postgres_create_table_preserves_defaults_identity_and_exact_types() { + let cols = vec![ + db::ColumnInfo { + data_type: "integer".to_string(), + column_default: Some("nextval('public.users_id_seq'::regclass)".to_string()), + is_primary_key: true, + is_nullable: false, + ..test_column("id", "integer") + }, + db::ColumnInfo { + data_type: "timestamp with time zone".to_string(), + column_default: Some("now()".to_string()), + is_nullable: false, + ..test_column("created_at", "timestamp with time zone") + }, + db::ColumnInfo { + data_type: "character varying(120)".to_string(), + column_default: Some("'guest'::character varying".to_string()), + ..test_column("name", "character varying(120)") + }, + ]; + + let ddl = generate_create_table_ddl( + &cols, + "users", + "public", + "public", + &DatabaseType::Postgres, + &DatabaseType::Postgres, + None, + ); + + assert!(ddl.contains("\"id\" integer GENERATED BY DEFAULT AS IDENTITY NOT NULL")); + assert!(ddl.contains("\"created_at\" timestamp with time zone DEFAULT now() NOT NULL")); + assert!(ddl.contains("\"name\" character varying(120) DEFAULT 'guest'::character varying")); + assert!(ddl.contains("PRIMARY KEY (\"id\")")); + } + + #[test] + fn postgres_create_table_rewrites_schema_qualified_custom_types_and_defaults() { + let cols = vec![db::ColumnInfo { + data_type: "\"public\".\"user_status\"".to_string(), + column_default: Some("'active'::public.user_status".to_string()), + is_nullable: false, + ..test_column("status", "\"public\".\"user_status\"") + }]; + + let ddl = generate_create_table_ddl( + &cols, + "users", + "public", + "archive", + &DatabaseType::Postgres, + &DatabaseType::Postgres, + None, + ); + + assert!( + ddl.contains("\"status\" \"archive\".\"user_status\" DEFAULT 'active'::\"archive\".user_status NOT NULL") + ); + } + #[test] fn mysql_create_table_includes_table_comment() { let cols = vec![db::ColumnInfo { is_primary_key: true, ..test_column("id", "int") }]; - let ddl = - generate_create_table_ddl(&cols, "users", "", &DatabaseType::Mysql, &DatabaseType::Mysql, Some("用户表")); + let ddl = generate_create_table_ddl( + &cols, + "users", + "", + "", + &DatabaseType::Mysql, + &DatabaseType::Mysql, + Some("用户表"), + ); assert!(ddl.contains(") COMMENT='用户表'")); } @@ -1324,7 +2588,7 @@ mod tests { let cols = vec![db::ColumnInfo { data_type: "text".to_string(), is_primary_key: true, ..test_column("id", "text") }]; - let ddl = generate_create_table_ddl(&cols, "logs", "", &DatabaseType::Mysql, &DatabaseType::Sqlite, None); + let ddl = generate_create_table_ddl(&cols, "logs", "", "", &DatabaseType::Mysql, &DatabaseType::Sqlite, None); assert!(ddl.contains("PRIMARY KEY (`id`(255))")); assert!(ddl.contains("`id` TEXT")); @@ -1334,7 +2598,7 @@ mod tests { fn mysql_int_pk_no_prefix() { let cols = vec![db::ColumnInfo { is_primary_key: true, ..test_column("id", "int") }]; - let ddl = generate_create_table_ddl(&cols, "users", "", &DatabaseType::Mysql, &DatabaseType::Sqlite, None); + let ddl = generate_create_table_ddl(&cols, "users", "", "", &DatabaseType::Mysql, &DatabaseType::Sqlite, None); assert!(ddl.contains("PRIMARY KEY (`id`)")); assert!(!ddl.contains("PRIMARY KEY (`id`(255))")); @@ -1382,10 +2646,168 @@ mod tests { let cols = vec![db::ColumnInfo { comment: Some("test".to_string()), ..test_column("col", "text") }]; // PostgreSQL target should NOT have inline COMMENT - let ddl = generate_create_table_ddl(&cols, "t", "", &DatabaseType::Postgres, &DatabaseType::Postgres, None); + let ddl = generate_create_table_ddl(&cols, "t", "", "", &DatabaseType::Postgres, &DatabaseType::Postgres, None); assert!(!ddl.contains("COMMENT")); } + #[test] + fn postgres_pagination_uses_stable_primary_key_order() { + let sql = pagination_sql_with_order( + &[String::from("id"), String::from("name")], + "users", + "public", + &DatabaseType::Postgres, + 200, + 100, + &[String::from("id")], + ); + + assert_eq!(sql, "SELECT \"id\", \"name\" FROM \"public\".\"users\" ORDER BY \"id\" LIMIT 100 OFFSET 200"); + } + + #[test] + fn postgres_generates_index_and_foreign_key_sql() { + let indexes = vec![db::IndexInfo { + name: "users_name_idx".to_string(), + columns: vec!["lower(name)".to_string()], + is_unique: false, + is_primary: false, + filter: Some("name IS NOT NULL".to_string()), + index_type: Some("btree".to_string()), + included_columns: Some(vec!["created_at".to_string()]), + comment: Some("lookup index".to_string()), + }]; + let foreign_keys = vec![ + db::ForeignKeyInfo { + name: "orders_user_id_fkey".to_string(), + column: "user_id".to_string(), + ref_table: "users".to_string(), + ref_column: "id".to_string(), + }, + db::ForeignKeyInfo { + name: "orders_user_id_fkey".to_string(), + column: "tenant_id".to_string(), + ref_table: "users".to_string(), + ref_column: "tenant_id".to_string(), + }, + ]; + + let index_sql = generate_postgres_index_ddl(&indexes, "users", "public"); + let foreign_key_sql = generate_postgres_foreign_key_ddl(&foreign_keys, "orders", "public"); + + assert_eq!( + index_sql, + vec![ + "CREATE INDEX IF NOT EXISTS \"users_name_idx\" ON \"public\".\"users\" USING btree (lower(name)) INCLUDE (\"created_at\") WHERE name IS NOT NULL".to_string(), + "COMMENT ON INDEX \"public\".\"users_name_idx\" IS 'lookup index'".to_string(), + ] + ); + assert_eq!( + foreign_key_sql, + vec![ + "ALTER TABLE \"public\".\"orders\" ADD CONSTRAINT \"orders_user_id_fkey\" FOREIGN KEY (\"user_id\", \"tenant_id\") REFERENCES \"public\".\"users\" (\"id\", \"tenant_id\")".to_string() + ] + ); + } + + #[test] + fn postgres_sequence_sync_sql_uses_table_max_values() { + let sql = generate_postgres_sequence_sync_sql( + &[db::ColumnInfo { + name: "id".to_string(), + data_type: "integer".to_string(), + is_nullable: false, + column_default: Some("nextval('public.users_id_seq'::regclass)".to_string()), + is_primary_key: true, + extra: None, + comment: None, + numeric_precision: None, + numeric_scale: None, + character_maximum_length: None, + }], + "users", + "public", + ); + + assert_eq!( + sql, + vec![ + "SELECT setval(pg_get_serial_sequence('\"public\".\"users\"', 'id'), GREATEST(COALESCE(MAX(\"id\"), 0), 1), MAX(\"id\") IS NOT NULL) FROM \"public\".\"users\"".to_string() + ] + ); + } + + #[test] + fn postgres_routine_schema_rewrite_targets_destination_schema() { + let rewritten = rewrite_postgres_routine_schema( + "CREATE OR REPLACE FUNCTION public.bump_counter(id integer)\nRETURNS integer\nLANGUAGE plpgsql\nAS $$ BEGIN RETURN id + 1; END; $$", + "archive", + ) + .unwrap(); + + assert!(rewritten.starts_with("CREATE OR REPLACE FUNCTION \"archive\".\"bump_counter\"(")); + } + + #[test] + fn postgres_trigger_schema_rewrite_targets_destination_table() { + let rewritten = rewrite_postgres_trigger_table_schema( + "CREATE TRIGGER bump BEFORE INSERT ON public.users FOR EACH ROW EXECUTE FUNCTION public.bump_counter()", + "public", + "users", + "archive", + ); + + assert!(rewritten.contains(" ON \"archive\".\"users\" ")); + } + + #[test] + fn postgres_extension_enum_and_domain_ddl_is_repeatable() { + let extension_sql = generate_postgres_extension_ddl( + &PostgresExtensionSource { extension_name: "pgcrypto".to_string() }, + "archive", + ); + let enum_sql = generate_postgres_enum_ddl( + &PostgresEnumSource { + type_name: "status".to_string(), + labels: vec!["pending".to_string(), "done".to_string()], + }, + "archive", + ); + let domain_sql = generate_postgres_domain_ddl( + &PostgresDomainSource { + domain_name: "email".to_string(), + base_type: "text".to_string(), + default_value: Some("'unknown@example.com'::text".to_string()), + not_null: true, + checks: vec!["CHECK ((VALUE ~* '^[^@]+@[^@]+$'::text))".to_string()], + }, + "archive", + ); + + assert_eq!(extension_sql, "CREATE EXTENSION IF NOT EXISTS \"pgcrypto\" WITH SCHEMA \"archive\""); + assert!(enum_sql.contains("DO $$ BEGIN IF NOT EXISTS")); + assert!(enum_sql.contains("CREATE TYPE \"archive\".\"status\" AS ENUM ('pending', 'done')")); + assert!(domain_sql.contains("CREATE DOMAIN \"archive\".\"email\" AS text DEFAULT 'unknown@example.com'::text NOT NULL CHECK ((VALUE ~* '^[^@]+@[^@]+$'::text))")); + } + + #[test] + fn postgres_materialized_view_ddls_drop_and_recreate_in_target_schema() { + let ddls = generate_postgres_materialized_view_ddls( + &PostgresMaterializedViewSource { + view_name: "active_users".to_string(), + source: "SELECT id, name FROM public.users WHERE active".to_string(), + }, + "archive", + ); + + assert_eq!(ddls.len(), 2); + assert_eq!(ddls[0], "DROP MATERIALIZED VIEW IF EXISTS \"archive\".\"active_users\""); + assert_eq!( + ddls[1], + "CREATE MATERIALIZED VIEW \"archive\".\"active_users\" AS\nSELECT id, name FROM public.users WHERE active;" + ); + } + #[test] fn mysql_insert_normalizes_rfc3339_datetime_strings() { let sql = generate_insert_typed( diff --git a/crates/dbx-core/tests/live_postgres_transfer.rs b/crates/dbx-core/tests/live_postgres_transfer.rs new file mode 100644 index 000000000..5e4908a36 --- /dev/null +++ b/crates/dbx-core/tests/live_postgres_transfer.rs @@ -0,0 +1,353 @@ +use dbx_core::connection::{AppState, PoolKind}; +use dbx_core::db::postgres; +use dbx_core::models::connection::{ConnectionConfig, DatabaseType, ProxyType}; +use dbx_core::storage::Storage; +use dbx_core::transfer::{ + get_db_type, transfer_postgres_schema_dependencies, transfer_postgres_schema_objects, transfer_table, TransferMode, + TransferRequest, +}; +use serde_json::json; + +fn postgres_test_config(id: &str, database: &str) -> ConnectionConfig { + ConnectionConfig { + id: id.to_string(), + name: id.to_string(), + db_type: DatabaseType::Postgres, + driver_profile: None, + driver_label: None, + url_params: None, + host: "127.0.0.1".to_string(), + port: 5432, + username: "postgres".to_string(), + password: String::new(), + database: Some(database.to_string()), + visible_databases: None, + attached_databases: Vec::new(), + color: None, + ssh_enabled: false, + ssh_host: String::new(), + ssh_port: 22, + ssh_user: String::new(), + ssh_password: String::new(), + ssh_key_path: String::new(), + ssh_key_passphrase: String::new(), + ssh_expose_lan: false, + ssh_connect_timeout_secs: 5, + connect_timeout_secs: 5, + query_timeout_secs: 30, + proxy_enabled: false, + proxy_type: ProxyType::Socks5, + proxy_host: String::new(), + proxy_port: 1080, + proxy_username: String::new(), + proxy_password: String::new(), + ssl: false, + ca_cert_path: String::new(), + sysdba: false, + oracle_connection_type: None, + connection_string: None, + redis_connection_mode: None, + redis_sentinel_master: String::new(), + redis_sentinel_nodes: String::new(), + redis_sentinel_username: String::new(), + redis_sentinel_password: String::new(), + redis_sentinel_tls: false, + redis_cluster_nodes: String::new(), + external_config: None, + jdbc_driver_class: None, + jdbc_driver_paths: Vec::new(), + one_time: false, + } +} + +async fn query_scalar(pool: &deadpool_postgres::Pool, sql: &str) -> serde_json::Value { + postgres::execute_query(pool, sql).await.unwrap().rows[0][0].clone() +} + +#[tokio::test] +#[ignore = "requires source/target PostgreSQL URLs via DBX_LIVE_PG_TRANSFER_SOURCE_URL and DBX_LIVE_PG_TRANSFER_TARGET_URL"] +async fn live_postgres_transfer_preserves_data_and_schema_objects() { + let source_url = std::env::var("DBX_LIVE_PG_TRANSFER_SOURCE_URL").expect("DBX_LIVE_PG_TRANSFER_SOURCE_URL"); + let target_url = std::env::var("DBX_LIVE_PG_TRANSFER_TARGET_URL").unwrap_or_else(|_| source_url.clone()); + + let source_pool = postgres::connect(&source_url, std::time::Duration::from_secs(5)).await.unwrap(); + let target_pool = postgres::connect(&target_url, std::time::Duration::from_secs(5)).await.unwrap(); + + let source_database = query_scalar(&source_pool, "SELECT current_database()").await.as_str().unwrap().to_string(); + let target_database = query_scalar(&target_pool, "SELECT current_database()").await.as_str().unwrap().to_string(); + + let suffix = uuid::Uuid::new_v4().simple().to_string(); + let source_schema = format!("dbx_src_{}", &suffix[..8]); + let target_schema = format!("dbx_dst_{}", &suffix[..8]); + + let cleanup_sql = vec![ + format!("DROP SCHEMA IF EXISTS \"{}\" CASCADE", source_schema), + format!("DROP SCHEMA IF EXISTS \"{}\" CASCADE", target_schema), + ]; + let _ = postgres::execute_batch(&source_pool, &[cleanup_sql[0].clone()]).await; + let _ = postgres::execute_batch(&target_pool, &[cleanup_sql[1].clone()]).await; + + let setup_sql = vec![ + format!("CREATE SCHEMA \"{}\"", source_schema), + format!("CREATE TYPE \"{}\".\"user_status\" AS ENUM ('active', 'disabled')", source_schema), + format!( + "CREATE DOMAIN \"{}\".\"email_text\" AS text CHECK (position('@' in VALUE) > 1)", + source_schema + ), + format!( + "CREATE TABLE \"{}\".\"users\" (\ + \"id\" integer GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY,\ + \"email\" \"{}\".\"email_text\" NOT NULL,\ + \"status\" \"{}\".\"user_status\" NOT NULL DEFAULT 'active',\ + \"created_at\" timestamptz NOT NULL DEFAULT now(),\ + \"active\" boolean NOT NULL DEFAULT true,\ + \"display_name\" text NOT NULL\ + )", + source_schema, source_schema, source_schema + ), + format!( + "CREATE TABLE \"{}\".\"audit_logs\" (\ + \"id\" integer GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY,\ + \"user_id\" integer NOT NULL REFERENCES \"{}\".\"users\"(\"id\"),\ + \"action\" text NOT NULL,\ + \"created_at\" timestamptz NOT NULL DEFAULT now()\ + )", + source_schema, source_schema + ), + format!( + "CREATE INDEX \"users_display_name_idx\" ON \"{}\".\"users\" USING btree (lower(display_name))", + source_schema + ), + format!("COMMENT ON INDEX \"{}\".\"users_display_name_idx\" IS 'lookup index'", source_schema), + format!( + "CREATE OR REPLACE FUNCTION \"{}\".\"log_user_insert\"() RETURNS trigger LANGUAGE plpgsql AS $$ \ + BEGIN \ + INSERT INTO \"{}\".\"audit_logs\" (\"user_id\", \"action\") VALUES (NEW.\"id\", 'insert'); \ + RETURN NEW; \ + END; \ + $$", + source_schema, source_schema + ), + format!( + "CREATE TRIGGER \"users_insert_audit\" AFTER INSERT ON \"{}\".\"users\" \ + FOR EACH ROW EXECUTE FUNCTION \"{}\".\"log_user_insert\"()", + source_schema, source_schema + ), + format!( + "INSERT INTO \"{}\".\"users\" (\"email\", \"status\", \"active\", \"display_name\") VALUES \ + ('alpha@example.com', 'active', true, 'Alpha'), \ + ('beta@example.com', 'disabled', false, 'Beta')", + source_schema + ), + format!( + "CREATE VIEW \"{}\".\"active_users\" AS \ + SELECT \"id\", \"email\", \"display_name\" FROM \"{}\".\"users\" WHERE \"active\"", + source_schema, source_schema + ), + format!( + "CREATE MATERIALIZED VIEW \"{}\".\"user_stats\" AS \ + SELECT \"status\", count(*)::bigint AS \"total\" FROM \"{}\".\"users\" GROUP BY \"status\"", + source_schema, source_schema + ), + format!("ALTER TABLE \"{}\".\"users\" ENABLE ROW LEVEL SECURITY", source_schema), + format!( + "CREATE POLICY \"users_public_read\" ON \"{}\".\"users\" AS PERMISSIVE FOR SELECT TO PUBLIC USING (\"active\")", + source_schema + ), + format!("GRANT USAGE ON SCHEMA \"{}\" TO PUBLIC", source_schema), + format!("GRANT SELECT ON TABLE \"{}\".\"users\" TO PUBLIC", source_schema), + format!("GRANT SELECT ON TABLE \"{}\".\"active_users\" TO PUBLIC", source_schema), + format!("GRANT EXECUTE ON FUNCTION \"{}\".\"log_user_insert\"() TO PUBLIC", source_schema), + ]; + postgres::execute_batch(&source_pool, &setup_sql).await.unwrap(); + + let dir = std::env::temp_dir().join(format!("dbx-live-transfer-{}", uuid::Uuid::new_v4())); + std::fs::create_dir_all(&dir).unwrap(); + let storage = Storage::open(&dir.join("storage.db")).await.unwrap(); + let state = AppState::new(storage); + + let source_connection_id = "live-source"; + let target_connection_id = "live-target"; + let source_pool_key = format!("{source_connection_id}:{source_database}"); + let target_pool_key = format!("{target_connection_id}:{target_database}"); + + state.connections.write().await.insert(source_pool_key.clone(), PoolKind::Postgres(source_pool.clone())); + state.connections.write().await.insert(target_pool_key.clone(), PoolKind::Postgres(target_pool.clone())); + state + .configs + .write() + .await + .insert(source_connection_id.to_string(), postgres_test_config(source_connection_id, &source_database)); + state + .configs + .write() + .await + .insert(target_connection_id.to_string(), postgres_test_config(target_connection_id, &target_database)); + + let request = TransferRequest { + transfer_id: format!("live-transfer-{suffix}"), + source_connection_id: source_connection_id.to_string(), + source_database: source_database.clone(), + source_schema: source_schema.clone(), + target_connection_id: target_connection_id.to_string(), + target_database: target_database.clone(), + target_schema: target_schema.clone(), + tables: vec!["users".to_string(), "audit_logs".to_string()], + create_table: true, + mode: TransferMode::Append, + batch_size: 100, + }; + + transfer_postgres_schema_dependencies(&state, &request, &source_pool_key, &target_pool_key, |_| {}).await.unwrap(); + + let source_db_type = get_db_type(&state, source_connection_id).await.unwrap(); + let target_db_type = get_db_type(&state, target_connection_id).await.unwrap(); + for (index, table) in request.tables.iter().enumerate() { + transfer_table( + &state, + &request, + table, + index, + &source_db_type, + &target_db_type, + &source_pool_key, + &target_pool_key, + |_| {}, + ) + .await + .unwrap(); + } + + transfer_postgres_schema_objects(&state, &request, &source_pool_key, &target_pool_key, |_| {}).await.unwrap(); + + assert_eq!( + query_scalar(&target_pool, &format!("SELECT count(*) FROM \"{}\".\"users\"", target_schema)).await, + json!(2) + ); + assert_eq!( + query_scalar(&target_pool, &format!("SELECT count(*) FROM \"{}\".\"audit_logs\"", target_schema)).await, + json!(2) + ); + assert_eq!( + query_scalar( + &target_pool, + &format!( + "SELECT column_default NOT LIKE '%{}%' AND column_default LIKE '%user_status%' \ + FROM information_schema.columns \ + WHERE table_schema = '{}' AND table_name = 'users' AND column_name = 'status'", + source_schema, target_schema + ) + ) + .await, + json!(true) + ); + assert_eq!( + query_scalar( + &target_pool, + &format!( + "SELECT is_identity FROM information_schema.columns \ + WHERE table_schema = '{}' AND table_name = 'users' AND column_name = 'id'", + target_schema + ) + ) + .await, + json!("YES") + ); + assert_eq!( + query_scalar( + &target_pool, + &format!( + "SELECT udt_name FROM information_schema.columns \ + WHERE table_schema = '{}' AND table_name = 'users' AND column_name = 'status'", + target_schema + ) + ) + .await, + json!("user_status") + ); + assert_eq!( + query_scalar( + &target_pool, + &format!( + "SELECT domain_name FROM information_schema.columns \ + WHERE table_schema = '{}' AND table_name = 'users' AND column_name = 'email'", + target_schema + ) + ) + .await, + json!("email_text") + ); + assert_eq!( + query_scalar(&target_pool, &format!("SELECT count(*) FROM \"{}\".\"active_users\"", target_schema)).await, + json!(1) + ); + assert_eq!( + query_scalar( + &target_pool, + &format!("SELECT count(*) FROM \"{}\".\"user_stats\" WHERE \"status\" = 'active'", target_schema) + ) + .await, + json!(1) + ); + assert_eq!( + query_scalar( + &target_pool, + &format!( + "SELECT relrowsecurity FROM pg_catalog.pg_class c \ + JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ + WHERE n.nspname = '{}' AND c.relname = 'users'", + target_schema + ) + ) + .await, + json!(true) + ); + assert_eq!( + query_scalar( + &target_pool, + &format!( + "SELECT count(*) FROM pg_catalog.pg_policy p \ + JOIN pg_catalog.pg_class c ON c.oid = p.polrelid \ + JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ + WHERE n.nspname = '{}' AND c.relname = 'users' AND p.polname = 'users_public_read'", + target_schema + ) + ) + .await, + json!(1) + ); + assert_eq!( + query_scalar( + &target_pool, + &format!( + "SELECT count(*) \ + FROM pg_catalog.pg_namespace n \ + JOIN LATERAL aclexplode(n.nspacl) a ON true \ + WHERE n.nspname = '{}' AND a.grantee = 0 AND a.privilege_type = 'USAGE'", + target_schema + ) + ) + .await, + json!(1) + ); + + postgres::execute_query( + &target_pool, + &format!( + "INSERT INTO \"{}\".\"users\" (\"email\", \"status\", \"active\", \"display_name\") \ + VALUES ('gamma@example.com', 'active', true, 'Gamma')", + target_schema + ), + ) + .await + .unwrap(); + + assert_eq!( + query_scalar(&target_pool, &format!("SELECT count(*) FROM \"{}\".\"audit_logs\"", target_schema)).await, + json!(3) + ); + + let _ = postgres::execute_batch(&source_pool, &[cleanup_sql[0].clone()]).await; + let _ = postgres::execute_batch(&target_pool, &[cleanup_sql[1].clone()]).await; + let _ = std::fs::remove_dir_all(dir); +} diff --git a/src-tauri/src/commands/transfer.rs b/src-tauri/src/commands/transfer.rs index 1955210ad..8ffbb1385 100644 --- a/src-tauri/src/commands/transfer.rs +++ b/src-tauri/src/commands/transfer.rs @@ -33,6 +33,56 @@ pub async fn start_transfer( let total_tables = request.tables.len(); log::info!("[transfer] starting transfer_id={} tables={}", transfer_id, total_tables); + if matches!(source_db_type, dbx_core::models::connection::DatabaseType::Postgres) + && matches!(target_db_type, dbx_core::models::connection::DatabaseType::Postgres) + { + match dbx_core::transfer::transfer_postgres_schema_dependencies( + &state, + &request, + &source_pool_key, + &target_pool_key, + |progress| emit_progress(&app, progress), + ) + .await + { + Ok(()) => {} + Err(e) if e == "Cancelled" => { + emit_progress( + &app, + TransferProgress { + transfer_id: transfer_id.clone(), + table: "schema dependencies".to_string(), + table_index: 0, + total_tables, + rows_transferred: 0, + total_rows: None, + status: TransferStatus::Cancelled, + error: None, + }, + ); + dbx_core::transfer::clear_cancelled(&transfer_id).await; + return; + } + Err(e) => { + emit_progress( + &app, + TransferProgress { + transfer_id: transfer_id.clone(), + table: "schema dependencies".to_string(), + table_index: 0, + total_tables, + rows_transferred: 0, + total_rows: None, + status: TransferStatus::Error, + error: Some(e), + }, + ); + dbx_core::transfer::clear_cancelled(&transfer_id).await; + return; + } + } + } + for (i, table) in request.tables.iter().enumerate() { if dbx_core::transfer::is_cancelled(&transfer_id).await { emit_progress( @@ -124,6 +174,54 @@ pub async fn start_transfer( } } + if matches!(source_db_type, dbx_core::models::connection::DatabaseType::Postgres) + && matches!(target_db_type, dbx_core::models::connection::DatabaseType::Postgres) + { + match dbx_core::transfer::transfer_postgres_schema_objects( + &state, + &request, + &source_pool_key, + &target_pool_key, + |progress| emit_progress(&app, progress), + ) + .await + { + Ok(()) => {} + Err(e) if e == "Cancelled" => { + emit_progress( + &app, + TransferProgress { + transfer_id: transfer_id.clone(), + table: "schema objects".to_string(), + table_index: total_tables, + total_tables, + rows_transferred: 0, + total_rows: None, + status: TransferStatus::Cancelled, + error: None, + }, + ); + dbx_core::transfer::clear_cancelled(&transfer_id).await; + return; + } + Err(e) => { + emit_progress( + &app, + TransferProgress { + transfer_id: transfer_id.clone(), + table: "schema objects".to_string(), + table_index: total_tables, + total_tables, + rows_transferred: 0, + total_rows: None, + status: TransferStatus::Error, + error: Some(e), + }, + ); + } + } + } + emit_progress( &app, TransferProgress {