diff --git a/crates/dbx-core/src/transfer.rs b/crates/dbx-core/src/transfer.rs index fffc05d8ed..68b70e2216 100644 --- a/crates/dbx-core/src/transfer.rs +++ b/crates/dbx-core/src/transfer.rs @@ -9,6 +9,15 @@ use crate::models::connection::DatabaseType; static CANCELLED: std::sync::LazyLock>> = std::sync::LazyLock::new(|| RwLock::new(HashSet::new())); +#[derive(Debug, Clone, Serialize, Deserialize, Default, PartialEq)] +#[serde(rename_all = "camelCase")] +pub enum TransferMode { + #[default] + Append, + Overwrite, + Upsert, +} + #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct TransferRequest { @@ -21,7 +30,8 @@ pub struct TransferRequest { pub target_schema: String, pub tables: Vec, pub create_table: bool, - pub truncate_before: bool, + #[serde(default)] + pub mode: TransferMode, pub batch_size: usize, } @@ -286,6 +296,150 @@ pub fn generate_insert( format!("INSERT INTO {full_table} ({col_list}) VALUES\n{}", value_rows.join(",\n")) } +pub fn generate_upsert( + columns: &[String], + rows: &[Vec], + table: &str, + schema: &str, + db_type: &DatabaseType, + pk_columns: &[String], +) -> String { + if rows.is_empty() || pk_columns.is_empty() { + return String::new(); + } + + let full_table = qualified_table(table, schema, db_type); + let col_list = columns.iter().map(|c| quote_identifier(c, db_type)).collect::>().join(", "); + + let value_rows: Vec = rows + .iter() + .map(|row| { + let vals: Vec = row.iter().map(|v| escape_value(v, db_type)).collect(); + format!("({})", vals.join(", ")) + }) + .collect(); + + let non_pk_columns: Vec<&String> = columns.iter().filter(|c| !pk_columns.contains(c)).collect(); + + match db_type { + DatabaseType::Postgres | DatabaseType::Sqlite | 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")); + if non_pk_columns.is_empty() { + sql.push_str(&format!("\nON CONFLICT ({pk_list}) DO NOTHING")); + } else { + let update_set = non_pk_columns + .iter() + .map(|c| { + let qc = quote_identifier(c, db_type); + format!("{qc} = EXCLUDED.{qc}") + }) + .collect::>() + .join(", "); + sql.push_str(&format!("\nON CONFLICT ({pk_list}) DO UPDATE SET {update_set}")); + } + sql + } + DatabaseType::Mysql | DatabaseType::Doris | DatabaseType::StarRocks => { + let mut sql = format!("INSERT INTO {full_table} ({col_list}) VALUES\n{}", value_rows.join(",\n")); + if non_pk_columns.is_empty() { + sql.push_str("\nON DUPLICATE KEY UPDATE "); + let first_pk = quote_identifier(&pk_columns[0], db_type); + sql.push_str(&format!("{first_pk} = {first_pk}")); + } else { + let update_set = non_pk_columns + .iter() + .map(|c| { + let qc = quote_identifier(c, db_type); + format!("{qc} = VALUES({qc})") + }) + .collect::>() + .join(", "); + sql.push_str(&format!("\nON DUPLICATE KEY UPDATE {update_set}")); + } + sql + } + DatabaseType::SqlServer => { + let src_col_list = columns.iter().map(|c| quote_identifier(c, db_type)).collect::>().join(", "); + let on_clause = pk_columns + .iter() + .map(|c| { + let qc = quote_identifier(c, db_type); + format!("target.{qc} = src.{qc}") + }) + .collect::>() + .join(" AND "); + + let mut sql = format!( + "MERGE INTO {full_table} AS target USING (VALUES\n{}\n) AS src ({src_col_list}) ON {on_clause}", + value_rows.join(",\n") + ); + + if !non_pk_columns.is_empty() { + let update_set = non_pk_columns + .iter() + .map(|c| { + let qc = quote_identifier(c, db_type); + format!("target.{qc} = src.{qc}") + }) + .collect::>() + .join(", "); + sql.push_str(&format!("\nWHEN MATCHED THEN UPDATE SET {update_set}")); + } + + let insert_cols = columns.iter().map(|c| quote_identifier(c, db_type)).collect::>().join(", "); + let insert_vals = + columns.iter().map(|c| format!("src.{}", quote_identifier(c, db_type))).collect::>().join(", "); + sql.push_str(&format!("\nWHEN NOT MATCHED THEN INSERT ({insert_cols}) VALUES ({insert_vals});")); + sql + } + DatabaseType::Oracle => { + let using_rows: Vec = rows + .iter() + .map(|row| { + let vals: Vec = row + .iter() + .zip(columns.iter()) + .map(|(v, c)| format!("{} AS {}", escape_value(v, db_type), quote_identifier(c, db_type))) + .collect(); + format!("SELECT {} FROM dual", vals.join(", ")) + }) + .collect(); + + let on_clause = pk_columns + .iter() + .map(|c| { + let qc = quote_identifier(c, db_type); + format!("t.{qc} = s.{qc}") + }) + .collect::>() + .join(" AND "); + + let mut sql = + format!("MERGE INTO {full_table} t USING ({}) s ON ({on_clause})", using_rows.join(" UNION ALL ")); + + if !non_pk_columns.is_empty() { + let update_set = non_pk_columns + .iter() + .map(|c| { + let qc = quote_identifier(c, db_type); + format!("t.{qc} = s.{qc}") + }) + .collect::>() + .join(", "); + sql.push_str(&format!("\nWHEN MATCHED THEN UPDATE SET {update_set}")); + } + + let insert_cols = columns.iter().map(|c| quote_identifier(c, db_type)).collect::>().join(", "); + let insert_vals = + columns.iter().map(|c| format!("s.{}", quote_identifier(c, db_type))).collect::>().join(", "); + sql.push_str(&format!("\nWHEN NOT MATCHED THEN INSERT ({insert_cols}) VALUES ({insert_vals})")); + sql + } + _ => generate_insert(columns, rows, table, schema, db_type), + } +} + pub fn pagination_sql( columns: &[String], table: &str, @@ -443,31 +597,7 @@ pub async fn get_columns_for_transfer( let table = table.to_string(); return tokio::task::spawn_blocking(move || { let con = con.lock().map_err(|e| e.to_string())?; - let mut stmt = con - .prepare( - "SELECT column_name, data_type, is_nullable, column_default - FROM information_schema.columns - WHERE table_schema = 'main' AND table_name = ? - ORDER BY ordinal_position", - ) - .map_err(|e| e.to_string())?; - let rows = stmt - .query_map([&table], |row| { - Ok(db::ColumnInfo { - name: row.get::<_, String>(0)?, - data_type: row.get::<_, String>(1)?, - is_nullable: row.get::<_, String>(2).unwrap_or_default() == "YES", - column_default: row.get::<_, Option>(3)?, - is_primary_key: false, - extra: None, - comment: None, - numeric_precision: None, - numeric_scale: None, - character_maximum_length: None, - }) - }) - .map_err(|e| e.to_string())?; - Ok(rows.filter_map(|r| r.ok()).collect()) + crate::schema::duckdb_query_columns(&con, &table) }) .await .map_err(|e| e.to_string())?; @@ -601,8 +731,8 @@ where } } - // Truncate target if requested - if request.truncate_before { + // Truncate target if overwrite mode + if request.mode == TransferMode::Overwrite { let full_table = qualified_table(table, &request.target_schema, target_db_type); let truncate_sql = match target_db_type { DatabaseType::Sqlite | DatabaseType::DuckDb => format!("DELETE FROM {full_table}"), @@ -611,6 +741,34 @@ where execute_on_pool(state, target_pool_key, &truncate_sql).await.map_err(|e| format!("Failed to truncate: {e}"))?; } + // 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) { + log::warn!("[transfer] upsert not supported for ClickHouse, falling back to append"); + (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, + table, + ) + .await + .unwrap_or_default(); + let pks: Vec = target_columns.iter().filter(|c| c.is_primary_key).map(|c| c.name.clone()).collect(); + if pks.is_empty() { + log::warn!("[transfer] table {} has no primary key, falling back to append", table); + (TransferMode::Append, vec![]) + } else { + (TransferMode::Upsert, pks) + } + } + } else { + (request.mode.clone(), vec![]) + }; + // Transfer data in batches let batch_size = if request.batch_size == 0 { 1000 } else { request.batch_size }; let mut offset: u64 = 0; @@ -629,9 +787,14 @@ where break; } - let insert_sql = generate_insert(&col_names, &result.rows, table, &request.target_schema, target_db_type); - if !insert_sql.is_empty() { - execute_on_pool(state, target_pool_key, &insert_sql) + let batch_sql = match effective_mode { + TransferMode::Upsert => { + generate_upsert(&col_names, &result.rows, table, &request.target_schema, target_db_type, &pk_columns) + } + _ => generate_insert(&col_names, &result.rows, table, &request.target_schema, target_db_type), + }; + if !batch_sql.is_empty() { + execute_on_pool(state, target_pool_key, &batch_sql) .await .map_err(|e| format!("Insert failed at offset {offset}: {e}"))?; } diff --git a/src/components/transfer/DataTransferDialog.vue b/src/components/transfer/DataTransferDialog.vue index a4ffc6fa1f..96fdf5c86f 100644 --- a/src/components/transfer/DataTransferDialog.vue +++ b/src/components/transfer/DataTransferDialog.vue @@ -10,7 +10,7 @@ import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@ import { useConnectionStore } from "@/stores/connectionStore"; import DatabaseIcon from "@/components/icons/DatabaseIcon.vue"; import * as api from "@/lib/api"; -import type { TransferProgress } from "@/lib/api"; +import type { TransferProgress, TransferMode } from "@/lib/api"; import type { DatabaseType } from "@/types/database"; import { nextTransferTerminalState } from "@/lib/transferProgressState"; import { ArrowRightLeft, Check, X, Loader2, Square, CheckSquare } from "lucide-vue-next"; @@ -47,7 +47,7 @@ const targetSchema = ref(""); // Options const createTable = ref(true); -const truncateBefore = ref(false); +const transferMode = ref("append"); const batchSize = ref(1000); // Transfer state @@ -188,7 +188,7 @@ function resetState() { targetDatabases.value = []; targetSchema.value = ""; createTable.value = true; - truncateBefore.value = false; + transferMode.value = "append"; batchSize.value = 1000; isTransferring.value = false; transferProgress.value.clear(); @@ -230,7 +230,7 @@ async function startTransfer() { targetSchema: targetSchema.value, tables: [...selectedTables.value], createTable: createTable.value, - truncateBefore: truncateBefore.value, + mode: transferMode.value, batchSize: batchSize.value, }; @@ -437,10 +437,18 @@ const totalTransferred = computed(() => {{ t("transfer.createTable") }} -
- - - {{ t("transfer.truncateBefore") }} +
+ +
diff --git a/src/i18n/locales/en.ts b/src/i18n/locales/en.ts index 56a8cf26c6..7002d0edf0 100644 --- a/src/i18n/locales/en.ts +++ b/src/i18n/locales/en.ts @@ -555,7 +555,10 @@ export default { selectAll: "Select All", deselectAll: "Deselect All", createTable: "Create table if not exists", - truncateBefore: "Truncate target table before insert", + transferMode: "Transfer mode", + modeAppend: "Append (INSERT)", + modeOverwrite: "Overwrite (TRUNCATE + INSERT)", + modeUpsert: "Upsert (INSERT or UPDATE by PK)", batchSize: "Batch size", start: "Start Transfer", cancel: "Cancel", diff --git a/src/i18n/locales/zh-CN.ts b/src/i18n/locales/zh-CN.ts index 9e2394b9a5..1d29414ec8 100644 --- a/src/i18n/locales/zh-CN.ts +++ b/src/i18n/locales/zh-CN.ts @@ -547,7 +547,10 @@ export default { selectAll: "全选", deselectAll: "取消全选", createTable: "目标表不存在时自动建表", - truncateBefore: "插入前清空目标表", + transferMode: "传输模式", + modeAppend: "追加 (INSERT)", + modeOverwrite: "覆盖 (TRUNCATE + INSERT)", + modeUpsert: "更新插入 (按主键 INSERT 或 UPDATE)", batchSize: "每批行数", start: "开始传输", cancel: "取消", diff --git a/src/lib/api.ts b/src/lib/api.ts index a5af8a87db..532e1e5f6d 100644 --- a/src/lib/api.ts +++ b/src/lib/api.ts @@ -137,6 +137,7 @@ export type { SqlFileProgress, TransferRequest, TransferProgress, + TransferMode, TableImportMode, TableImportStatus, TableImportColumnMapping, diff --git a/src/lib/tauri.ts b/src/lib/tauri.ts index 168f5397b3..5c947a0248 100644 --- a/src/lib/tauri.ts +++ b/src/lib/tauri.ts @@ -450,6 +450,8 @@ export async function listenSqlFileProgress(handler: (progress: SqlFileProgress) } // --- Data Transfer --- +export type TransferMode = "append" | "overwrite" | "upsert"; + export interface TransferRequest { transferId: string; sourceConnectionId: string; @@ -460,7 +462,7 @@ export interface TransferRequest { targetSchema: string; tables: string[]; createTable: boolean; - truncateBefore: boolean; + mode: TransferMode; batchSize: number; }