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

fix(transfer): preserve PostgreSQL always identity values

Browse Source 
Closes #5942
This commit is contained in:
t8y2 committed 2026-08-12 02:11:19 +00:00
1 parent 02a56e4f64
commit e4bd07c535
2 files changed
+456 -36
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

crates/dbx-core/src/transfer.rs
+271 -36
View File
@@ -1275,6 +1275,16 @@ fn is_postgres_identity_extra(extra: Option<&str>) -> bool {
})
}
fn is_postgres_generated_always_identity_extra(extra: Option<&str>) -> bool {
extra.is_some_and(|value| {
let mut parts = value.split_whitespace();
parts.next().is_some_and(|part| part.eq_ignore_ascii_case("generated"))
&& parts.next().is_some_and(|part| part.eq_ignore_ascii_case("always"))
&& parts.next().is_some_and(|part| part.eq_ignore_ascii_case("as"))
&& parts.next().is_some_and(|part| part.eq_ignore_ascii_case("identity"))
})
}
pub(crate) fn is_identity_column_extra(extra: Option<&str>) -> bool {
extra.is_some_and(|value| {
let normalized = value.trim().to_ascii_lowercase();
@@ -1311,6 +1321,16 @@ fn selected_columns_include_identity_columns(columns: &[String], all_columns: &[
})
}
fn selected_columns_include_postgres_generated_always_identity_columns(
columns: &[String],
all_columns: &[db::ColumnInfo],
) -> bool {
all_columns.iter().any(|column| {
is_postgres_generated_always_identity_extra(column.extra.as_deref())
&& columns.iter().any(|name| name.eq_ignore_ascii_case(&column.name))
})
}
fn is_sqlserver_rowversion_type(data_type: &str) -> bool {
let normalized = data_type.trim().to_ascii_lowercase();
matches!(normalized.split(['(', ' ', '\t', '\n']).next().unwrap_or(""), "timestamp" | "rowversion")
@@ -2951,7 +2971,7 @@ pub fn generate_insert_typed(
return String::new();
}
let template = InsertSqlTemplate::new(columns, table, schema, db_type, catalog);
let template = InsertSqlTemplate::new(columns, table, schema, db_type, catalog, false);
let value_rows = value_rows_sql(rows, column_types, db_type);
template.build(&value_rows)
}
@@ -2963,11 +2983,23 @@ struct InsertSqlTemplate {
}
impl InsertSqlTemplate {
fn new(columns: &[String], table: &str, schema: &str, db_type: &DatabaseType, catalog: Option<&str>) -> Self {
fn new(
columns: &[String],
table: &str,
schema: &str,
db_type: &DatabaseType,
catalog: Option<&str>,
overrides_postgres_system_values: bool,
) -> Self {
let full_table = qualified_table(table, schema, db_type, catalog);
let col_list = columns.iter().map(|column| quote_identifier(column, db_type)).collect::<Vec<_>>().join(", ");
let overriding = if overrides_postgres_system_values && matches!(db_type, DatabaseType::Postgres) {
" OVERRIDING SYSTEM VALUE"
} else {
""
};
Self {
standard_prefix: format!("INSERT INTO {full_table} ({col_list}) VALUES\n"),
standard_prefix: format!("INSERT INTO {full_table} ({col_list}){overriding} VALUES\n"),
oracle_into_prefix: matches!(db_type, DatabaseType::Oracle)
.then(|| format!("INTO {full_table} ({col_list}) VALUES ")),
}
@@ -3054,6 +3086,21 @@ pub fn generate_upsert_typed(
db_type: &DatabaseType,
pk_columns: &[String],
catalog: Option<&str>,
) -> String {
generate_upsert_typed_for_transfer(columns, column_types, rows, table, schema, db_type, pk_columns, catalog, false)
}
#[allow(clippy::too_many_arguments)]
fn generate_upsert_typed_for_transfer(
columns: &[String],
column_types: &[Option<String>],
rows: &[Vec<serde_json::Value>],
table: &str,
schema: &str,
db_type: &DatabaseType,
pk_columns: &[String],
catalog: Option<&str>,
overrides_postgres_system_values: bool,
) -> String {
if rows.is_empty() || pk_columns.is_empty() {
return String::new();
@@ -3077,7 +3124,13 @@ pub fn generate_upsert_typed(
|| matches!(db_type, DatabaseType::Sqlite | DatabaseType::CloudflareD1 | DatabaseType::DuckDb) =>
{
let pk_list = pk_columns.iter().map(|c| quote_identifier(c, db_type)).collect::<Vec<_>>().join(", ");
let mut sql = format!("INSERT INTO {full_table} ({col_list}) VALUES\n{}", value_rows.join(",\n"));
let overriding = if overrides_postgres_system_values && matches!(db_type, DatabaseType::Postgres) {
" OVERRIDING SYSTEM VALUE"
} else {
""
};
let mut sql =
format!("INSERT INTO {full_table} ({col_list}){overriding} VALUES\n{}", value_rows.join(",\n"));
if non_pk_columns.is_empty() {
sql.push_str(&format!("\nON CONFLICT ({pk_list}) DO NOTHING"));
} else {
@@ -3259,12 +3312,28 @@ fn generate_transfer_write_sql(
db_type: &DatabaseType,
pk_columns: &[String],
catalog: Option<&str>,
overrides_postgres_system_values: bool,
) -> String {
match mode {
TransferMode::Upsert => {
generate_upsert_typed(columns, column_types, rows, table, schema, db_type, pk_columns, catalog)
TransferMode::Upsert => generate_upsert_typed_for_transfer(
columns,
column_types,
rows,
table,
schema,
db_type,
pk_columns,
catalog,
overrides_postgres_system_values,
),
_ => {
if rows.is_empty() {
return String::new();
}
let template =
InsertSqlTemplate::new(columns, table, schema, db_type, catalog, overrides_postgres_system_values);
template.build(&value_rows_sql(rows, column_types, db_type))
}
_ => generate_insert_typed(columns, column_types, rows, table, schema, db_type, catalog),
}
}
@@ -3278,6 +3347,31 @@ pub(crate) fn generate_insert_typed_sql_batches(
db_type: &DatabaseType,
catalog: Option<&str>,
limits: SqlBatchLimits,
) -> Result<Vec<(String, usize)>, String> {
generate_insert_typed_sql_batches_for_transfer(
columns,
column_types,
rows,
table,
schema,
db_type,
catalog,
limits,
false,
)
}
#[allow(clippy::too_many_arguments)]
fn generate_insert_typed_sql_batches_for_transfer(
columns: &[String],
column_types: &[Option<String>],
rows: &[Vec<serde_json::Value>],
table: &str,
schema: &str,
db_type: &DatabaseType,
catalog: Option<&str>,
limits: SqlBatchLimits,
overrides_postgres_system_values: bool,
) -> Result<Vec<(String, usize)>, String> {
if rows.is_empty() {
return Ok(Vec::new());
@@ -3290,7 +3384,7 @@ pub(crate) fn generate_insert_typed_sql_batches(
});
let target_sql_bytes = limits.target_sql_bytes.max(1);
let batch_sql_bytes = limits.hard_sql_bytes.map_or(target_sql_bytes, |hard| target_sql_bytes.min(hard));
let template = InsertSqlTemplate::new(columns, table, schema, db_type, catalog);
let template = InsertSqlTemplate::new(columns, table, schema, db_type, catalog, overrides_postgres_system_values);
let value_rows = value_rows_sql(rows, column_types, db_type);
let value_row_bytes = value_rows.iter().map(|row| sql_text_bytes(row, db_type)).collect::<Vec<_>>();
let mut statements = Vec::new();
@@ -3339,13 +3433,14 @@ fn generate_transfer_write_sql_batches(
db_type: &DatabaseType,
pk_columns: &[String],
catalog: Option<&str>,
overrides_postgres_system_values: bool,
) -> Result<Vec<String>, String> {
if rows.is_empty() {
return Ok(Vec::new());
}
if matches!(mode, TransferMode::Append | TransferMode::Overwrite) {
return Ok(generate_insert_typed_sql_batches(
return Ok(generate_insert_typed_sql_batches_for_transfer(
columns,
column_types,
rows,
@@ -3354,6 +3449,7 @@ fn generate_transfer_write_sql_batches(
db_type,
catalog,
SqlBatchLimits::for_database(db_type, max_transfer_write_rows(db_type, mode)),
overrides_postgres_system_values,
)?
.into_iter()
.map(|(sql, _)| sql)
@@ -3380,6 +3476,7 @@ fn generate_transfer_write_sql_batches(
db_type,
pk_columns,
catalog,
overrides_postgres_system_values,
);
while end < rows.len() && end - start < max_rows {
@@ -3393,6 +3490,7 @@ fn generate_transfer_write_sql_batches(
db_type,
pk_columns,
catalog,
overrides_postgres_system_values,
);
if candidate.len() > max_sql_bytes && !accepted.is_empty() {
break;
@@ -5891,6 +5989,7 @@ where
target_db_type,
&[],
request.target_catalog.as_deref(),
false,
)?;
for (statement_index, batch_sql) in write_statements.iter().enumerate() {
execute_on_pool(state, target_pool_key, batch_sql).await.map_err(|e| {
@@ -6265,23 +6364,31 @@ where
execute_on_pool(state, target_pool_key, &truncate_sql).await.map_err(|e| format!("Failed to truncate: {e}"))?;
}
let target_columns = if (request.mode == TransferMode::Upsert
&& !matches!(target_db_type, DatabaseType::ClickHouse | DatabaseType::Hive))
|| matches!(target_db_type, DatabaseType::Postgres | DatabaseType::Dameng)
{
get_columns_for_transfer(
state,
target_pool_key,
&request.target_connection_id,
&request.target_database,
&request.target_schema,
&target_table,
request.target_catalog.as_deref(),
)
.await
.unwrap_or_default()
} else {
Vec::new()
};
// Determine effective mode and PK columns for upsert
let (effective_mode, pk_columns) = if request.mode == TransferMode::Upsert {
if matches!(target_db_type, DatabaseType::ClickHouse | DatabaseType::Hive) {
log::warn!("[transfer] upsert not supported for {:?}, falling back to append", target_db_type);
(TransferMode::Append, vec![])
} else {
let target_columns = get_columns_for_transfer(
state,
target_pool_key,
&request.target_connection_id,
&request.target_database,
&request.target_schema,
&target_table,
request.target_catalog.as_deref(),
)
.await
.unwrap_or_default();
let pks: Vec<String> = transfer_key_columns(&target_columns, target_db_type)
.into_iter()
.filter(|name| col_names.iter().any(|column_name| column_name.eq_ignore_ascii_case(name)))
@@ -6297,22 +6404,10 @@ where
(request.mode.clone(), vec![])
};
let writes_dameng_identity_columns = if matches!(target_db_type, DatabaseType::Dameng) {
let target_columns = get_columns_for_transfer(
state,
target_pool_key,
&request.target_connection_id,
&request.target_database,
&request.target_schema,
&target_table,
request.target_catalog.as_deref(),
)
.await
.unwrap_or_default();
selected_columns_include_identity_columns(&col_names, &target_columns)
} else {
false
};
let writes_dameng_identity_columns = matches!(target_db_type, DatabaseType::Dameng)
&& selected_columns_include_identity_columns(&col_names, &target_columns);
let overrides_postgres_system_values = matches!(target_db_type, DatabaseType::Postgres)
&& selected_columns_include_postgres_generated_always_identity_columns(&col_names, &target_columns);
// Transfer data in batches
let batch_size = if request.batch_size == 0 { 1000 } else { request.batch_size };
@@ -6351,6 +6446,7 @@ where
target_db_type,
&pk_columns,
request.target_catalog.as_deref(),
overrides_postgres_system_values,
)?;
for (statement_index, batch_sql) in write_statements.iter().enumerate() {
execute_transfer_write_statement(
@@ -7985,6 +8081,43 @@ mod tests {
assert!(!selected_columns_include_identity_columns(&[String::from("name")], &target_columns));
}
#[test]
fn detects_selected_postgres_generated_always_identity_columns() {
let target_columns = vec![
db::ColumnInfo {
name: "ID".to_string(),
extra: Some(" GeNeRaTeD\tALWAYS AS\nIDENTITY (start with 1 increment by 1)".to_string()),
..test_column("ID", "bigint")
},
db::ColumnInfo {
extra: Some("generated by default as identity".to_string()),
..test_column("by_default_id", "bigint")
},
db::ColumnInfo {
extra: Some("generated always as (quantity * 2) stored".to_string()),
..test_column("total", "bigint")
},
test_column("name", "text"),
];
assert!(selected_columns_include_postgres_generated_always_identity_columns(
&[String::from("id")],
&target_columns
));
assert!(!selected_columns_include_postgres_generated_always_identity_columns(
&[String::from("by_default_id")],
&target_columns
));
assert!(!selected_columns_include_postgres_generated_always_identity_columns(
&[String::from("total")],
&target_columns
));
assert!(!selected_columns_include_postgres_generated_always_identity_columns(
&[String::from("name")],
&target_columns
));
}
#[test]
fn sqlserver_writable_transfer_columns_skip_rowversion_types() {
let columns = vec![
@@ -9680,6 +9813,7 @@ SELECT 1 FROM dual"#
&DatabaseType::Oracle,
&[],
None,
false,
)
.unwrap();
@@ -9702,6 +9836,7 @@ SELECT 1 FROM dual"#
&DatabaseType::Mysql,
&[],
None,
false,
)
.unwrap();
@@ -9807,6 +9942,7 @@ SELECT 1 FROM dual"#
&DatabaseType::Mysql,
&[String::from("id")],
None,
false,
)
.unwrap();
@@ -9814,6 +9950,105 @@ SELECT 1 FROM dual"#
assert!(statements[0].contains("ON DUPLICATE KEY UPDATE"));
}
#[test]
fn postgres_transfer_insert_overrides_generated_always_identity_values() {
for mode in [TransferMode::Append, TransferMode::Overwrite] {
let statements = generate_transfer_write_sql_batches(
&mode,
&[String::from("id"), String::from("name")],
&[Some(String::from("bigint")), Some(String::from("text"))],
&[vec![json!(42), json!("Ada")]],
"users",
"public",
&DatabaseType::Postgres,
&[],
None,
true,
)
.unwrap();
assert_eq!(statements.len(), 1);
assert_eq!(
statements[0],
"INSERT INTO \"public\".\"users\" (\"id\", \"name\") OVERRIDING SYSTEM VALUE VALUES\n(42, 'Ada')"
);
}
}
#[test]
fn postgres_transfer_upsert_overrides_generated_always_identity_values() {
let statements = generate_transfer_write_sql_batches(
&TransferMode::Upsert,
&[String::from("id"), String::from("name")],
&[Some(String::from("bigint")), Some(String::from("text"))],
&[vec![json!(42), json!("Ada")]],
"users",
"public",
&DatabaseType::Postgres,
&[String::from("id")],
None,
true,
)
.unwrap();
assert_eq!(statements.len(), 1);
assert_eq!(
statements[0],
"INSERT INTO \"public\".\"users\" (\"id\", \"name\") OVERRIDING SYSTEM VALUE VALUES\n(42, 'Ada')\nON CONFLICT (\"id\") DO UPDATE SET \"name\" = EXCLUDED.\"name\""
);
}
#[test]
fn postgres_transfer_without_generated_always_identity_keeps_sql_shape() {
let statements = generate_transfer_write_sql_batches(
&TransferMode::Append,
&[String::from("name")],
&[Some(String::from("text"))],
&[vec![json!("Ada")]],
"users",
"public",
&DatabaseType::Postgres,
&[],
None,
false,
)
.unwrap();
assert_eq!(statements, vec!["INSERT INTO \"public\".\"users\" (\"name\") VALUES\n('Ada')"]);
}
#[test]
fn postgres_system_value_override_is_not_applied_to_other_dialects() {
let statements = generate_transfer_write_sql_batches(
&TransferMode::Append,
&[String::from("id")],
&[Some(String::from("bigint"))],
&[vec![json!(42)]],
"users",
"public",
&DatabaseType::Kingbase,
&[],
None,
true,
)
.unwrap();
assert_eq!(statements, vec!["INSERT INTO \"public\".\"users\" (\"id\") VALUES\n(42)"]);
}
#[test]
fn postgres_non_transfer_insert_keeps_existing_sql_shape() {
let sql = generate_insert(
&[String::from("id"), String::from("name")],
&[vec![json!(42), json!("Ada")]],
"users",
"public",
&DatabaseType::Postgres,
);
assert_eq!(sql, "INSERT INTO \"public\".\"users\" (\"id\", \"name\") VALUES\n(42, 'Ada')");
}
#[test]
fn database_from_pool_key_handles_session_scoped_keys() {
assert_eq!(database_from_pool_key("conn:analytics"), Some("analytics"));
crates/dbx-core/tests/live_postgres_transfer.rs
+185
View File
@@ -71,6 +71,191 @@ async fn query_scalar(pool: &deadpool_postgres::Pool, sql: &str) -> serde_json::
postgres::execute_query(pool, sql).await.unwrap().rows[0][0].clone()
}
#[tokio::test]
#[ignore = "requires PostgreSQL URLs via DBX_LIVE_PG_TRANSFER_SOURCE_URL and DBX_LIVE_PG_TRANSFER_TARGET_URL"]
async fn live_postgres_transfer_upserts_generated_always_identity_values() {
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_always_{}", &suffix[..8]);
let target_schema = format!("dbx_dst_always_{}", &suffix[..8]);
let cleanup_sql = [
format!("DROP SCHEMA IF EXISTS \"{}\" CASCADE", source_schema),
format!("DROP SCHEMA IF EXISTS \"{}\" CASCADE", target_schema),
];
postgres::execute_batch(
&source_pool,
&[
format!("CREATE SCHEMA \"{}\"", source_schema),
format!(
"CREATE TABLE \"{}\".\"items\" (\"id\" bigint GENERATED ALWAYS AS IDENTITY PRIMARY KEY, \"name\" text NOT NULL)",
source_schema
),
format!(
"INSERT INTO \"{}\".\"items\" (\"id\", \"name\") OVERRIDING SYSTEM VALUE VALUES (42, 'Ada')",
source_schema
),
],
)
.await
.unwrap();
postgres::execute_batch(
&target_pool,
&[
format!("CREATE SCHEMA \"{}\"", target_schema),
format!(
"CREATE TABLE \"{}\".\"items\" (\"id\" bigint GENERATED ALWAYS AS IDENTITY PRIMARY KEY, \"name\" text NOT NULL)",
target_schema
),
],
)
.await
.unwrap();
let old_sql = format!(
"INSERT INTO \"{}\".\"items\" (\"id\", \"name\") VALUES (43, 'old-writer') ON CONFLICT (\"id\") DO UPDATE SET \"name\" = EXCLUDED.\"name\"",
target_schema
);
let old_error = postgres::execute_query(&target_pool, &old_sql).await.unwrap_err();
assert!(old_error.contains("identity column defined as GENERATED ALWAYS"), "{old_error}");
postgres::execute_query(
&target_pool,
&format!(
"INSERT INTO \"{}\".\"items\" (\"id\", \"name\") OVERRIDING SYSTEM VALUE VALUES (43, 'append')",
target_schema
),
)
.await
.unwrap();
assert_eq!(
query_scalar(&target_pool, &format!("SELECT \"name\" FROM \"{}\".\"items\" WHERE \"id\" = 43", target_schema))
.await,
json!("append")
);
postgres::execute_batch(
&target_pool,
&[
format!("TRUNCATE TABLE \"{}\".\"items\"", target_schema),
format!(
"INSERT INTO \"{}\".\"items\" (\"id\", \"name\") OVERRIDING SYSTEM VALUE VALUES (44, 'overwrite')",
target_schema
),
],
)
.await
.unwrap();
assert_eq!(
query_scalar(&target_pool, &format!("SELECT \"name\" FROM \"{}\".\"items\" WHERE \"id\" = 44", target_schema))
.await,
json!("overwrite")
);
postgres::execute_query(&target_pool, &format!("TRUNCATE TABLE \"{}\".\"items\"", target_schema)).await.unwrap();
let dir = std::env::temp_dir().join(format!("dbx-live-always-transfer-{suffix}"));
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-always-source";
let target_connection_id = "live-always-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-always-transfer-{suffix}"),
source_connection_id: source_connection_id.to_string(),
source_database: source_database.clone(),
source_schema: source_schema.clone(),
source_catalog: None,
target_connection_id: target_connection_id.to_string(),
target_database: target_database.clone(),
target_schema: target_schema.clone(),
target_catalog: None,
tables: vec!["items".to_string()],
create_table: false,
content: dbx_core::transfer::TransferContent::default(),
objects: Vec::new(),
mode: TransferMode::Upsert,
target_table_name_case: TransferTableNameCase::Preserve,
ownership_policy: TransferOwnershipPolicy::Preserve,
batch_size: 100,
};
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();
let transferred = transfer_table(
&state,
&request,
"items",
0,
&source_db_type,
&target_db_type,
&source_pool_key,
&target_pool_key,
|_| {},
)
.await
.unwrap();
assert_eq!(transferred, 1);
assert_eq!(
query_scalar(&target_pool, &format!("SELECT \"name\" FROM \"{}\".\"items\" WHERE \"id\" = 42", target_schema))
.await,
json!("Ada")
);
assert_eq!(
query_scalar(&target_pool, &format!("SELECT count(*) FROM \"{}\".\"items\"", target_schema)).await,
json!(1)
);
postgres::execute_query(
&source_pool,
&format!("UPDATE \"{}\".\"items\" SET \"name\" = 'Grace' WHERE \"id\" = 42", source_schema),
)
.await
.unwrap();
let updated = transfer_table(
&state,
&request,
"items",
0,
&source_db_type,
&target_db_type,
&source_pool_key,
&target_pool_key,
|_| {},
)
.await
.unwrap();
assert_eq!(updated, 1);
assert_eq!(
query_scalar(&target_pool, &format!("SELECT \"name\" FROM \"{}\".\"items\" WHERE \"id\" = 42", target_schema))
.await,
json!("Grace")
);
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);
}
#[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() {