From 2f06ca5fe3e583501047d142cd347dbcdda92b23 Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Fri, 26 Jun 2026 17:09:47 +0800 Subject: [PATCH] fix(data-compare): batch sync execution --- .../java/com/dbx/agent/AgentProtocol.java | 2 + .../java/com/dbx/agent/BaseDatabaseAgent.java | 5 + .../java/com/dbx/agent/BatchExecutor.java | 117 ++++++ .../java/com/dbx/agent/DatabaseAgent.java | 8 + .../java/com/dbx/agent/JsonRpcServer.java | 5 + .../src/main/resources/agent-protocol-v1.json | 1 + .../java/com/dbx/agent/BatchExecutorTest.java | 100 +++++ .../src/components/diff/DataCompareDialog.vue | 18 +- crates/dbx-core/assets/agent-protocol-v1.json | 1 + crates/dbx-core/src/data_compare.rs | 351 +++++++++++++++--- crates/dbx-core/src/db/agent_driver.rs | 22 +- crates/dbx-core/src/query.rs | 44 +++ 12 files changed, 611 insertions(+), 63 deletions(-) create mode 100644 agents/common/src/main/java/com/dbx/agent/BatchExecutor.java create mode 100644 agents/common/src/test/java/com/dbx/agent/BatchExecutorTest.java diff --git a/agents/common/src/main/java/com/dbx/agent/AgentProtocol.java b/agents/common/src/main/java/com/dbx/agent/AgentProtocol.java index a04f7eac3..411a02ea6 100644 --- a/agents/common/src/main/java/com/dbx/agent/AgentProtocol.java +++ b/agents/common/src/main/java/com/dbx/agent/AgentProtocol.java @@ -31,6 +31,7 @@ public final class AgentProtocol { public static final String METHOD_FETCH_TABLE_READ_PAGE = "fetch_table_read_page"; public static final String METHOD_CLOSE_TABLE_READ_SESSION = "close_table_read_session"; public static final String METHOD_GET_EXPLAIN_INFO = "get_explain_info"; + public static final String METHOD_EXECUTE_BATCH = "execute_batch"; public static final String METHOD_EXECUTE_TRANSACTION = "execute_transaction"; public static final String METHOD_DISCONNECT = "disconnect"; public static final String METHOD_SHUTDOWN = "shutdown"; @@ -102,6 +103,7 @@ public final class AgentProtocol { METHOD_FETCH_TABLE_READ_PAGE, METHOD_CLOSE_TABLE_READ_SESSION, METHOD_GET_EXPLAIN_INFO, + METHOD_EXECUTE_BATCH, METHOD_EXECUTE_TRANSACTION, METHOD_DISCONNECT, METHOD_SHUTDOWN diff --git a/agents/common/src/main/java/com/dbx/agent/BaseDatabaseAgent.java b/agents/common/src/main/java/com/dbx/agent/BaseDatabaseAgent.java index 685cb824b..13b353116 100644 --- a/agents/common/src/main/java/com/dbx/agent/BaseDatabaseAgent.java +++ b/agents/common/src/main/java/com/dbx/agent/BaseDatabaseAgent.java @@ -94,6 +94,11 @@ public abstract class BaseDatabaseAgent implements DatabaseAgent { return TransactionExecutor.executeUpdateStatements(requireConnected(), statements, schema, this::setSchemaSQL); } + @Override + public QueryResult executeBatch(List statements, String schema) { + return BatchExecutor.executeBatchStatements(requireConnected(), statements, schema, this::setSchemaSQL); + } + protected Connection requireConnected() { Connection conn = getConnection(); if (conn == null) { diff --git a/agents/common/src/main/java/com/dbx/agent/BatchExecutor.java b/agents/common/src/main/java/com/dbx/agent/BatchExecutor.java new file mode 100644 index 000000000..b8c973f17 --- /dev/null +++ b/agents/common/src/main/java/com/dbx/agent/BatchExecutor.java @@ -0,0 +1,117 @@ +package com.dbx.agent; + +import java.sql.BatchUpdateException; +import java.sql.Connection; +import java.sql.SQLFeatureNotSupportedException; +import java.sql.Statement; +import java.util.Collections; +import java.util.List; +import java.util.function.Function; + +public final class BatchExecutor { + private BatchExecutor() { + } + + public static QueryResult executeBatchStatements( + Connection conn, + List statements, + String schema, + Function setSchemaSql + ) { + return unchecked(() -> { + long start = System.currentTimeMillis(); + applySchema(conn, schema, setSchemaSql); + long totalAffected = 0; + int statementCount = 0; + try (Statement stmt = conn.createStatement()) { + for (String statement : statements) { + String trimmed = JdbcExecutor.trimSql(statement); + if (trimmed.isEmpty()) { + continue; + } + stmt.addBatch(trimmed); + statementCount++; + } + if (statementCount > 0) { + totalAffected = affectedRows(executeBatch(stmt)); + } + } catch (BatchUpdateException e) { + long[] counts = e.getLargeUpdateCounts(); + int failedIndex = counts == null ? 1 : counts.length + 1; + throw new RuntimeException("Statement " + failedIndex + " failed: " + e.getMessage(), e); + } + return new QueryResult( + Collections.emptyList(), + Collections.emptyList(), + totalAffected, + System.currentTimeMillis() - start, + false + ); + }); + } + + private static long affectedRows(long[] updateCounts) { + long total = 0; + if (updateCounts == null) { + return total; + } + for (long count : updateCounts) { + if (count >= 0) { + total += count; + } else if (count == Statement.SUCCESS_NO_INFO) { + total += 1; + } + } + return total; + } + + private static long[] executeBatch(Statement stmt) throws Exception { + try { + return stmt.executeLargeBatch(); + } catch (SQLFeatureNotSupportedException | UnsupportedOperationException | AbstractMethodError e) { + int[] counts = stmt.executeBatch(); + long[] largeCounts = new long[counts.length]; + for (int i = 0; i < counts.length; i++) { + largeCounts[i] = counts[i]; + } + return largeCounts; + } + } + + private static void applySchema(Connection conn, String schema, Function setSchemaSql) throws Exception { + if (schema == null || schema.trim().isEmpty()) { + return; + } + try { + conn.setSchema(schema); + return; + } catch (Exception | AbstractMethodError ignored) { + } + try { + conn.setCatalog(schema); + return; + } catch (Exception | AbstractMethodError ignored) { + } + String schemaSql = setSchemaSql.apply(schema); + if (schemaSql == null || schemaSql.trim().isEmpty()) { + return; + } + try (Statement stmt = conn.createStatement()) { + stmt.execute(schemaSql); + } + } + + private static T unchecked(ThrowingSupplier supplier) { + try { + return supplier.get(); + } catch (RuntimeException e) { + throw e; + } catch (Exception e) { + throw new RuntimeException(e); + } + } + + private interface ThrowingSupplier { + T get() throws Exception; + } +} diff --git a/agents/common/src/main/java/com/dbx/agent/DatabaseAgent.java b/agents/common/src/main/java/com/dbx/agent/DatabaseAgent.java index bd3f41057..27c313842 100644 --- a/agents/common/src/main/java/com/dbx/agent/DatabaseAgent.java +++ b/agents/common/src/main/java/com/dbx/agent/DatabaseAgent.java @@ -151,6 +151,14 @@ public interface DatabaseAgent { return TransactionExecutor.executeUpdateStatements(conn, statements, schema, this::setSchemaSQL); } + default QueryResult executeBatch(List statements, String schema) { + Connection conn = getConnection(); + if (conn == null) { + throw new IllegalStateException("Not connected"); + } + return BatchExecutor.executeBatchStatements(conn, statements, schema, this::setSchemaSQL); + } + default String setSchemaSQL(String schema) { return "SET SCHEMA " + JdbcIdentifiers.INSTANCE.doubleQuote(schema); } diff --git a/agents/common/src/main/java/com/dbx/agent/JsonRpcServer.java b/agents/common/src/main/java/com/dbx/agent/JsonRpcServer.java index 3d3aed70a..569e1b24b 100644 --- a/agents/common/src/main/java/com/dbx/agent/JsonRpcServer.java +++ b/agents/common/src/main/java/com/dbx/agent/JsonRpcServer.java @@ -222,6 +222,11 @@ public final class JsonRpcServer { List statements = gson.fromJson(params.get("statements"), statementsType); return agent.executeTransaction(statements, stringOrNull(params, "schema")); } + if (AgentProtocol.METHOD_EXECUTE_BATCH.equals(method)) { + Type statementsType = new TypeToken>() {}.getType(); + List statements = gson.fromJson(params.get("statements"), statementsType); + return agent.executeBatch(statements, stringOrNull(params, "schema")); + } if (AgentProtocol.METHOD_DISCONNECT.equals(method)) { JdbcExecutor.INSTANCE.closeAllQuerySessions(); JdbcExecutor.INSTANCE.closeAllTableReadSessions(); diff --git a/agents/common/src/main/resources/agent-protocol-v1.json b/agents/common/src/main/resources/agent-protocol-v1.json index 10ea69f17..bd5a31e59 100644 --- a/agents/common/src/main/resources/agent-protocol-v1.json +++ b/agents/common/src/main/resources/agent-protocol-v1.json @@ -59,6 +59,7 @@ "fetch_table_read_page", "close_table_read_session", "get_explain_info", + "execute_batch", "execute_transaction", "disconnect", "shutdown" diff --git a/agents/common/src/test/java/com/dbx/agent/BatchExecutorTest.java b/agents/common/src/test/java/com/dbx/agent/BatchExecutorTest.java new file mode 100644 index 000000000..4500e342a --- /dev/null +++ b/agents/common/src/test/java/com/dbx/agent/BatchExecutorTest.java @@ -0,0 +1,100 @@ +package com.dbx.agent; + +import org.junit.jupiter.api.Test; + +import java.lang.reflect.InvocationHandler; +import java.lang.reflect.Method; +import java.lang.reflect.Proxy; +import java.sql.Connection; +import java.sql.Statement; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.concurrent.atomic.AtomicInteger; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +class BatchExecutorTest { + @Test + void executeBatchStatementsUsesJdbcBatch() { + List batchedSql = new ArrayList<>(); + AtomicInteger executeLargeBatchCalls = new AtomicInteger(); + AtomicInteger executeUpdateCalls = new AtomicInteger(); + + Statement statement = statementProxy(batchedSql, executeLargeBatchCalls, executeUpdateCalls); + Connection connection = connectionProxy(statement); + + QueryResult result = BatchExecutor.executeBatchStatements( + connection, + Arrays.asList(" INSERT INTO items VALUES (1); ", "UPDATE items SET name = 'Ada' WHERE id = 1;"), + null, + schema -> null + ); + + assertEquals(Arrays.asList("INSERT INTO items VALUES (1)", "UPDATE items SET name = 'Ada' WHERE id = 1"), batchedSql); + assertEquals(1, executeLargeBatchCalls.get()); + assertEquals(0, executeUpdateCalls.get()); + assertEquals(2L, result.getAffected_rows()); + } + + private static Statement statementProxy( + List batchedSql, + AtomicInteger executeLargeBatchCalls, + AtomicInteger executeUpdateCalls + ) { + InvocationHandler handler = (Object unused, Method method, Object[] args) -> { + switch (method.getName()) { + case "addBatch": + batchedSql.add((String) args[0]); + return null; + case "executeLargeBatch": + executeLargeBatchCalls.incrementAndGet(); + return new long[]{1L, Statement.SUCCESS_NO_INFO}; + case "executeUpdate": + executeUpdateCalls.incrementAndGet(); + throw new AssertionError("executeBatchStatements must not execute statements one by one"); + default: + return defaultValue(method.getReturnType()); + } + }; + return (Statement) Proxy.newProxyInstance(Statement.class.getClassLoader(), new Class[]{Statement.class}, handler); + } + + private static Connection connectionProxy(Statement statement) { + InvocationHandler handler = (Object unused, Method method, Object[] args) -> { + if ("createStatement".equals(method.getName())) { + return statement; + } + return defaultValue(method.getReturnType()); + }; + return (Connection) Proxy.newProxyInstance(Connection.class.getClassLoader(), new Class[]{Connection.class}, handler); + } + + private static Object defaultValue(Class type) { + if (type == Boolean.TYPE) { + return false; + } + if (type == Byte.TYPE) { + return (byte) 0; + } + if (type == Short.TYPE) { + return (short) 0; + } + if (type == Integer.TYPE) { + return 0; + } + if (type == Long.TYPE) { + return 0L; + } + if (type == Float.TYPE) { + return 0f; + } + if (type == Double.TYPE) { + return 0.0d; + } + if (type == Character.TYPE) { + return '\0'; + } + return null; + } +} diff --git a/apps/desktop/src/components/diff/DataCompareDialog.vue b/apps/desktop/src/components/diff/DataCompareDialog.vue index b952ac1fa..ae8d34c45 100644 --- a/apps/desktop/src/components/diff/DataCompareDialog.vue +++ b/apps/desktop/src/components/diff/DataCompareDialog.vue @@ -64,6 +64,7 @@ interface DataCompareTableResult { } const PREVIEW_LIMIT_OPTIONS = [50, 100, 200, 500]; +const SYNC_EXECUTE_BATCH_SIZE = 500; const { t } = useI18n(); const { toast } = useToast(); @@ -717,13 +718,22 @@ async function executeSql() { executedCount.value = 0; try { await store.ensureConnected(targetConnectionId.value); - for (const stmt of syncPlan.value.syncStatements) { + const statements = syncPlan.value.syncStatements; + for (let index = 0; index < statements.length; index += SYNC_EXECUTE_BATCH_SIZE) { + const batch = statements.slice(index, index + SYNC_EXECUTE_BATCH_SIZE); try { - await api.executeQuery(targetConnectionId.value, targetDatabase.value, stmt, targetSchema.value); + await api.executeBatch(targetConnectionId.value, targetDatabase.value, batch, targetSchema.value); + executedCount.value += batch.length; } catch (e: any) { - syncErrors.value.push({ sql: stmt, error: e?.message || String(e) }); + for (const stmt of batch) { + try { + await api.executeBatch(targetConnectionId.value, targetDatabase.value, [stmt], targetSchema.value); + } catch (singleError: any) { + syncErrors.value.push({ sql: stmt, error: singleError?.message || String(singleError) }); + } + executedCount.value++; + } } - executedCount.value++; } const failed = syncErrors.value.length; if (failed === 0) { diff --git a/crates/dbx-core/assets/agent-protocol-v1.json b/crates/dbx-core/assets/agent-protocol-v1.json index 10ea69f17..bd5a31e59 100644 --- a/crates/dbx-core/assets/agent-protocol-v1.json +++ b/crates/dbx-core/assets/agent-protocol-v1.json @@ -59,6 +59,7 @@ "fetch_table_read_page", "close_table_read_session", "get_explain_info", + "execute_batch", "execute_transaction", "disconnect", "shutdown" diff --git a/crates/dbx-core/src/data_compare.rs b/crates/dbx-core/src/data_compare.rs index 9673e964d..545259cf8 100644 --- a/crates/dbx-core/src/data_compare.rs +++ b/crates/dbx-core/src/data_compare.rs @@ -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::>() .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::>() - .join(", "); - format!("INSERT INTO {table} ({columns}) VALUES ({values});") - }) - .collect::>(); - 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::>() - .join(", "); - format!( - "UPDATE {table} SET {assignments} WHERE {};", - where_by_key(&row.key_values, options.key_columns, options.database_type, column_info) - ) - }) - .collect::>(); - 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::>(); + 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 { + 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::>() + .join(", "); + format!("({row_values})") + }) + .collect::>() + .join(", "); + format!("INSERT INTO {table} ({columns}) VALUES {values};") + }) + .collect() +} + +fn generate_update_sync_statements( + options: &GenerateDataSyncSqlOptions<'_>, + table: &str, + column_info: &[DataGridColumnInfo], +) -> Vec { + 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::>(); + 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::>() + .join(" "); + format!("{quoted_column} = CASE {cases} ELSE {quoted_column} END") + }) + .collect::>() + .join(", "); + let where_clause = chunk + .iter() + .map(|row| { + format!( + "({})", + where_by_key(&row.key_values, options.key_columns, options.database_type, column_info) + ) + }) + .collect::>() + .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::>() + .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 { + 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::>() + .join(" OR "); + format!("DELETE FROM {table} WHERE {where_clause};") + }) + .collect() +} + async fn connection_database_type(state: &AppState, connection_id: &str) -> Result { 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()]; diff --git a/crates/dbx-core/src/db/agent_driver.rs b/crates/dbx-core/src/db/agent_driver.rs index 918176a20..5f0c3f013 100644 --- a/crates/dbx-core/src/db/agent_driver.rs +++ b/crates/dbx-core/src/db/agent_driver.rs @@ -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( + &mut self, + database: Option<&str>, + statements: &[String], + schema: Option<&str>, + timeout_duration: Option, + ) -> Result { + self.call_method_with_timeout( + AgentMethod::ExecuteBatch, + agent_transaction_params(database, statements, schema), + timeout_duration, + ) + .await + } + pub async fn call_mongo_method( &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::; let _fetch_query_page = AgentDriverClient::fetch_query_page::; let _close_query_session = AgentDriverClient::close_query_session::; + let _execute_batch = AgentDriverClient::execute_batch::; let _execute_transaction = AgentDriverClient::execute_transaction::; } diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index d78253079..f1239da5d 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -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::>(); + 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 = + 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,