diff --git a/crates/dbx-core/src/transfer.rs b/crates/dbx-core/src/transfer.rs index 19ffafaeb3..e73e667e8c 100644 --- a/crates/dbx-core/src/transfer.rs +++ b/crates/dbx-core/src/transfer.rs @@ -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::>().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], + rows: &[Vec], + 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::>().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, 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], + rows: &[Vec], + table: &str, + schema: &str, + db_type: &DatabaseType, + catalog: Option<&str>, + limits: SqlBatchLimits, + overrides_postgres_system_values: bool, ) -> Result, 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::>(); 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, 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 = 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")); diff --git a/crates/dbx-core/tests/live_postgres_transfer.rs b/crates/dbx-core/tests/live_postgres_transfer.rs index 4467e5fa49..47f25bee30 100644 --- a/crates/dbx-core/tests/live_postgres_transfer.rs +++ b/crates/dbx-core/tests/live_postgres_transfer.rs @@ -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() {