feat(transfer): improve postgres cross-version migration

This commit is contained in:
t8y2 2026-05-29 00:49:29 +08:00
parent e2959420fc
commit 2e431d6e33
5 changed files with 1887 additions and 33 deletions

View File

@ -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());

View File

@ -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 {

File diff suppressed because it is too large Load Diff

View File

@ -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);
}

View File

@ -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 {