fix(sql): harden SQL recognition, duplicate commits, and batch execution edges

This commit is contained in:
efeire
2026-10-13 17:19:14 +08:00
committed by GitHub
parent 79acc066e1
commit 9cade98737
19 changed files with 1101 additions and 98 deletions
@@ -2,6 +2,8 @@ package com.dbx.agent;
final class AgentExecutionContext {
private static final ThreadLocal<JdbcExecutor> JDBC_EXECUTOR = new ThreadLocal<>();
private record StatementExecution(int timeoutSecs, JdbcExecutor.NativeOperation operation, AbstractJdbcAgent owner) { }
private static final ThreadLocal<StatementExecution> STATEMENT_EXECUTION = new ThreadLocal<>();
private AgentExecutionContext() {
}
@@ -11,6 +13,35 @@ final class AgentExecutionContext {
return executor == null ? JdbcExecutor.INSTANCE : executor;
}
static int statementTimeoutSecs() {
StatementExecution execution = STATEMENT_EXECUTION.get();
return execution == null ? 0 : execution.timeoutSecs();
}
static JdbcExecutor.NativeOperation statementOperation() {
StatementExecution execution = STATEMENT_EXECUTION.get();
return execution == null ? null : execution.operation();
}
static void discardConnection(java.sql.Connection connection) throws java.sql.SQLException {
StatementExecution execution = STATEMENT_EXECUTION.get();
if (execution != null && execution.owner() != null) execution.owner().discardConnection(connection);
else connection.close();
}
static <T> T withStatementTimeout(int timeoutSecs, JdbcExecutor.NativeOperation operation, AbstractJdbcAgent owner,
DatabaseAgent.ThrowingSupplier<T> supplier) throws Exception {
if (timeoutSecs < 0) throw new IllegalArgumentException("Query timeout cannot be negative");
StatementExecution previous = STATEMENT_EXECUTION.get();
STATEMENT_EXECUTION.set(new StatementExecution(timeoutSecs, operation, owner));
try {
return supplier.get();
} finally {
if (previous == null) STATEMENT_EXECUTION.remove();
else STATEMENT_EXECUTION.set(previous);
}
}
static <T> T withJdbcExecutor(JdbcExecutor executor, DatabaseAgent.ThrowingSupplier<T> supplier) throws Exception {
JdbcExecutor previous = JDBC_EXECUTOR.get();
JDBC_EXECUTOR.set(executor);
@@ -30,25 +30,35 @@ public final class BatchExecutor {
Supplier<String> resetSchemaSql
) {
return unchecked(() -> {
long start = System.currentTimeMillis();
applySchema(conn, schema, setSchemaSql, resetSchemaSql);
Long batchAffected = tryExecuteBatch(conn, statements);
long totalAffected = batchAffected == null ? executeIndividually(conn, statements) : batchAffected;
return new QueryResult(
Collections.emptyList(),
Collections.emptyList(),
totalAffected,
System.currentTimeMillis() - start,
false
);
try (JdbcExecutor.NativeOperation operation = JdbcExecutor.current().beginNativeOperation()) {
long start = System.currentTimeMillis();
operation.checkCancelled();
applySchema(conn, schema, setSchemaSql, resetSchemaSql);
operation.checkCancelled();
Long batchAffected = tryExecuteBatch(conn, statements, operation);
long totalAffected = batchAffected == null ? executeIndividually(conn, statements, operation) : batchAffected;
operation.checkCancelled();
return new QueryResult(
Collections.emptyList(),
Collections.emptyList(),
totalAffected,
System.currentTimeMillis() - start,
false
);
}
});
}
private static Long tryExecuteBatch(Connection conn, List<String> statements) throws Exception {
try (Statement stmt = conn.createStatement()) {
private static Long tryExecuteBatch(Connection conn, List<String> statements, JdbcExecutor.NativeOperation operation) throws Exception {
try (var tracked = JdbcExecutor.current().trackStatement(conn.createStatement())) {
Statement stmt = tracked.statement();
if (AgentExecutionContext.statementTimeoutSecs() > 0) {
stmt.setQueryTimeout(AgentExecutionContext.statementTimeoutSecs());
}
try {
int statementCount = 0;
for (String statement : statements) {
operation.checkCancelled();
String trimmed = JdbcExecutor.trimSql(statement);
if (trimmed.isEmpty()) {
continue;
@@ -56,7 +66,8 @@ public final class BatchExecutor {
stmt.addBatch(trimmed);
statementCount++;
}
return statementCount == 0 ? 0L : affectedRows(executeBatch(stmt));
operation.checkCancelled();
return statementCount == 0 ? 0L : affectedRows(executeBatch(stmt, operation));
} catch (BatchUpdateException e) {
long[] counts = e.getLargeUpdateCounts();
int failedIndex = counts == null ? 1 : counts.length + 1;
@@ -67,11 +78,16 @@ public final class BatchExecutor {
}
}
private static long executeIndividually(Connection conn, List<String> statements) throws Exception {
private static long executeIndividually(Connection conn, List<String> statements, JdbcExecutor.NativeOperation operation) throws Exception {
long totalAffected = 0;
int statementIndex = 0;
try (Statement stmt = conn.createStatement()) {
try (var tracked = JdbcExecutor.current().trackStatement(conn.createStatement())) {
Statement stmt = tracked.statement();
if (AgentExecutionContext.statementTimeoutSecs() > 0) {
stmt.setQueryTimeout(AgentExecutionContext.statementTimeoutSecs());
}
for (String statement : statements) {
operation.checkCancelled();
String trimmed = JdbcExecutor.trimSql(statement);
if (trimmed.isEmpty()) {
continue;
@@ -107,10 +123,11 @@ public final class BatchExecutor {
return total;
}
private static long[] executeBatch(Statement stmt) throws Exception {
private static long[] executeBatch(Statement stmt, JdbcExecutor.NativeOperation operation) throws Exception {
try {
return stmt.executeLargeBatch();
} catch (SQLFeatureNotSupportedException | UnsupportedOperationException | AbstractMethodError e) {
operation.checkCancelled();
int[] counts = stmt.executeBatch();
long[] largeCounts = new long[counts.length];
for (int i = 0; i < counts.length; i++) {
@@ -47,11 +47,13 @@ public final class JdbcExecutor {
public final class NativeOperation implements AutoCloseable {
private final java.util.concurrent.atomic.AtomicBoolean cancelled = new java.util.concurrent.atomic.AtomicBoolean();
private final NativeOperation parent = AgentExecutionContext.statementOperation();
private NativeOperation() {
}
public void checkCancelled() {
if (parent != null) parent.checkCancelled();
if (cancelled.get()) {
throw new java.util.concurrent.CancellationException("Native JDBC operation cancelled");
}
@@ -62,6 +64,12 @@ public final class JdbcExecutor {
nativeOperations.remove(this);
}
}
/** Apply the current RPC timeout to a native-driver statement. */
public void applyStatementTimeout(Statement statement) throws SQLException {
int timeoutSecs = AgentExecutionContext.statementTimeoutSecs();
if (timeoutSecs > 0) statement.setQueryTimeout(timeoutSecs);
}
/** Register native-driver statements with this session's cancellation boundary. */
public <S extends Statement> TrackedStatement<S> trackStatement(S statement) {
java.util.Objects.requireNonNull(statement, "statement");
@@ -375,7 +375,8 @@ public final class JsonRpcServer {
if (params.has("boundStatements")) {
return executeBlobBoundStatements(params, statements, true);
}
return agent.executeTransaction(statements, stringOrNull(params, "schema"));
return executeStatementsWithTimeout(intOrDefault(params, "timeoutSecs", 0),
() -> agent.executeTransaction(statements, stringOrNull(params, "schema")));
}
if (AgentProtocol.METHOD_BEGIN_MANUAL_TRANSACTION.equals(method)) {
agent.invalidateLargeValues();
@@ -396,7 +397,8 @@ public final class JsonRpcServer {
if (params.has("boundStatements")) {
return executeBlobBoundStatements(params, statements, false);
}
return agent.executeBatch(statements, stringOrNull(params, "schema"));
return executeStatementsWithTimeout(intOrDefault(params, "timeoutSecs", 0),
() -> agent.executeBatch(statements, stringOrNull(params, "schema")));
}
if (AgentProtocol.METHOD_DISCONNECT.equals(method)) {
jdbcExecutor.closeAllQuerySessions();
@@ -415,6 +417,20 @@ public final class JsonRpcServer {
throw new IllegalArgumentException("Unknown method: " + method);
}
private QueryResult executeStatementsWithTimeout(int timeoutSecs, DatabaseAgent.ThrowingSupplier<QueryResult> execute) throws Exception {
try (JdbcExecutor.NativeOperation operation = jdbcExecutor.beginNativeOperation()) {
return AgentExecutionContext.withStatementTimeout(timeoutSecs, operation,
agent instanceof AbstractJdbcAgent jdbcAgent ? jdbcAgent : null, () -> {
operation.checkCancelled();
if (agent instanceof AbstractJdbcAgent jdbcAgent) {
jdbcAgent.beforeQueryExecution(jdbcAgent.requireConnected(), timeoutSecs);
}
operation.checkCancelled();
return execute.get();
});
}
}
private QueryResult executeBlobBoundStatements(JsonObject params, List<String> previews, boolean transaction) {
if (!agent.supportsBlobBindStatements()) {
throw new UnsupportedOperationException("Agent does not support blob_bind_statements_v1");
@@ -2,6 +2,7 @@ package com.dbx.agent;
import java.sql.Connection;
import java.sql.Statement;
import java.sql.SQLTransientConnectionException;
import java.util.Collections;
import java.util.List;
import java.util.function.Function;
@@ -54,24 +55,42 @@ public final class TransactionExecutor {
StatementRunner runner
) {
return unchecked(() -> {
long start = System.currentTimeMillis();
if (!supportsTransactions(conn)) {
// An explicit transaction request must not degrade to auto-commit:
// callers rely on it to avoid partial writes when a later statement fails.
throw new UnsupportedOperationException("Transactions are not supported by this JDBC driver");
}
try (JdbcExecutor.NativeOperation operation = JdbcExecutor.current().beginNativeOperation()) {
long start = System.currentTimeMillis();
operation.checkCancelled();
if (!supportsTransactions(conn)) {
// An explicit transaction request must not degrade to auto-commit:
// callers rely on it to avoid partial writes when a later statement fails.
throw new UnsupportedOperationException("Transactions are not supported by this JDBC driver");
}
boolean savedAutoCommit = conn.getAutoCommit();
conn.setAutoCommit(false);
try {
long totalAffected = executeAll(conn, statements, schema, setSchemaSql, resetSchemaSql, runner);
conn.commit();
return result(totalAffected, start);
} catch (Exception e) {
conn.rollback();
throw e;
} finally {
conn.setAutoCommit(savedAutoCommit);
boolean savedAutoCommit = conn.getAutoCommit();
conn.setAutoCommit(false);
boolean finished = false;
boolean discarded = false;
try {
long totalAffected = executeAll(conn, statements, schema, setSchemaSql, resetSchemaSql, runner, operation);
operation.checkCancelled();
conn.commit();
finished = true;
return result(totalAffected, start);
} catch (Exception e) {
try {
conn.rollback();
finished = true;
} catch (Exception rollbackFailure) {
SQLTransientConnectionException unsafe = new SQLTransientConnectionException(
"Transaction rollback was not confirmed; discard this connection and treat the operation outcome as unknown", "08007", e);
unsafe.addSuppressed(rollbackFailure);
discarded = true;
try { AgentExecutionContext.discardConnection(conn); } catch (Exception closeFailure) { unsafe.addSuppressed(closeFailure); }
throw unsafe;
}
throw e;
} finally {
if (finished) conn.setAutoCommit(savedAutoCommit);
else if (!discarded) AgentExecutionContext.discardConnection(conn);
}
}
});
}
@@ -82,12 +101,19 @@ public final class TransactionExecutor {
String schema,
Function<String, String> setSchemaSql,
Supplier<String> resetSchemaSql,
StatementRunner runner
StatementRunner runner,
JdbcExecutor.NativeOperation operation
) throws Exception {
applySchema(conn, schema, setSchemaSql, resetSchemaSql);
long totalAffected = 0;
for (String statement : statements) {
try (Statement stmt = conn.createStatement()) {
operation.checkCancelled();
try (var tracked = JdbcExecutor.current().trackStatement(conn.createStatement())) {
Statement stmt = tracked.statement();
if (AgentExecutionContext.statementTimeoutSecs() > 0) {
stmt.setQueryTimeout(AgentExecutionContext.statementTimeoutSecs());
}
operation.checkCancelled();
totalAffected += runner.run(stmt, JdbcExecutor.trimSql(statement));
}
}
@@ -37,6 +37,75 @@ import static org.junit.jupiter.api.Assertions.fail;
class JdbcConnectionPoolingTest {
private static final Gson GSON = new Gson();
@Test
void unconfirmedRollbackCannotReturnItsTransactionToAnotherSession() throws Exception {
for (String sqlState : List.of("HY000", "HYT00", "08006")) {
AtomicInteger physicalOpens = new AtomicInteger();
CountDownLatch transactionFinished = new CountDownLatch(1);
CountDownLatch finishRequest = new CountDownLatch(1);
ExecutorService worker = Executors.newSingleThreadExecutor();
String url = h2Url("unconfirmed_rollback");
try (Connection setup = DriverManager.getConnection(url, "sa", "");
Statement statement = setup.createStatement();
JdbcConnectionPoolRegistry registry = new JdbcConnectionPoolRegistry(poolSettings(2))) {
statement.execute("CREATE TABLE rollback_probe (id INT)");
List<JsonRpcServer> sessions = new java.util.ArrayList<>();
for (int index = 0; index < 2; index++) {
H2TestAgent agent = new H2TestAgent(url, physicalOpens) {
@Override
protected Connection openConnection(ConnectParams params) throws Exception {
Connection physical = openH2(url, physicalOpens);
if (physicalOpens.get() != 1) return physical;
return (Connection) Proxy.newProxyInstance(getClass().getClassLoader(),
new Class<?>[]{Connection.class}, (proxy, method, args) -> {
if (method.getName().equals("rollback")) {
throw new SQLException("Rollback was not acknowledged", sqlState);
}
try {
return method.invoke(physical, args);
} catch (InvocationTargetException error) {
throw error.getCause();
}
});
}
@Override
public QueryResult executeTransaction(List<String> statements, String schema) {
try {
return super.executeTransaction(statements, schema);
} finally {
// Allow another session to borrow before the RPC's outer cleanup.
transactionFinished.countDown();
awaitUninterruptibly(finishRequest);
}
}
};
agent.attachConnectionPoolRegistry(registry);
JsonRpcServer server = new JsonRpcServer(agent);
server.dispatchForRuntime(AgentProtocol.METHOD_CONNECT,
JsonParser.parseString("{\"database\":\"pooling\",\"username\":\"sa\",\"password\":\"\"}").getAsJsonObject());
sessions.add(server);
}
Future<Exception> failure = worker.submit(() -> assertThrows(Exception.class,
() -> sessions.get(0).dispatchForRuntime(AgentProtocol.METHOD_EXECUTE_TRANSACTION,
JsonParser.parseString("{\"statements\":[\"INSERT INTO rollback_probe VALUES (1)\",\"INVALID SQL\"]}").getAsJsonObject())));
try {
assertTrue(transactionFinished.await(3, TimeUnit.SECONDS));
JsonObject query = JsonParser.parseString("{\"sql\":\"SELECT COUNT(*) FROM rollback_probe\"}").getAsJsonObject();
QueryResult result = (QueryResult) sessions.get(1).dispatchForRuntime(AgentProtocol.METHOD_EXECUTE_QUERY, query);
assertEquals(0L, ((Number) result.getRows().get(0).get(0)).longValue(), sqlState);
assertEquals(2, physicalOpens.get(), sqlState);
} finally {
finishRequest.countDown();
}
assertTrue(failure.get(3, TimeUnit.SECONDS).toString().contains("rollback was not confirmed"));
} finally {
finishRequest.countDown();
worker.shutdownNow();
}
}
}
@Test
void blobRollbackFailureCannotReturnDirtyConnectionToAnotherSession() throws Exception {
for (boolean manual : List.of(false, true)) {
@@ -0,0 +1,212 @@
package com.dbx.agent;
import com.google.gson.JsonParser;
import org.junit.jupiter.api.Test;
import java.lang.reflect.Proxy;
import java.sql.*;
import java.util.List;
import java.util.Set;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;
import static org.junit.jupiter.api.Assertions.*;
class JsonRpcBatchLifecycleTest {
@Test
void ordinaryBatchHonorsRpcStatementTimeout() throws Exception {
Fixture fixture = new Fixture();
fixture.expectedTimeout = 7;
QueryResult result = (QueryResult) fixture.server.dispatchForRuntime("execute_batch", params(7));
assertEquals(2, result.getAffected_rows());
}
@Test
void cancellationAtAgentDispatchBoundaryPreventsBatchAndTransactionWrites() throws Exception {
for (String method : List.of("execute_batch", "execute_transaction")) {
Fixture fixture = new Fixture();
fixture.cancelBeforeExecutor = true;
assertThrows(Exception.class, () -> fixture.server.dispatchForRuntime(method, params(0)));
assertEquals(0, fixture.writes.get());
assertEquals(0, fixture.commits.get());
}
}
@Test
void cancellationBeforeUnsupportedLargeBatchFallbackDoesNotExecuteBatchAgain() throws Exception {
Fixture fixture = new Fixture();
fixture.cancelWhenLargeBatchUnsupported = true;
assertThrows(Exception.class, () -> fixture.server.dispatchForRuntime("execute_batch", params(0)));
assertEquals(0, fixture.fallbackWrites.get());
}
@Test
void cancelledTransactionWithUnconfirmedRollbackClosesWithoutEnablingAutoCommit() throws Exception {
Fixture fixture = new Fixture();
fixture.cancelBeforeExecutor = false;
fixture.cancelDuringWrite = true;
fixture.rollbackFails = true;
Exception failure = assertThrows(Exception.class, () -> fixture.server.dispatchForRuntime("execute_transaction", params(0)));
assertEquals(0, fixture.autoCommitRestored.get());
assertEquals(1, fixture.connectionCloses.get());
assertTrue(failure.toString().contains("rollback was not confirmed"));
}
@Test
void transactionAndFallbackStatementsHonorTimeoutWithoutLeakingToLaterRpc() throws Exception {
Fixture fixture = new Fixture();
fixture.expectedTimeout = 7;
assertEquals(2, ((QueryResult) fixture.server.dispatchForRuntime("execute_transaction", params(7))).getAffected_rows());
fixture.fallback = true;
fixture.server.dispatchForRuntime("execute_batch", params(7));
fixture.expectedTimeout = 0;
fixture.server.dispatchForRuntime("execute_transaction", params(0));
}
@Test
void cancellingJdbcBatchRejectsItsLateSuccessAndLeavesNextRpcUsable() throws Exception {
Fixture fixture = new Fixture();
fixture.block = true;
ExecutorService worker = Executors.newSingleThreadExecutor();
try {
Future<?> result = worker.submit(() -> fixture.server.dispatchForRuntime("execute_batch", params(0)));
assertTrue(fixture.started.await(3, TimeUnit.SECONDS));
fixture.server.cancelActiveStatements();
fixture.release.countDown();
assertThrows(ExecutionException.class, () -> result.get(3, TimeUnit.SECONDS));
assertEquals(1, fixture.cancels.get());
fixture.block = false;
assertEquals(2, ((QueryResult) fixture.server.dispatchForRuntime("execute_batch", params(0))).getAffected_rows());
} finally {
fixture.release.countDown();
worker.shutdownNow();
}
}
@Test
void cancellationStopsFallbackBatchBeforeTheNextWrite() throws Exception {
Fixture fixture = new Fixture();
fixture.fallback = true;
fixture.block = true;
ExecutorService worker = Executors.newSingleThreadExecutor();
try {
Future<?> result = worker.submit(() -> fixture.server.dispatchForRuntime("execute_batch", params(0)));
assertTrue(fixture.started.await(3, TimeUnit.SECONDS));
fixture.server.cancelActiveStatements();
fixture.release.countDown();
assertThrows(ExecutionException.class, () -> result.get(3, TimeUnit.SECONDS));
assertEquals(1, fixture.writes.get());
assertEquals(1, fixture.cancels.get());
} finally {
fixture.release.countDown();
worker.shutdownNow();
}
}
@Test
void cancellationRollsBackTransactionInsteadOfCommittingOrStartingNextWrite() throws Exception {
Fixture fixture = new Fixture();
fixture.block = true;
ExecutorService worker = Executors.newSingleThreadExecutor();
try {
Future<?> result = worker.submit(() -> fixture.server.dispatchForRuntime("execute_transaction", params(0)));
assertTrue(fixture.started.await(3, TimeUnit.SECONDS));
fixture.server.cancelActiveStatements();
fixture.release.countDown();
assertThrows(ExecutionException.class, () -> result.get(3, TimeUnit.SECONDS));
assertEquals(1, fixture.writes.get());
assertEquals(1, fixture.cancels.get());
assertEquals(0, fixture.commits.get());
assertEquals(1, fixture.rollbacks.get());
} finally {
fixture.release.countDown();
worker.shutdownNow();
}
}
private static com.google.gson.JsonObject params(int timeout) {
return JsonParser.parseString("{\"statements\":[\"UPDATE T SET C=1\",\"UPDATE T SET C=2\"],\"timeoutSecs\":" + timeout + "}").getAsJsonObject();
}
private static final class Fixture {
int expectedTimeout;
boolean fallback;
boolean block;
boolean cancelBeforeExecutor;
boolean cancelWhenLargeBatchUnsupported;
boolean cancelDuringWrite;
boolean rollbackFails;
final AtomicInteger fallbackWrites = new AtomicInteger();
final AtomicInteger autoCommitRestored = new AtomicInteger();
final AtomicInteger connectionCloses = new AtomicInteger();
final CountDownLatch started = new CountDownLatch(1);
final CountDownLatch release = new CountDownLatch(1);
final AtomicInteger writes = new AtomicInteger();
final AtomicInteger cancels = new AtomicInteger();
final AtomicInteger commits = new AtomicInteger();
final AtomicInteger rollbacks = new AtomicInteger();
final JsonRpcServer server;
Fixture() {
Connection connection = (Connection) Proxy.newProxyInstance(getClass().getClassLoader(), new Class<?>[]{Connection.class}, (p, m, a) -> {
if (m.getName().equals("createStatement")) return statement();
if (m.getName().equals("getAutoCommit")) return true;
if (m.getName().equals("isValid")) return true;
if (m.getName().equals("commit")) { commits.incrementAndGet(); return null; }
if (m.getName().equals("rollback")) { rollbacks.incrementAndGet(); if (rollbackFails) throw new SQLException("rollback network failure", "08006"); return null; }
if (m.getName().equals("setAutoCommit") && Boolean.TRUE.equals(a[0])) { autoCommitRestored.incrementAndGet(); return null; }
if (m.getName().equals("close")) { connectionCloses.incrementAndGet(); return null; }
if (m.getName().equals("getMetaData")) return Proxy.newProxyInstance(getClass().getClassLoader(), new Class<?>[]{DatabaseMetaData.class}, (q, n, b) -> {
if (n.getName().equals("supportsTransactions")) return true;
if (n.getName().equals("getIdentifierQuoteString")) return "\"";
return defaultValue(n.getReturnType());
});
return defaultValue(m.getReturnType());
});
ConfiguredJdbcAgent agent = new ConfiguredJdbcAgent(new JdbcAgentProfile("unused", "jdbc:test", 1, false, Set.of(), List.of())) {
@Override protected void loadDriver(ConnectParams params) { }
@Override protected Connection openConnection(ConnectParams params) { return connection; }
@Override public QueryResult executeBatch(List<String> statements, String schema) {
if (cancelBeforeExecutor) server.cancelActiveStatements();
return super.executeBatch(statements, schema);
}
@Override public QueryResult executeTransaction(List<String> statements, String schema) {
if (cancelBeforeExecutor) server.cancelActiveStatements();
return super.executeTransaction(statements, schema);
}
};
agent.connect(new ConnectParams());
server = new JsonRpcServer(agent);
}
Statement statement() {
AtomicInteger timeout = new AtomicInteger();
return (Statement) Proxy.newProxyInstance(getClass().getClassLoader(), new Class<?>[]{Statement.class}, (p, m, a) -> {
if (m.getName().equals("setQueryTimeout")) { timeout.set((int) a[0]); return null; }
if (m.getName().equals("addBatch") && fallback) throw new SQLFeatureNotSupportedException();
if (m.getName().equals("cancel")) { cancels.incrementAndGet(); return null; }
if (m.getName().equals("execute") || m.getName().equals("executeUpdate")) {
if (timeout.get() != expectedTimeout) throw new SQLTimeoutException("Statement has no requested execution timeout");
writes.incrementAndGet();
if (cancelDuringWrite) server.cancelActiveStatements();
if (block && writes.get() == 1) { started.countDown(); assertTrue(release.await(3, TimeUnit.SECONDS)); }
return m.getName().equals("execute") ? false : 1;
}
if (m.getName().equals("executeLargeBatch")) {
writes.incrementAndGet();
if (cancelWhenLargeBatchUnsupported) {
server.cancelActiveStatements();
throw new SQLFeatureNotSupportedException();
}
if (timeout.get() != expectedTimeout) throw new SQLTimeoutException("Statement has no requested execution timeout");
if (block) { started.countDown(); assertTrue(release.await(3, TimeUnit.SECONDS)); }
return new long[]{1, 1};
}
if (m.getName().equals("executeBatch")) { fallbackWrites.incrementAndGet(); return new int[]{1, 1}; }
return defaultValue(m.getReturnType());
});
}
}
private static Object defaultValue(Class<?> type) {
if (type == boolean.class) return false;
if (type == int.class) return 0;
if (type == long.class) return 0L;
return null;
}
}
@@ -322,6 +322,14 @@ public final class OceanBaseOracleAgent extends ConfiguredJdbcAgent {
return buildUrl(params);
}
@Override
protected java.util.Properties buildConnectionProperties(ConnectParams params) {
java.util.Properties properties = super.buildConnectionProperties(params);
// Connector/J must stop after a failed or timed-out batch item.
properties.setProperty("continueBatchOnError", "false");
return properties;
}
static String buildUrl(ConnectParams params) {
return appendDefaultCompatibilityOption(OCEANBASE_ORACLE_PROFILE.buildUrl(params));
}
@@ -0,0 +1,17 @@
package com.dbx.agent.oceanbaseoracle;
import com.dbx.agent.ConnectParams;
import com.oceanbase.jdbc.UrlParser;
import org.junit.jupiter.api.Test;
import static org.junit.jupiter.api.Assertions.*;
class OceanBaseBatchOptionsTest {
@Test
void agentFailFastPropertyOverridesExplicitContinueBatchOnErrorUrl() throws Exception {
ConnectParams params = new ConnectParams();
params.setConnection_string("jdbc:oceanbase://localhost:2881/SYS?continueBatchOnError=true");
OceanBaseOracleAgent agent = new OceanBaseOracleAgent();
assertFalse(UrlParser.parse(OceanBaseOracleAgent.buildUrl(params), agent.buildConnectionProperties(params))
.getOptions().continueBatchOnError);
}
}
+1
View File
@@ -1,6 +1,7 @@
dependencies {
implementation files('libs/inceptor-sdk-transwarp-6.1.0-SNAPSHOT.jar')
testImplementation 'com.h2database:h2:2.4.240'
testImplementation 'com.google.code.gson:gson:2.12.1'
}
tasks.named('shadowJar') {
@@ -7,6 +7,7 @@ import com.dbx.agent.DatabaseInfo;
import com.dbx.agent.ForeignKeyInfo;
import com.dbx.agent.IndexInfo;
import com.dbx.agent.JdbcIdentifiers;
import com.dbx.agent.JdbcExecutor;
import com.dbx.agent.MetadataListConstraints;
import com.dbx.agent.MultiSessionJsonRpcServer;
import com.dbx.agent.ObjectInfo;
@@ -357,28 +358,37 @@ public final class TranswarpAgent extends AbstractJdbcAgent {
long started = System.currentTimeMillis();
long affected = 0;
Connection connection = requireConnected();
applySchemaContext(connection, schema);
try (Statement statement = connection.createStatement()) {
if (connection instanceof HiveConnection) {
statement.execute("SET transaction.type = inceptor");
}
statement.execute("BEGIN TRANSACTION");
try {
for (String sql : statements) {
if (sql != null && !sql.isBlank()) {
statement.execute(sql);
affected += Math.max(statement.getUpdateCount(), 0);
}
try (JdbcExecutor.NativeOperation operation = JdbcExecutor.current().beginNativeOperation()) {
operation.checkCancelled();
applySchemaContext(connection, schema);
try (var tracked = JdbcExecutor.current().trackStatement(connection.createStatement())) {
Statement statement = tracked.statement();
JdbcExecutor.current().applyStatementTimeout(statement);
operation.checkCancelled();
if (connection instanceof HiveConnection) {
statement.execute("SET transaction.type = inceptor");
}
statement.execute("COMMIT");
} catch (Exception error) {
operation.checkCancelled();
try {
statement.execute("ROLLBACK");
} catch (Exception rollbackError) {
transactionStateUnknown = true;
error.addSuppressed(rollbackError);
statement.execute("BEGIN TRANSACTION");
for (String sql : statements) {
operation.checkCancelled();
if (sql != null && !sql.isBlank()) {
statement.execute(sql);
affected += Math.max(statement.getUpdateCount(), 0);
}
}
operation.checkCancelled();
statement.execute("COMMIT");
} catch (Exception error) {
try {
statement.execute("ROLLBACK");
} catch (Exception rollbackError) {
transactionStateUnknown = true;
error.addSuppressed(rollbackError);
}
throw error;
}
throw error;
}
}
return new QueryResult(List.of(), List.of(), affected, System.currentTimeMillis() - started);
@@ -0,0 +1,112 @@
package com.dbx.agent;
import com.dbx.agent.transwarp.TranswarpAgent;
import com.google.gson.JsonParser;
import org.junit.jupiter.api.Test;
import java.lang.reflect.Proxy;
import java.sql.Connection;
import java.sql.Statement;
import java.util.List;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;
import static org.junit.jupiter.api.Assertions.*;
class TranswarpTransactionCancellationTest {
@Test
void rpcTimeoutReachesTheCustomTransactionStatement() throws Exception {
Fixture fixture = new Fixture();
fixture.run(7);
assertEquals(7, fixture.timeout.get());
fixture.run(0);
assertEquals(0, fixture.timeout.get());
assertEquals(2, fixture.commits.get());
}
@Test
void cancellationRejectsLateBeginOrWriteSuccessWithoutCommitting() throws Exception {
assertCancellation(false);
assertCancellation(true);
}
private void assertCancellation(boolean failOnCancel) throws Exception {
for (String blockedSql : List.of("BEGIN TRANSACTION", "UPDATE T SET C=1")) {
Fixture fixture = new Fixture();
fixture.blockedSql = blockedSql;
fixture.failOnCancel = failOnCancel;
ExecutorService worker = Executors.newSingleThreadExecutor();
try {
Future<?> result = worker.submit(() -> { fixture.run(7); return null; });
assertTrue(fixture.started.await(3, TimeUnit.SECONDS));
fixture.server.cancelActiveStatements();
fixture.release.countDown();
assertThrows(ExecutionException.class, () -> result.get(3, TimeUnit.SECONDS));
assertEquals(1, fixture.cancels.get());
assertEquals(blockedSql.startsWith("BEGIN") ? 0 : 1, fixture.writes.get());
assertEquals(0, fixture.commits.get());
assertEquals(1, fixture.rollbacks.get());
fixture.blockedSql = null;
fixture.run(0);
assertEquals(1, fixture.commits.get());
} finally {
fixture.release.countDown();
worker.shutdownNow();
}
}
}
private static final class Fixture {
String blockedSql;
boolean failOnCancel;
final AtomicInteger timeout = new AtomicInteger();
final AtomicInteger writes = new AtomicInteger();
final AtomicInteger commits = new AtomicInteger();
final AtomicInteger rollbacks = new AtomicInteger();
final AtomicInteger cancels = new AtomicInteger();
final CountDownLatch started = new CountDownLatch(1);
final CountDownLatch release = new CountDownLatch(1);
final JsonRpcServer server;
Fixture() throws Exception {
Connection connection = (Connection) Proxy.newProxyInstance(getClass().getClassLoader(), new Class<?>[]{Connection.class}, (proxy, method, args) -> {
if (method.getName().equals("createStatement")) {
timeout.set(0);
return Proxy.newProxyInstance(getClass().getClassLoader(), new Class<?>[]{Statement.class}, (stmt, call, values) -> {
if (call.getName().equals("setQueryTimeout")) { timeout.set((int) values[0]); return null; }
if (call.getName().equals("cancel")) { cancels.incrementAndGet(); return null; }
if (call.getName().equals("getUpdateCount")) return 1;
if (call.getName().equals("execute")) {
String sql = (String) values[0];
if (sql.startsWith("UPDATE")) writes.incrementAndGet();
if (sql.equals("COMMIT")) commits.incrementAndGet();
if (sql.equals("ROLLBACK")) rollbacks.incrementAndGet();
if (sql.equals(blockedSql)) {
started.countDown();
assertTrue(release.await(3, TimeUnit.SECONDS));
if (failOnCancel) throw new java.sql.SQLTimeoutException("Statement cancelled");
}
}
return defaultValue(call.getReturnType());
});
}
return defaultValue(method.getReturnType());
});
TranswarpAgent agent = new TranswarpAgent();
// Inject only the external JDBC boundary; dispatch and driver override are real.
var field = AbstractJdbcAgent.class.getDeclaredField("connection");
field.setAccessible(true);
field.set(agent, connection);
server = new JsonRpcServer(agent);
}
void run(int seconds) throws Exception {
server.dispatchForRuntime("execute_transaction", JsonParser.parseString("{\"statements\":[\"UPDATE T SET C=1\",\"UPDATE T SET C=2\"],\"timeoutSecs\":" + seconds + "}").getAsJsonObject());
}
}
private static Object defaultValue(Class<?> type) {
if (type == boolean.class) return false;
if (type == int.class) return 0;
if (type == long.class) return 0L;
return null;
}
}
@@ -0,0 +1,19 @@
import { describe, expect, it } from "vitest";
import { extractSqlParameters, substituteSqlParameters } from "@/lib/sql/sqlParameters";
import { expandSqlVariables } from "@/lib/sql/sqlVariables";
describe.each(["oracle", "oceanbase-oracle"] as const)("%s remote object parameters", (databaseType) => {
it.each([
"BEGIN MERGE INTO local_t t USING app.remote_t@LINK s ON (t.id=s.id) WHEN MATCHED THEN UPDATE SET t.v=@value; END;",
"SELECT seq.NEXTVAL@LINK, @value FROM dual",
"SELECT app.seq.CURRVAL@LINK, @value FROM dual",
"CREATE SYNONYM local_t FOR app.remote_t@LINK; SELECT @value FROM dual",
"MERGE INTO local_t t USING app.remote_t@LINK s ON (t.id=s.id) WHEN MATCHED THEN UPDATE SET t.v=@value",
])("preserves database links in %s", (sql) => {
const options = { databaseType };
expect(extractSqlParameters(sql, options)).toEqual(["value"]);
expect(substituteSqlParameters(sql, { value: { kind: "number", value: "7" }, LINK: { kind: "string", value: "wrong" } }, options)).toContain("@LINK");
const expanded = expandSqlVariables(`@set LINK = wrong;\n@set value = 7;\n${sql}`, options);
expect(expanded.sql).toBe(sql.replace("@value", "7"));
});
});
@@ -80,6 +80,20 @@ export function collectSqlDatabaseLinkStarts(sql: string, databaseType?: Databas
const sqlXmlExpression = objectStart === i - 1 && sqlXmlKeywords.has(objectStart);
const routineCall = !proceduralExpression && !sqlXmlExpression && tokens[i + 1]?.text === "(" && (object.kind === "quoted_identifier" || !isReservedKeyword(object.text));
let objectContext = (ORACLE_OBJECT_PREFIX_KEYWORDS.has(prefix) && !(prefix === "update" && tokens[objectStart - 2]?.normalized === "for")) || routineCall;
const sequenceValue = object.kind === "word" && (object.normalized === "nextval" || object.normalized === "currval") && tokens[i - 2]?.text === "." && isObjectIdentifier(i - 3);
if (sequenceValue) objectContext = true;
if (prefix === "for" || prefix === "using") {
const statementWords: string[] = [];
for (let j = objectStart - 2; j >= 0; j -= 1) {
if (tokens[j].text === ";" || tokens[j].depth < token.depth) break;
if (tokens[j].depth === token.depth && tokens[j].kind === "word") statementWords.unshift(tokens[j].normalized);
}
if (prefix === "for") objectContext ||= statementWords[0] === "create" && statementWords.includes("synonym");
else {
const mergeIndex = statementWords.lastIndexOf("merge");
objectContext ||= mergeIndex >= 0 && statementWords[mergeIndex + 1] === "into";
}
}
// A comma can separate tables or expressions. Only a FROM list makes the
// following non-reserved keyword an object name, including after subqueries.
if (prefix === ",") {
@@ -58,7 +58,101 @@ describe.each(["oracle", "oceanbase-oracle"] as const)("queryStore %s commit lif
afterEach(() => vi.unstubAllGlobals());
it("preserves a synchronous backend error and permits the next transaction to commit", async () => {
const { store, id, tab } = await setupManualTab(dbType);
const error = new Error("Synchronous commit boundary failure");
mocks.commitManualTransaction.mockImplementationOnce(() => {
throw error;
});
await expect(store.commitTransaction(id)).rejects.toBe(error);
expect(tab.txnSessionId).toBeUndefined();
expect(tab.txnPossiblyDirty).toBe(false);
mocks.beginManualTransaction.mockResolvedValueOnce("txn-next");
await store.ensureManualTransactionSession(id, "ORCL", "APP");
mocks.commitManualTransaction.mockResolvedValueOnce(undefined);
await expect(store.commitTransaction(id)).resolves.toBeUndefined();
expect(mocks.commitManualTransaction.mock.calls.map(([sessionId]) => sessionId)).toEqual(["txn-old", "txn-next"]);
expect(tab.txnSessionId).toBeUndefined();
});
it.each(["success", "failure"] as const)("shares an in-flight commit until its %s result arrives", async (outcome) => {
const { store, id, tab } = await setupManualTab(dbType);
const commit = deferred<void>();
mocks.commitManualTransaction.mockReturnValueOnce(commit.promise).mockRejectedValue(new Error("Transaction session not found"));
const first = store.commitTransaction(id);
const second = store.commitTransaction(id);
const secondResult = second.then(
() => undefined,
(error) => error,
);
await Promise.resolve();
await Promise.resolve();
expect(mocks.commitManualTransaction).toHaveBeenCalledTimes(1);
expect(tab.txnSessionId).toBe("txn-old");
expect(tab.txnPossiblyDirty).toBe(true);
await finishCommit(commit, first, outcome);
expect(await secondResult).toEqual(outcome === "success" ? undefined : new Error("Commit response lost"));
expect(tab.txnSessionId).toBeUndefined();
});
it("tracks a replacement commit independently while the old commit finishes", async () => {
const { store, id, tab } = await setupManualTab(dbType);
const oldCommit = deferred<void>();
const newCommit = deferred<void>();
mocks.commitManualTransaction.mockReturnValueOnce(oldCommit.promise).mockReturnValueOnce(newCommit.promise);
const oldCommitting = store.commitTransaction(id);
store.updateSchema(id, "OTHER");
mocks.beginManualTransaction.mockResolvedValueOnce("txn-new");
await store.ensureManualTransactionSession(id, "ORCL", "OTHER");
store.markManualTransactionDirty(id);
const newCommitting = store.commitTransaction(id);
oldCommit.resolve();
await oldCommitting;
const duplicate = store.commitTransaction(id);
expect(mocks.commitManualTransaction.mock.calls.map(([sessionId]) => sessionId)).toEqual(["txn-old", "txn-new"]);
expect(tab.txnSessionId).toBe("txn-new");
expect(tab.txnPossiblyDirty).toBe(true);
newCommit.resolve();
await Promise.all([newCommitting, duplicate]);
expect(tab.txnSessionId).toBeUndefined();
});
describe.each(["success", "failure"] as const)("old commit %s", (outcome) => {
it("keeps a reopened tab's new transaction independent of the old pending commit", async () => {
const { store, id, tab } = await setupManualTab(dbType);
const oldCommit = deferred<void>();
const newCommit = deferred<void>();
mocks.commitManualTransaction.mockReturnValueOnce(oldCommit.promise).mockReturnValueOnce(newCommit.promise);
const oldCommitting = store.commitTransaction(id);
store.closeTab(id);
expect(store.reopenClosedTab()).toBe(id);
const reopened = store.tabs.find((candidate) => candidate.id === id)!;
expect(reopened).not.toBe(tab);
expect(reopened.txnSessionId).toBeUndefined();
store.setAutoCommit(id, false);
mocks.beginManualTransaction.mockResolvedValueOnce("txn-reopened");
await store.ensureManualTransactionSession(id, "ORCL", "APP");
store.markManualTransactionDirty(id);
const newCommitting = store.commitTransaction(id);
await finishCommit(oldCommit, oldCommitting, outcome);
expect(reopened.txnSessionId).toBe("txn-reopened");
expect(reopened.txnPossiblyDirty).toBe(true);
const duplicate = store.commitTransaction(id);
expect(mocks.commitManualTransaction.mock.calls.map(([sessionId]) => sessionId)).toEqual(["txn-old", "txn-reopened"]);
newCommit.resolve();
await Promise.all([newCommitting, duplicate]);
expect(reopened.txnSessionId).toBeUndefined();
expect(reopened.txnPossiblyDirty).toBe(false);
});
it.each(["mode", "schema", "database", "connection"] as const)("preserves a replacement transaction after a %s change", async (change) => {
const { store, id, tab } = await setupManualTab(dbType);
const commit = deferred<void>();
+22 -6
View File
@@ -5565,6 +5565,8 @@ export const useQueryStore = defineStore("query", () => {
clearAutoCommitOpenTransaction(tab);
}
const pendingManualTransactionCommits = new WeakMap<QueryTab, { sessionId: string; promise: Promise<void> }>();
async function commitTransaction(id: string) {
const tab = tabs.value.find((t) => t.id === id);
if (!tab) return;
@@ -5605,12 +5607,26 @@ export const useQueryStore = defineStore("query", () => {
return ending;
}
const sessionId = tab.txnSessionId;
try {
await api.commitManualTransaction(sessionId);
} finally {
// A mode or target switch may have started a replacement transaction.
if (tab.txnSessionId === sessionId) clearManualTransactionSession(tab);
}
const pending = pendingManualTransactionCommits.get(tab);
if (pending?.sessionId === sessionId) return pending.promise;
let resolveCommit!: () => void;
let rejectCommit!: (error: unknown) => void;
const committing = new Promise<void>((resolve, reject) => {
resolveCommit = resolve;
rejectCommit = reject;
});
const entry = { sessionId, promise: committing };
pendingManualTransactionCommits.set(tab, entry);
void (async () => {
try {
await api.commitManualTransaction(sessionId);
} finally {
// A mode or target switch may have started a replacement transaction.
if (tab.txnSessionId === sessionId) clearManualTransactionSession(tab);
if (pendingManualTransactionCommits.get(tab) === entry) pendingManualTransactionCommits.delete(tab);
}
})().then(resolveCommit, rejectCommit);
return committing;
}
async function rollbackTransaction(id: string) {
+52 -4
View File
@@ -2825,7 +2825,7 @@ impl AgentDriverClient {
self.invalidate_cached_query();
self.call_method_with_timeout(
AgentMethod::ExecuteTransaction,
agent_transaction_params(database, statements, schema),
agent_statements_with_timeout_params(database, statements, schema, timeout_duration),
timeout_duration,
)
.await
@@ -2841,7 +2841,7 @@ impl AgentDriverClient {
self.invalidate_cached_query();
self.call_method_typed_with_timeout(
AgentMethod::ExecuteTransaction,
agent_transaction_params(database, statements, schema),
agent_statements_with_timeout_params(database, statements, schema, timeout_duration),
timeout_duration,
)
.await
@@ -2875,7 +2875,7 @@ impl AgentDriverClient {
self.invalidate_cached_query();
self.call_method_with_timeout(
AgentMethod::ExecuteBatch,
agent_transaction_params(database, statements, schema),
agent_statements_with_timeout_params(database, statements, schema, timeout_duration),
timeout_duration,
)
.await
@@ -2891,7 +2891,7 @@ impl AgentDriverClient {
self.invalidate_cached_query();
self.call_method_typed_with_timeout(
AgentMethod::ExecuteBatch,
agent_transaction_params(database, statements, schema),
agent_statements_with_timeout_params(database, statements, schema, timeout_duration),
timeout_duration,
)
.await
@@ -3340,6 +3340,18 @@ pub fn agent_transaction_params(database: Option<&str>, statements: &[String], s
})
}
fn agent_statements_with_timeout_params(
database: Option<&str>,
statements: &[String],
schema: Option<&str>,
timeout: Option<Duration>,
) -> Value {
let mut params = agent_transaction_params(database, statements, schema);
params["timeoutSecs"] =
serde_json::json!(timeout.map(|value| value.as_secs().clamp(1, i32::MAX as u64)).unwrap_or(0));
params
}
pub fn agent_manual_transaction_params(schema: Option<&str>) -> Value {
let schema = schema.map(str::trim).filter(|schema| !schema.is_empty());
match schema {
@@ -5471,6 +5483,42 @@ for line in sys.stdin:
let _ = std::fs::remove_file(script_path);
}
#[tokio::test]
async fn batch_and_transaction_send_statement_timeout_to_agent() {
let dir = tempfile::tempdir().unwrap();
let script = dir.path().join("echo.py");
std::fs::write(
&script,
r#"import json, sys
print(json.dumps({'ready': True}), flush=True)
for line in sys.stdin:
req = json.loads(line)
print(json.dumps({'jsonrpc': '2.0', 'id': req['id'], 'result': req['params']}), flush=True)
"#,
)
.unwrap();
let mut client = AgentDriverClient::spawn(
AgentLaunchSpec::new(test_python()).with_args([script.to_string_lossy().to_string()]),
)
.await
.unwrap();
let statements = ["UPDATE t SET a = 1".to_string()];
for timeout in [Some(Duration::from_secs(7)), None] {
let expected = if timeout.is_some() { 7 } else { 0 };
let results: [serde_json::Value; 4] = [
client.execute_batch(Some("db"), &statements, Some("schema"), timeout).await.unwrap(),
client.execute_batch_typed(Some("db"), &statements, Some("schema"), timeout).await.unwrap(),
client.execute_transaction(Some("db"), &statements, Some("schema"), timeout).await.unwrap(),
client.execute_transaction_typed(Some("db"), &statements, Some("schema"), timeout).await.unwrap(),
];
for result in results {
assert_eq!(result["timeoutSecs"], expected);
assert_eq!(result["statements"], serde_json::json!(statements));
}
}
client.kill();
}
#[tokio::test]
async fn legacy_client_tolerates_post_startup_driver_stdout_noise() {
let script_path =
+235 -31
View File
@@ -144,6 +144,7 @@ pub struct SqlFileImportStatement {
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct SqlDialectProfile {
supports_hash_line_comments: bool,
supports_backtick_identifiers: bool,
/// Whether `/* ... */` comments nest. PostgreSQL and SQL Server document
/// nesting (`/* /* */ */` needs both closers, so commenting out a block that
/// already contains a comment keeps the whole block commented); MySQL
@@ -188,6 +189,7 @@ impl Default for SqlDialectProfile {
fn default() -> Self {
Self {
supports_hash_line_comments: false,
supports_backtick_identifiers: true,
supports_nested_block_comments: false,
supports_oracle_q_quotes: false,
supports_backslash_escaped_quotes: true,
@@ -228,6 +230,15 @@ impl SqlDialectProfile {
return Self::postgres_family();
}
if matches!(db_type, DatabaseType::OceanbaseOracle) {
return Self {
supports_backtick_identifiers: false,
supports_dollar_quoted_strings: false,
supports_custom_delimiter_commands: false,
..Self::oracle_like()
};
}
if Self::is_oracle_like_database(db_type) {
return Self::oracle_like();
}
@@ -588,28 +599,25 @@ fn ends_with_xml_entity_reference(buffer: &str) -> bool {
matches!(name, "lt" | "gt" | "amp" | "quot" | "apos")
}
/// Tracks an Oracle `q'...'` literal across characters, so a semicolon inside its
/// Tracks an Oracle `q'...'` or `nq'...'` literal across characters, so a semicolon inside its
/// text never splits the statement.
#[derive(Default, Clone, Copy)]
struct OracleQQuoteState {
/// 1 while the introducer's quote is expected, 2 while its delimiter is.
awaiting: u8,
awaiting_delimiter: bool,
closer: Option<char>,
}
impl OracleQQuoteState {
fn is_open(&self) -> bool {
self.awaiting > 0 || self.closer.is_some()
self.awaiting_delimiter || self.closer.is_some()
}
/// Consumes one character of an open literal; `previous` is the character
/// before it, as tracked by the caller.
fn consume(&mut self, ch: char, previous: Option<char>) {
if self.awaiting == 1 {
self.awaiting = if ch == '\'' { 2 } else { 0 };
} else if self.awaiting == 2 {
if self.awaiting_delimiter {
self.closer = oracle_q_quote_closer(ch);
self.awaiting = 0;
self.awaiting_delimiter = false;
} else if let Some(closer) = self.closer {
if previous == Some(closer) && ch == '\'' {
self.closer = None;
@@ -756,11 +764,10 @@ impl SqlStatementSplitter {
if !self.in_single_quote && !self.in_double_quote && !self.in_backtick {
if self.options.profile.supports_oracle_q_quotes
&& matches!(ch, 'q' | 'Q')
&& next == Some('\'')
&& !self.previous.is_some_and(is_identifier_continue_char)
&& ch == '\''
&& ends_with_oracle_q_quote_prefix(&self.buffer)
{
self.oracle_q_quote.awaiting = 1;
self.oracle_q_quote.awaiting_delimiter = true;
self.buffer.push(ch);
self.previous = Some(ch);
i += 1;
@@ -849,7 +856,10 @@ impl SqlStatementSplitter {
self.in_double_quote = !self.in_double_quote;
self.buffer.push(ch);
}
'`' if !self.in_single_quote && !self.in_double_quote => {
'`' if self.options.profile.supports_backtick_identifiers
&& !self.in_single_quote
&& !self.in_double_quote =>
{
self.in_backtick = !self.in_backtick;
self.buffer.push(ch);
}
@@ -1073,6 +1083,19 @@ fn is_identifier_continue_char(ch: char) -> bool {
ch.is_alphanumeric() || ch == '_' || ch == '$'
}
/// Whether the text before a quote ends in an Oracle `q`/`nq` introducer token.
fn ends_with_oracle_q_quote_prefix(text: &str) -> bool {
let mut chars = text.chars().rev();
if !matches!(chars.next(), Some('q' | 'Q')) {
return false;
}
let mut previous = chars.next();
if matches!(previous, Some('n' | 'N')) {
previous = chars.next();
}
!previous.is_some_and(is_identifier_continue_char)
}
/// The closing character of an Oracle alternative-quoting delimiter: brackets
/// pair up (`q'[a]'`) and every other character delimits itself (`q'!a!'`).
/// A quote or whitespace cannot delimit the literal.
@@ -1388,12 +1411,8 @@ fn split_sql_statement_ranges_with_options(sql: &str, options: SqlParsingOptions
}
if !in_single_quote && !in_double_quote && !in_backtick {
if options.profile.supports_oracle_q_quotes
&& matches!(ch, 'q' | 'Q')
&& next == Some('\'')
&& !sql[..i].chars().next_back().is_some_and(is_identifier_continue_char)
{
if let Some(end) = oracle_q_quote_end(sql, i + ch.len_utf8()) {
if options.profile.supports_oracle_q_quotes && ch == '\'' && ends_with_oracle_q_quote_prefix(&sql[..i]) {
if let Some(end) = oracle_q_quote_end(sql, i) {
i = end;
continue;
}
@@ -1488,7 +1507,7 @@ fn split_sql_statement_ranges_with_options(sql: &str, options: SqlParsingOptions
in_double_quote = !in_double_quote;
i += ch.len_utf8();
}
'`' if !in_single_quote && !in_double_quote => {
'`' if options.profile.supports_backtick_identifiers && !in_single_quote && !in_double_quote => {
in_backtick = !in_backtick;
i += ch.len_utf8();
}
@@ -2470,6 +2489,7 @@ fn parse_insert_values_tail(tail: &str) -> Option<String> {
struct SqlScanner {
profile: SqlDialectProfile,
in_single_quote: bool,
single_quote_escape_string: bool,
in_double_quote: bool,
in_backtick: bool,
in_line_comment: bool,
@@ -2521,12 +2541,8 @@ impl SqlScanner {
}
if !self.in_single_quote && !self.in_double_quote && !self.in_backtick {
if self.profile.supports_oracle_q_quotes
&& matches!(ch, 'q' | 'Q')
&& next == Some('\'')
&& !self.previous.is_some_and(is_identifier_continue_char)
{
self.oracle_q_quote.awaiting = 1;
if self.profile.supports_oracle_q_quotes && ch == '\'' && ends_with_oracle_q_quote_prefix(&sql[..idx]) {
self.oracle_q_quote.awaiting_delimiter = true;
self.previous = Some(ch);
return;
}
@@ -2544,14 +2560,29 @@ impl SqlScanner {
}
}
let single_quote_backslash_escapes = if self.in_single_quote {
self.single_quote_escape_string
} else {
self.profile.supports_backslash_escaped_quotes
};
match ch {
'\'' if !self.in_double_quote && !self.in_backtick && self.previous != Some('\\') => {
'\'' if !self.in_double_quote
&& !self.in_backtick
&& !(single_quote_backslash_escapes && self.previous == Some('\\')) =>
{
self.in_single_quote = !self.in_single_quote;
self.single_quote_escape_string = self.in_single_quote
&& (self.profile.supports_backslash_escaped_quotes
|| (self.profile.supports_postgres_escape_strings
&& ends_with_escape_string_prefix(&sql[..idx])));
}
'"' if !self.in_single_quote && !self.in_backtick && self.previous != Some('\\') => {
'"' if !self.in_single_quote
&& !self.in_backtick
&& !(self.profile.supports_backslash_escaped_quotes && self.previous == Some('\\')) =>
{
self.in_double_quote = !self.in_double_quote;
}
'`' if !self.in_single_quote && !self.in_double_quote => {
'`' if self.profile.supports_backtick_identifiers && !self.in_single_quote && !self.in_double_quote => {
self.in_backtick = !self.in_backtick;
}
_ => {}
@@ -4056,9 +4087,36 @@ mod tests {
);
}
#[test]
fn oracle_national_alternative_quoting_preserves_statement_boundaries() {
for database in [DatabaseType::Oracle, DatabaseType::OceanbaseOracle] {
let sql = "UPDATE events SET name = nq'{it's; DELETE FROM events; --}'; SELECT 1 FROM dual;";
assert_eq!(
split_sql_statements_for_database(sql, database),
vec!["UPDATE events SET name = nq'{it's; DELETE FROM events; --}'", "SELECT 1 FROM dual"]
);
let sql = "UPDATE events SET name = nq'[a'b]'; DELETE FROM events; --'";
assert_eq!(
split_sql_statements_for_database(sql, database),
vec!["UPDATE events SET name = nq'[a'b]'", "DELETE FROM events"]
);
}
}
#[test]
fn oracle_alternative_quoting_supports_every_delimiter_form() {
for oracle_literal in ["q'{a;b}'", "q'(a;b)'", "q'<a;b>'", "q'!a;b!'", "q'Aa;bA'", "Q'[i's; x]'"] {
for oracle_literal in [
"q'{a;b}'",
"q'(a;b)'",
"q'<a;b>'",
"q'!a;b!'",
"q'Aa;bA'",
"Q'[i's; x]'",
"nq'[i's; x]'",
"NQ'<i's; x>'",
"nq'!i's; x!'",
"nQ'(i's; x)'",
] {
let sql = format!("SELECT {oracle_literal} FROM dual; SELECT 1 FROM dual;");
assert_eq!(
@@ -4069,6 +4127,101 @@ mod tests {
}
}
#[test]
fn oracle_alternative_quoting_stays_intact_across_chunk_boundaries() {
for literal in ["q'[it's; one]'", "nq'{it's; one}'", "NQ'<it's; one>'", "nQ'!it's; one!'"] {
let statement = format!("UPDATE events SET name = {literal}");
let sql = format!("{statement}; DELETE FROM events;");
for boundary in 0..=sql.len() {
let mut splitter =
SqlStatementSplitter::with_options(SqlParsingOptions::for_database_type(DatabaseType::Oracle));
let mut statements = splitter.push_chunk(&sql[..boundary]);
statements.extend(splitter.push_chunk(&sql[boundary..]));
statements.extend(splitter.finish());
assert_eq!(
statements,
vec![statement.clone(), "DELETE FROM events".to_string()],
"boundary {boundary}"
);
}
}
}
#[test]
fn oracle_alternative_quoting_requires_an_introducer_token() {
for prefix in ["aq", "anq", "_nq", "column$nq"] {
let sql = format!("SELECT {prefix}'[a'b]'; DELETE FROM events; --'");
assert_eq!(split_sql_statements_for_database(&sql, DatabaseType::Oracle), vec![sql]);
}
}
#[test]
fn non_oracle_dialects_do_not_interpret_national_alternative_quoting() {
let sql = "SELECT nq'[a'b]'; DELETE FROM events; --'";
for database in [DatabaseType::Mysql, DatabaseType::Postgres, DatabaseType::Sqlite, DatabaseType::SqlServer] {
assert_eq!(split_sql_statements_for_database(sql, database), vec![sql]);
}
}
#[test]
fn oceanbase_oracle_unsupported_quoting_does_not_hide_statement_boundaries() {
for (sql, expected) in [
("UPDATE t SET x=$$;DELETE FROM t;$$", vec!["UPDATE t SET x=$$", "DELETE FROM t", "$$"]),
("UPDATE t SET x='a\\'; DELETE FROM t; --'", vec!["UPDATE t SET x='a\\'", "DELETE FROM t"]),
("UPDATE t SET x=`a;DELETE FROM t;`", vec!["UPDATE t SET x=`a", "DELETE FROM t", "`"]),
(
"DELIMITER $$\nUPDATE t SET x=1; DELETE FROM t;$$",
vec!["DELIMITER $$\nUPDATE t SET x=1", "DELETE FROM t", "$$"],
),
] {
assert_eq!(split_sql_statements_for_database(sql, DatabaseType::OceanbaseOracle), expected, "{sql}");
}
}
#[test]
fn oceanbase_oracle_cursor_uses_its_own_quote_boundaries() {
for sql in [
"UPDATE t SET x=$$; DELETE FROM t;$$",
"UPDATE t SET x='a\\'; DELETE FROM t; --'",
"UPDATE t SET x=`a; DELETE FROM t;`",
"DELIMITER $$\nUPDATE t SET x=1; DELETE FROM t;$$",
] {
let cursor = sql.find("DELETE").unwrap() + 3;
assert_eq!(
find_statement_at_cursor_for_database(sql, cursor, DatabaseType::OceanbaseOracle),
"DELETE FROM t",
"{sql}"
);
}
for sql in [
"SELECT 'a\\'\n\n\nUPDATE t SET x=2;",
"SELECT `x\n\n\nUPDATE t SET x=2;",
"SELECT $$\n\n\nUPDATE t SET x=2;",
] {
let cursor = sql.find("UPDATE").unwrap() + 3;
assert_eq!(
find_statement_at_cursor_for_database(sql, cursor, DatabaseType::OceanbaseOracle),
"UPDATE t SET x=2",
"{sql}"
);
}
}
#[test]
fn oceanbase_oracle_quote_rules_do_not_change_other_dialects() {
for database in [DatabaseType::Oracle, DatabaseType::Dameng, DatabaseType::Gaussdb, DatabaseType::Mysql] {
for sql in ["UPDATE t SET x=$$;DELETE FROM t;$$", "UPDATE t SET x=`a;DELETE FROM t;`"] {
assert_eq!(split_sql_statements_for_database(sql, database), vec![sql]);
}
}
let statement = "SELECT E'a\\'\n\n\nDELETE FROM t' FROM t";
let sql = format!("{statement}; SELECT 2;");
for database in [DatabaseType::Postgres, DatabaseType::Gaussdb] {
let cursor = sql.find("DELETE").unwrap() + 3;
assert_eq!(find_statement_at_cursor_for_database(&sql, cursor, database), statement);
}
}
#[test]
fn dameng_alternative_quoting_matches_oracle() {
let sql = "SELECT q'[it's; x]' FROM dual; SELECT 1 FROM dual;";
@@ -4097,6 +4250,37 @@ mod tests {
assert_eq!(find_statement_at_cursor_for_database(sql, cursor, DatabaseType::Oracle), "SELECT 1 FROM dual");
}
#[test]
fn oracle_like_cursor_recognizes_blank_line_boundary_after_backslash() {
let sql = "SELECT 'a\\'\n\n\nUPDATE t SET x=2;";
for database in [
DatabaseType::Oracle,
DatabaseType::OceanbaseOracle,
DatabaseType::Dameng,
DatabaseType::Yashandb,
DatabaseType::Oscar,
] {
let cursor = sql.find("UPDATE").unwrap() + 3;
assert_eq!(
find_statement_at_cursor_for_database(sql, cursor, database),
"UPDATE t SET x=2",
"{database:?}"
);
}
}
#[test]
fn cursor_statement_keeps_national_alternative_quoting_and_blank_lines_together() {
let statement = "UPDATE events SET name = NQ'<it's;\n\n\nDELETE FROM events; -->' WHERE id = 1";
let sql = format!("{statement};\nSELECT 1 FROM dual;");
for database in [DatabaseType::Oracle, DatabaseType::OceanbaseOracle] {
let cursor = sql.find("DELETE").unwrap() + 3;
assert_eq!(find_statement_at_cursor_for_database(&sql, cursor, database), statement);
let cursor = sql.find("SELECT 1").unwrap() + 3;
assert_eq!(find_statement_at_cursor_for_database(&sql, cursor, database), "SELECT 1 FROM dual");
}
}
#[test]
fn emits_trailing_statement_without_semicolon() {
assert_eq!(
@@ -4765,7 +4949,7 @@ SELECT 2;";
DatabaseType::OceanbaseOracle,
] {
let profile = SqlDialectProfile::for_database_type(db_type);
assert_eq!(profile, SqlDialectProfile::oracle_like());
assert!(profile.supports_oracle_q_quotes);
assert!(!profile.supports_backslash_escaped_quotes);
assert!(profile.supports_oracle_plsql_blocks);
assert!(profile.supports_slash_line_block_delimiter);
@@ -4915,6 +5099,26 @@ END \"T\";";
);
}
#[test]
fn oracle_split_keeps_compound_trigger_sections_in_one_execution_unit() {
let trigger = "CREATE OR REPLACE TRIGGER \"S\".\"tr\" FOR INSERT ON \"S\".\"t\" COMPOUND TRIGGER\nBEFORE STATEMENT IS BEGIN NULL; END BEFORE STATEMENT;\nAFTER EACH ROW IS BEGIN :NEW.x := q'[a;b]'; END AFTER EACH ROW;\nEND tr;";
let state = "ALTER TRIGGER \"S\".\"tr\" DISABLE";
for database in [DatabaseType::Oracle, DatabaseType::OceanbaseOracle] {
assert_eq!(split_sql_statements_for_database(trigger, database), vec![trigger.to_string()]);
assert_eq!(
split_sql_statements_for_database(
&format!("{trigger}\n/\n{state};\nCREATE TABLE t2 (id NUMBER);"),
database
),
vec![trigger.to_string(), state.to_string(), "CREATE TABLE t2 (id NUMBER)".to_string()]
);
let execution_units = [trigger.to_string(), format!("{state};")];
let split_units: Vec<_> =
execution_units.iter().flat_map(|unit| split_sql_statements_for_database(unit, database)).collect();
assert_eq!(split_units, vec![trigger.to_string(), state.to_string()]);
}
}
#[test]
fn oracle_like_split_keeps_anonymous_plsql_block_together() {
let sql = "\
+81
View File
@@ -0,0 +1,81 @@
# SQL 识别与执行回归
当前整合基线:`becccf108f2a91c99e16d8e4951dd15dd0369492`。初始复现基线为 `da281de1d147fac2d9242b44384c5cde370c3600`。范围为 SQL 识别、手工事务提交、JDBC 批量超时与取消、Oracle/OB Oracle 脚本切分。未包含对象管理、LOB 或过程多结果功能。
## 与上游实现的比较
- #11849 已补齐复合触发器和存储例程声明区的嵌套子程序。采用上游限定在触发器头部的识别方式,删除本 PR 扫描整个 token 集合的重复分支;保留 Oracle/OB Oracle、独占行 `/`、后续 ALTER/CREATE 的回归用例。原复合触发器用例已在当前 main 上通过。
- #11951 已在公共 Oracle 方言配置中关闭反斜杠转义,删除本 PR 的 OB 重复配置。光标选句使用的 SqlScanner 仍需读取该配置,不能只保留配置而撤销入口修复;新增 Oracle 系空行选句用例,当前 main 失败、候选通过。
- 当前 main 仍无法通过 nq、跨块 q/nq、OB 专属引号、DBLink、重复提交、Rust RPC 超时、JDBC 批量取消/回滚以及 Transwarp 取消回归。保留这些增量。
- main 新增关闭后重开标签的能力。补测同一标签 ID 对应新对象时,旧提交成功或失败均不清除新事务,且新事务仍独立去重。复用现有 Store 流程,不增加新的交互。
## 1. DBLink 被误识别为参数
Oracle 和 OceanBase Oracle 的以下脚本中,参数识别把 `LINK` 加入参数列表;同名 `@set` 还可能替换对象链接。
```sql
SELECT seq.NEXTVAL@LINK, @value FROM dual;
CREATE SYNONYM local_t FOR app.remote_t@LINK;
MERGE INTO local_t t USING app.remote_t@LINK s
ON (t.id=s.id) WHEN MATCHED THEN UPDATE SET t.v=@value;
```
复现:将单条脚本交给 `extractSqlParameters`,数据库类型选择 `oracle` 或 `oceanbase-oracle`。SELECT、MERGE 预期只返回 `value`,同义词语句预期无参数;基线实际多返回 `LINK`。在脚本前声明 `@set LINK = wrong;`,检查展开结果仍保留链接。MERGE 同时覆盖匿名块中的第一条语句。
验收:远程序列、同义词目标、MERGE 来源保留 `@LINK`;真正参数正常识别和替换;已有普通 FROM、表达式和其他方言用例通过。复用共享 `sqlDatabaseLinks.ts`,同时覆盖参数替换和变量展开。
## 2. 同一事务重复提交
复现:建立 Oracle/OB Oracle 手工事务,使后端提交 Promise 暂不完成,连续调用真实 store 的 `commitTransaction` 两次。基线调用后端两次,第二次可能报事务不存在。
验收:同一标签、同一事务只发出一次后端请求,所有调用者收到相同成功或失败结果;切换连接、库、schema、关闭标签或关闭后重开同 ID 标签后的迟到结果不清除新事务;同步异常后仍可提交下一事务。实际工具栏连点验收单独记录。
## 3. 批量失败、超时或取消后继续执行
复现:通过 `JsonRpcServer.dispatchForRuntime("execute_batch", ...)` 传入两条写入和 `timeoutSecs`。JDBC 边界记录 Statement 超时与取消;第一条阻塞时取消,再放行第一条。基线不设置指定超时、未追踪活动 Statement,fallback 可能继续第二条。OceanBase Connector/J 还允许 `continueBatchOnError=true` 导致失败后继续。
此外,Rust `AgentDriverClient` 的四个批量/事务入口未把超时放入 RPC 参数。使用真实客户端连接返回收到参数的测试 Agent,指定 7 秒后,基线收到的 `timeoutSecs` 为 NULL;修复后为 7,后续无超时请求为 0。
实库复现步骤:在隔离表准备两行,连接 A 锁定第一行;连接 B 执行“更新第一行、更新第二行”的批量并设置短超时。等待失败后用独立连接读回第二行,再释放 A 的锁。另用无效第一条检查错误即停,用取消检查迟到成功不会触发后续项。结束后删除本次创建的隔离对象并读回确认。
验收:RPC 超时传到 Statement;准备和执行阶段的取消被保留;fallback 和事务执行不启动下一条;下一次 RPC 不继承上一次超时/取消;事务回滚未确认时,池连接先驱逐租约再关闭,非池连接直接关闭,不通过恢复 auto-commit 意外提交;OceanBase 连接属性强制失败即停,覆盖 URL 中相反配置。普通批量此前已成功的语句不承诺自动回滚。
连接池回归使用真实 RPC、Hikari 和隔离 H2 表,只在 JDBC 回滚边界注入 `HY000`、`HYT00`、`08006` 异常。第一会话插入后失败,在外层请求清理前让第二会话查询;修复前能读到未提交的一行,修复后读到零行并使用新的物理连接。此用例覆盖 `close()` 归还池连接与外层驱逐之间的窗口。
共享 RPC 的 Transwarp 自定义事务另有回归:通过真实 `execute_transaction` dispatch 检查指定超时,在 `BEGIN TRANSACTION` 或首条写入阻塞期间取消,迟到成功后不得继续写入或 COMMIT,应尝试 ROLLBACK。使用 JDBC 边界夹具,修复前超时与取消两项失败,修复后通过,下一请求不继承取消状态。Transwarp 实库尚未验收。
## 4. Oracle/OB Oracle 脚本切分错误
复现:将以下文本交给 `split_sql_statements_for_database`,基线会把字符串里的 `DELETE` 当成独立语句。
```sql
UPDATE events SET name = nq'{it's; DELETE FROM events; --}';
SELECT 1 FROM dual;
```
复合触发器历史复现:一个 `COMPOUND TRIGGER` 包含 `BEFORE STATEMENT` 和 `AFTER EACH ROW` 两个区段,随后以独占一行的 `/` 结束;初始基线把区段分为多条,当前 main 已解决。流式导入将 `q`/`nq` 前缀与引号跨 chunk 切开的缺陷在当前 main 仍可复现。
验收:引号内分号不拆句,复合触发器整体发送;流式导入与光标选句一致;普通多语句正常切分;OB Oracle 不采用 MySQL 反引号和 DELIMITER;其他方言及上游 XML 实体回归保持通过。解析测试不证明目标数据库支持所有触发器语法。
## 自动化命令与证据
```sh
pnpm exec vitest run apps/desktop/src/lib/__tests__/sql/sqlDatabaseLinks.spec.ts apps/desktop/src/lib/__tests__/sql/sqlParameters.spec.ts apps/desktop/src/lib/__tests__/sql/sqlVariables.spec.ts apps/desktop/src/stores/__tests__/queryStore.manualTransactionCommit.spec.ts
cargo test -p dbx-sql-core --lib --locked
cargo test -p dbx-driver-agent --lib --locked batch_and_transaction_send_statement_timeout_to_agent
cd agents
./gradlew :common:test :oceanbase-oracle:test :oceanbase-oracle:shadowJar
```
当前 main 的行为复现失败包括 DBLink 10 项、事务并发 6 项、JDBC 批量生命周期 8 项及回滚连接池 1 项、SQL 切分 6 项和另测的光标反斜杠 1 项、Rust 四入口超时传递 1 项、Transwarp 2 项。测试调用生产入口;基线工作树仅加入回归测试及其测试依赖。修复后的定向测试通过,完整验证和实库、GUI、CI 状态以 PR 最新 SHA 的记录为准,未执行不得记为通过。
2026-10-13 使用本次独立构建的 OceanBase Agent,在 OceanBase Oracle 4.2.5.7 / Connector/J 2.4.19 上验证批量与事务的超时、取消四项,以及首条 SQL 失败后不执行下一条。五项通过,独立连接读回两行均未修改;专属临时表已删除并读回不存在。JAR SHA-256 为 `ea92c7cf7c10f872ca25a5f7db19374e1424fa64241493401f702898b74a9f09`。未替换现有运行实例的 Agent。此验证不替代原生 Oracle、Transwarp 实库及桌面 GUI 验收。
## 交付验收清单
- [ ] 本文四类行为回归通过,Core 到 Agent 的超时传递完整。
- [ ] 格式、lint、类型、相关 Rust/Java 测试及生成产物检查通过;基线失败单列证据。
- [ ] OceanBase Oracle 实库完成批量/事务超时及取消,读回后续行与清理结果。
- [ ] 产品工具栏连点、SQL 编辑器参数提示及执行反馈完成 GUI 验收。
- [ ] 原生 Oracle 相关入口完成实库兼容验收,或明确保留具体缺口。
- [ ] 修复 PR 对应的最新 SHA 必需 CI 通过;代码、实库、GUI 和部署状态分别记录。