fix(data-compare): batch sync execution

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

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<String> statements, String schema) {
+        return BatchExecutor.executeBatchStatements(requireConnected(), statements, schema, this::setSchemaSQL);
+    }
+
     protected Connection requireConnected() {
         Connection conn = getConnection();
         if (conn == null) {

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<String> statements,
+        String schema,
+        Function<String, String> 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<String, String> 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> T unchecked(ThrowingSupplier<T> supplier) {
+        try {
+            return supplier.get();
+        } catch (RuntimeException e) {
+            throw e;
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    private interface ThrowingSupplier<T> {
+        T get() throws Exception;
+    }
+}

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<String> 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);
     }

agents/common/src/main/java/com/dbx/agent/JsonRpcServer.java

@@ -222,6 +222,11 @@ public final class JsonRpcServer {
             List<String> 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<List<String>>() {}.getType();
+            List<String> 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();

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"