diff --git a/agents/common/src/main/java/com/dbx/agent/BatchExecutor.java b/agents/common/src/main/java/com/dbx/agent/BatchExecutor.java index 8075fc9b8..332863d78 100644 --- a/agents/common/src/main/java/com/dbx/agent/BatchExecutor.java +++ b/agents/common/src/main/java/com/dbx/agent/BatchExecutor.java @@ -32,25 +32,8 @@ public final class BatchExecutor { return unchecked(() -> { long start = System.currentTimeMillis(); applySchema(conn, schema, setSchemaSql, resetSchemaSql); - 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); - } + Long batchAffected = tryExecuteBatch(conn, statements); + long totalAffected = batchAffected == null ? executeIndividually(conn, statements) : batchAffected; return new QueryResult( Collections.emptyList(), Collections.emptyList(), @@ -61,6 +44,54 @@ public final class BatchExecutor { }); } + private static Long tryExecuteBatch(Connection conn, List statements) throws Exception { + try (Statement stmt = conn.createStatement()) { + try { + int statementCount = 0; + for (String statement : statements) { + String trimmed = JdbcExecutor.trimSql(statement); + if (trimmed.isEmpty()) { + continue; + } + stmt.addBatch(trimmed); + statementCount++; + } + return statementCount == 0 ? 0L : 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); + } catch (SQLFeatureNotSupportedException | UnsupportedOperationException | AbstractMethodError e) { + return null; + } + } + } + + private static long executeIndividually(Connection conn, List statements) throws Exception { + long totalAffected = 0; + int statementIndex = 0; + try (Statement stmt = conn.createStatement()) { + for (String statement : statements) { + String trimmed = JdbcExecutor.trimSql(statement); + if (trimmed.isEmpty()) { + continue; + } + statementIndex++; + try { + if (!stmt.execute(trimmed)) { + long updateCount = updateCount(stmt); + if (updateCount >= 0) { + totalAffected += updateCount; + } + } + } catch (Exception e) { + throw new RuntimeException("Statement " + statementIndex + " failed: " + e.getMessage(), e); + } + } + } + return totalAffected; + } + private static long affectedRows(long[] updateCounts) { long total = 0; if (updateCounts == null) { @@ -89,6 +120,14 @@ public final class BatchExecutor { } } + private static long updateCount(Statement stmt) throws Exception { + try { + return stmt.getLargeUpdateCount(); + } catch (SQLFeatureNotSupportedException | UnsupportedOperationException | AbstractMethodError e) { + return stmt.getUpdateCount(); + } + } + private static void applySchema( Connection conn, String schema, diff --git a/agents/common/src/test/java/com/dbx/agent/BatchExecutorTest.java b/agents/common/src/test/java/com/dbx/agent/BatchExecutorTest.java index 4500e342a..8f6e29bef 100644 --- a/agents/common/src/test/java/com/dbx/agent/BatchExecutorTest.java +++ b/agents/common/src/test/java/com/dbx/agent/BatchExecutorTest.java @@ -6,6 +6,7 @@ import java.lang.reflect.InvocationHandler; import java.lang.reflect.Method; import java.lang.reflect.Proxy; import java.sql.Connection; +import java.sql.SQLFeatureNotSupportedException; import java.sql.Statement; import java.util.ArrayList; import java.util.Arrays; @@ -37,6 +38,31 @@ class BatchExecutorTest { assertEquals(2L, result.getAffected_rows()); } + @Test + void executeBatchStatementsFallsBackWhenStatementBatchIsUnsupported() { + List executedSql = new ArrayList<>(); + AtomicInteger addBatchCalls = new AtomicInteger(); + AtomicInteger executeLargeBatchCalls = new AtomicInteger(); + + Statement statement = unsupportedBatchStatementProxy(executedSql, addBatchCalls, executeLargeBatchCalls); + Connection connection = connectionProxy(statement); + + QueryResult result = BatchExecutor.executeBatchStatements( + connection, + Arrays.asList(" UPDATE items SET name = 'Ada' WHERE id = 1; ", " DELETE FROM items WHERE id = 2; "), + null, + schema -> null + ); + + assertEquals(1, addBatchCalls.get()); + assertEquals(0, executeLargeBatchCalls.get()); + assertEquals( + Arrays.asList("UPDATE items SET name = 'Ada' WHERE id = 1", "DELETE FROM items WHERE id = 2"), + executedSql + ); + assertEquals(0L, result.getAffected_rows()); + } + private static Statement statementProxy( List batchedSql, AtomicInteger executeLargeBatchCalls, @@ -60,6 +86,31 @@ class BatchExecutorTest { return (Statement) Proxy.newProxyInstance(Statement.class.getClassLoader(), new Class[]{Statement.class}, handler); } + private static Statement unsupportedBatchStatementProxy( + List executedSql, + AtomicInteger addBatchCalls, + AtomicInteger executeLargeBatchCalls + ) { + InvocationHandler handler = (Object unused, Method method, Object[] args) -> { + switch (method.getName()) { + case "addBatch": + addBatchCalls.incrementAndGet(); + throw new SQLFeatureNotSupportedException("Batches not supported"); + case "executeLargeBatch": + executeLargeBatchCalls.incrementAndGet(); + throw new AssertionError("Unsupported JDBC batches must not be executed"); + case "execute": + executedSql.add((String) args[0]); + return false; + case "getLargeUpdateCount": + return -1L; + 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())) {