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

fix(data-compare): batch sync execution

Browse Source
This commit is contained in:
t8y2 committed 2026-06-26 17:10:26 +08:00
1 parent 895f116b04
commit 2f06ca5fe3
12 files changed
+611 -63
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

crates/dbx-core/src/data_compare.rs
+293 -58
View File
@@ -15,6 +15,9 @@ use crate::sql_dialect::{
};
use crate::transfer::{generate_comment_ddl, generate_create_table_ddl};
const DATA_SYNC_INSERT_BATCH_SIZE: usize = 500;
const DATA_SYNC_CONDITION_BATCH_SIZE: usize = 200;
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct CompareDataRowsOptions {
@@ -648,64 +651,9 @@ fn generate_data_sync_statements(options: &GenerateDataSyncSqlOptions<'_>) -> Ve
.collect::<Vec<_>>()
.join(", ");
let column_info = options.column_info;
let added = options
.diff
.added
.par_iter()
.map(|row| {
let values = options
.columns
.iter()
.map(|column| {
format_grid_sql_literal(
row.values.get(column).unwrap_or(&Value::Null),
options.database_type,
column_info_for(column_info, column),
)
})
.collect::<Vec<_>>()
.join(", ");
format!("INSERT INTO {table} ({columns}) VALUES ({values});")
})
.collect::<Vec<_>>();
let modified = options
.diff
.modified
.par_iter()
.map(|row| {
let assignments = row
.changes
.iter()
.map(|change| {
format!(
"{} = {}",
quote_table_identifier(options.database_type, &change.column),
format_grid_sql_literal(
&change.source,
options.database_type,
column_info_for(column_info, &change.column),
)
)
})
.collect::<Vec<_>>()
.join(", ");
format!(
"UPDATE {table} SET {assignments} WHERE {};",
where_by_key(&row.key_values, options.key_columns, options.database_type, column_info)
)
})
.collect::<Vec<_>>();
let removed = options
.diff
.removed
.par_iter()
.map(|row| {
format!(
"DELETE FROM {table} WHERE {};",
where_by_key(&row.key_values, options.key_columns, options.database_type, column_info)
)
})
.collect::<Vec<_>>();
let added = generate_insert_sync_statements(options, &table, &columns, column_info);
let modified = generate_update_sync_statements(options, &table, column_info);
let removed = generate_delete_sync_statements(options, &table, column_info);
let mut statements = Vec::with_capacity(added.len() + modified.len() + removed.len());
statements.extend(added);
@@ -714,6 +662,160 @@ fn generate_data_sync_statements(options: &GenerateDataSyncSqlOptions<'_>) -> Ve
statements
}
fn generate_insert_sync_statements(
options: &GenerateDataSyncSqlOptions<'_>,
table: &str,
columns: &str,
column_info: &[DataGridColumnInfo],
) -> Vec<String> {
options
.diff
.added
.par_chunks(DATA_SYNC_INSERT_BATCH_SIZE)
.map(|chunk| {
let values = chunk
.iter()
.map(|row| {
let row_values = options
.columns
.iter()
.map(|column| {
format_grid_sql_literal(
row.values.get(column).unwrap_or(&Value::Null),
options.database_type,
column_info_for(column_info, column),
)
})
.collect::<Vec<_>>()
.join(", ");
format!("({row_values})")
})
.collect::<Vec<_>>()
.join(", ");
format!("INSERT INTO {table} ({columns}) VALUES {values};")
})
.collect()
}
fn generate_update_sync_statements(
options: &GenerateDataSyncSqlOptions<'_>,
table: &str,
column_info: &[DataGridColumnInfo],
) -> Vec<String> {
options
.diff
.modified
.par_chunks(DATA_SYNC_CONDITION_BATCH_SIZE)
.flat_map_iter(|chunk| {
if chunk.len() == 1 {
return vec![generate_single_update_statement(options, table, column_info, &chunk[0])];
}
let changed_columns = options
.columns
.iter()
.filter(|column| chunk.iter().any(|row| row.changes.iter().any(|change| change.column == **column)))
.collect::<Vec<_>>();
if changed_columns.is_empty() {
return Vec::new();
}
let assignments = changed_columns
.iter()
.map(|column| {
let quoted_column = quote_table_identifier(options.database_type, column);
let cases = chunk
.iter()
.filter_map(|row| {
let change = row.changes.iter().find(|change| change.column == **column)?;
Some(format!(
"WHEN {} THEN {}",
where_by_key(&row.key_values, options.key_columns, options.database_type, column_info),
format_grid_sql_literal(
&change.source,
options.database_type,
column_info_for(column_info, column),
)
))
})
.collect::<Vec<_>>()
.join(" ");
format!("{quoted_column} = CASE {cases} ELSE {quoted_column} END")
})
.collect::<Vec<_>>()
.join(", ");
let where_clause = chunk
.iter()
.map(|row| {
format!(
"({})",
where_by_key(&row.key_values, options.key_columns, options.database_type, column_info)
)
})
.collect::<Vec<_>>()
.join(" OR ");
vec![format!("UPDATE {table} SET {assignments} WHERE {where_clause};")]
})
.collect()
}
fn generate_single_update_statement(
options: &GenerateDataSyncSqlOptions<'_>,
table: &str,
column_info: &[DataGridColumnInfo],
row: &DataCompareModifiedRow,
) -> String {
let assignments = row
.changes
.iter()
.map(|change| {
format!(
"{} = {}",
quote_table_identifier(options.database_type, &change.column),
format_grid_sql_literal(
&change.source,
options.database_type,
column_info_for(column_info, &change.column),
)
)
})
.collect::<Vec<_>>()
.join(", ");
format!(
"UPDATE {table} SET {assignments} WHERE {};",
where_by_key(&row.key_values, options.key_columns, options.database_type, column_info)
)
}
fn generate_delete_sync_statements(
options: &GenerateDataSyncSqlOptions<'_>,
table: &str,
column_info: &[DataGridColumnInfo],
) -> Vec<String> {
options
.diff
.removed
.par_chunks(DATA_SYNC_CONDITION_BATCH_SIZE)
.map(|chunk| {
if chunk.len() == 1 {
return format!(
"DELETE FROM {table} WHERE {};",
where_by_key(&chunk[0].key_values, options.key_columns, options.database_type, column_info)
);
}
let where_clause = chunk
.iter()
.map(|row| {
format!(
"({})",
where_by_key(&row.key_values, options.key_columns, options.database_type, column_info)
)
})
.collect::<Vec<_>>()
.join(" OR ");
format!("DELETE FROM {table} WHERE {where_clause};")
})
.collect()
}
async fn connection_database_type(state: &AppState, connection_id: &str) -> Result<DatabaseType, String> {
state
.configs
@@ -1086,6 +1188,139 @@ mod tests {
assert_eq!(plan.statement_count, 2);
}
#[test]
fn batches_added_rows_into_multi_value_insert_statements() {
let plan = build_data_compare_sync_plan(DataCompareSyncPlanOptions {
tables: vec![DataCompareSyncPlanTableOptions {
table_name: "users".to_string(),
schema: Some("public".to_string()),
columns: vec!["id".to_string(), "name".to_string()],
key_columns: vec!["id".to_string()],
column_info: Vec::new(),
diff: DataCompareResult {
added: vec![
DataCompareRow {
key: "1".to_string(),
key_values: HashMap::from([(String::from("id"), json!(1))]),
values: HashMap::from([
(String::from("id"), json!(1)),
(String::from("name"), json!("Ada")),
]),
},
DataCompareRow {
key: "2".to_string(),
key_values: HashMap::from([(String::from("id"), json!(2))]),
values: HashMap::from([
(String::from("id"), json!(2)),
(String::from("name"), json!("Bob")),
]),
},
],
removed: Vec::new(),
modified: Vec::new(),
},
database_type: Some(DatabaseType::Postgres),
pre_sync_statements: Vec::new(),
}],
});
assert_eq!(plan.insert_count, 2);
assert_eq!(plan.statement_count, 1);
assert_eq!(plan.sync_sql, "INSERT INTO \"public\".\"users\" (\"id\", \"name\") VALUES (1, 'Ada'), (2, 'Bob');");
}
#[test]
fn batches_modified_rows_into_case_update_statements() {
let plan = build_data_compare_sync_plan(DataCompareSyncPlanOptions {
tables: vec![DataCompareSyncPlanTableOptions {
table_name: "users".to_string(),
schema: None,
columns: vec!["id".to_string(), "name".to_string(), "active".to_string()],
key_columns: vec!["id".to_string()],
column_info: Vec::new(),
diff: DataCompareResult {
added: Vec::new(),
removed: Vec::new(),
modified: vec![
DataCompareModifiedRow {
key: "1".to_string(),
key_values: HashMap::from([(String::from("id"), json!(1))]),
source_values: HashMap::new(),
target_values: HashMap::new(),
changes: vec![DataCompareChangedCell {
column: "name".to_string(),
source: json!("Ada"),
target: json!("Ada old"),
}],
},
DataCompareModifiedRow {
key: "2".to_string(),
key_values: HashMap::from([(String::from("id"), json!(2))]),
source_values: HashMap::new(),
target_values: HashMap::new(),
changes: vec![
DataCompareChangedCell {
column: "name".to_string(),
source: json!("Bob"),
target: json!("Bob old"),
},
DataCompareChangedCell {
column: "active".to_string(),
source: json!(false),
target: json!(true),
},
],
},
],
},
database_type: Some(DatabaseType::Mysql),
pre_sync_statements: Vec::new(),
}],
});
assert_eq!(plan.update_count, 2);
assert_eq!(plan.statement_count, 1);
assert_eq!(
plan.sync_sql,
"UPDATE `users` SET `name` = CASE WHEN `id` = 1 THEN 'Ada' WHEN `id` = 2 THEN 'Bob' ELSE `name` END, `active` = CASE WHEN `id` = 2 THEN FALSE ELSE `active` END WHERE (`id` = 1) OR (`id` = 2);"
);
}
#[test]
fn batches_removed_rows_into_or_delete_statements() {
let plan = build_data_compare_sync_plan(DataCompareSyncPlanOptions {
tables: vec![DataCompareSyncPlanTableOptions {
table_name: "users".to_string(),
schema: Some("public".to_string()),
columns: vec!["id".to_string(), "name".to_string()],
key_columns: vec!["id".to_string()],
column_info: Vec::new(),
diff: DataCompareResult {
added: Vec::new(),
removed: vec![
DataCompareRow {
key: "1".to_string(),
key_values: HashMap::from([(String::from("id"), json!(1))]),
values: HashMap::new(),
},
DataCompareRow {
key: "2".to_string(),
key_values: HashMap::from([(String::from("id"), json!(2))]),
values: HashMap::new(),
},
],
modified: Vec::new(),
},
database_type: Some(DatabaseType::Postgres),
pre_sync_statements: Vec::new(),
}],
});
assert_eq!(plan.delete_count, 2);
assert_eq!(plan.statement_count, 1);
assert_eq!(plan.sync_sql, "DELETE FROM \"public\".\"users\" WHERE (\"id\" = 1) OR (\"id\" = 2);");
}
#[test]
fn borrowed_sync_plan_builder_matches_owned_plan() {
let columns = vec!["id".to_string(), "name".to_string()];
crates/dbx-core/src/db/agent_driver.rs
+21 -1
View File
@@ -136,13 +136,14 @@ pub enum AgentMethod {
FetchTableReadPage,
CloseTableReadSession,
GetExplainInfo,
ExecuteBatch,
ExecuteTransaction,
Disconnect,
Shutdown,
}
impl AgentMethod {
pub const ALL: [Self; 27] = [
pub const ALL: [Self; 28] = [
Self::Handshake,
Self::Connect,
Self::TestConnection,
@@ -167,6 +168,7 @@ impl AgentMethod {
Self::FetchTableReadPage,
Self::CloseTableReadSession,
Self::GetExplainInfo,
Self::ExecuteBatch,
Self::ExecuteTransaction,
Self::Disconnect,
Self::Shutdown,
@@ -198,6 +200,7 @@ impl AgentMethod {
Self::FetchTableReadPage => "fetch_table_read_page",
Self::CloseTableReadSession => "close_table_read_session",
Self::GetExplainInfo => "get_explain_info",
Self::ExecuteBatch => "execute_batch",
Self::ExecuteTransaction => "execute_transaction",
Self::Disconnect => "disconnect",
Self::Shutdown => "shutdown",
@@ -857,6 +860,21 @@ impl AgentDriverClient {
self.call_method(AgentMethod::ExecuteTransaction, agent_transaction_params(database, statements, schema)).await
}
pub async fn execute_batch<T: DeserializeOwned + Send + 'static>(
&mut self,
database: Option<&str>,
statements: &[String],
schema: Option<&str>,
timeout_duration: Option<Duration>,
) -> Result<T, String> {
self.call_method_with_timeout(
AgentMethod::ExecuteBatch,
agent_transaction_params(database, statements, schema),
timeout_duration,
)
.await
}
pub async fn call_mongo_method<T: DeserializeOwned + Send + 'static>(
&mut self,
method: MongoAgentMethod,
@@ -1369,6 +1387,7 @@ mod tests {
assert_eq!(AgentMethod::StartTableRead.as_str(), "start_table_read");
assert_eq!(AgentMethod::FetchTableReadPage.as_str(), "fetch_table_read_page");
assert_eq!(AgentMethod::CloseTableReadSession.as_str(), "close_table_read_session");
assert_eq!(AgentMethod::ExecuteBatch.as_str(), "execute_batch");
assert_eq!(AgentMethod::ExecuteTransaction.as_str(), "execute_transaction");
assert_eq!(AgentMethod::Disconnect.as_str(), "disconnect");
assert_eq!(AgentMethod::Shutdown.as_str(), "shutdown");
@@ -1409,6 +1428,7 @@ mod tests {
let _execute_query_page = AgentDriverClient::execute_query_page::<serde_json::Value>;
let _fetch_query_page = AgentDriverClient::fetch_query_page::<serde_json::Value>;
let _close_query_session = AgentDriverClient::close_query_session::<serde_json::Value>;
let _execute_batch = AgentDriverClient::execute_batch::<serde_json::Value>;
let _execute_transaction = AgentDriverClient::execute_transaction::<serde_json::Value>;
}
crates/dbx-core/src/query.rs
+44
View File
@@ -1731,6 +1731,50 @@ pub async fn execute_statements(
let start = std::time::Instant::now();
let mysql_dialect = connection_mysql_query_dialect(state, connection_id).await;
let agent_client = {
let conns = state.connections.read().await;
match conns.get(&pool_key) {
Some(PoolKind::Agent(client)) => Some(client.clone()),
_ => None,
}
};
if let Some(client) = agent_client {
check_read_only_for_connection_multi(state, &pool_key, statements).await?;
let db_type = connection_database_type_for_pool_key(state, &pool_key).await;
let execution_schema = schema_for_execution_context(db_type, schema);
let rewritten_statements;
let statements = if matches!(db_type, Some(DatabaseType::Iris)) {
rewritten_statements =
statements.iter().map(|sql| sql_for_execution_context(db_type, sql, schema)).collect::<Vec<_>>();
rewritten_statements.as_slice()
} else {
statements
};
let mut client = client.lock().await;
let database = if database.trim().is_empty() { None } else { Some(database) };
let timeout_duration = timeout_secs.map(Duration::from_secs);
let result: Result<db::QueryResult, String> =
client.execute_batch(database, statements, execution_schema, timeout_duration).await;
match result {
Ok(result) => return Ok(db::QueryResult { execution_time_ms: start.elapsed().as_millis(), ..result }),
Err(err) => {
if err.contains("Unknown method: execute_batch") {
log::warn!(
"Agent does not support execute_batch; falling back to statement-by-statement execution"
);
} else {
match pool_error_action(connection_database_type(state, connection_id).await, &err) {
PoolErrorAction::ReconnectAndRetry | PoolErrorAction::Discard => {
let _ = state.remove_pool_by_key(&pool_key).await;
}
PoolErrorAction::Keep => {}
}
return Err(err);
}
}
}
}
for (i, sql) in statements.iter().enumerate() {
match do_execute(
state,