xiaoxue_admin
/
dbx
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix(jdbc): fall back when statement batches are unsupported 

Closes #5556
This commit is contained in:
t8y2 2026-08-07 01:51:22 +08:00
parent cda6fd2600
commit 4402faac1c
No known key found for this signature in database
2 changed files with 109 additions and 19 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

77
agents/common/src/main/java/com/dbx/agent/BatchExecutor.java
View File

@ -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<String> 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<String> 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,

51
agents/common/src/test/java/com/dbx/agent/BatchExecutorTest.java
View File

@ -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<String> 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<String> 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<String> 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())) {