github/dbx
4
0 Fork 0
mirror of https://github.com/t8y2/dbx.git synced 2026-11-15 06:51:27 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(oceanbase): support manual transaction control in oracle mode

Browse Source
This commit is contained in:
zipg authored and GitHub committed 2026-09-15 15:37:36 +08:00
1 parent 39042f7917
commit 2b1c5fa71e
15 files changed
+391 -52
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

agents/common/src/main/java/com/dbx/agent/AbstractJdbcAgent.java
+22 -9
View File
@@ -255,13 +255,19 @@ public abstract class AbstractJdbcAgent extends BaseDatabaseAgent {
@Override
public QueryResult executeTransaction(List<String> statements, String schema) {
return TransactionExecutor.executeUpdateStatements(
requireConnected(),
statements,
schema,
this::setSchemaSQL,
this::resetSchemaSQL
);
Connection conn = requireConnected();
return unchecked(() -> {
if (!conn.getAutoCommit()) {
throw new IllegalStateException("Cannot start a one-shot transaction while a manual transaction is open");
}
return TransactionExecutor.executeUpdateStatements(
conn,
statements,
schema,
this::setSchemaSQL,
this::resetSchemaSQL
);
});
}
@Override
@@ -377,7 +383,9 @@ public abstract class AbstractJdbcAgent extends BaseDatabaseAgent {
JdbcExecutor executor,
boolean succeeded,
boolean requiresSessionAffinity,
boolean evictAfterRequest
boolean evictAfterRequest,
boolean endsSessionAffinity,
boolean preservesSchemaContext
) {
if (poolRegistry == null) {
return;
@@ -389,7 +397,12 @@ public abstract class AbstractJdbcAgent extends BaseDatabaseAgent {
}
if (succeeded && requiresSessionAffinity) {
sessionAffinity = true;
JdbcSchemaSwitcher.forget(connection);
if (!preservesSchemaContext) {
JdbcSchemaSwitcher.forget(connection);
}
}
if (succeeded && endsSessionAffinity) {
sessionAffinity = false;
}
if (pooledLease == null) {
connection = null;
agents/common/src/main/java/com/dbx/agent/DatabaseAgent.java
+65 -7
View File
@@ -224,13 +224,71 @@ public interface DatabaseAgent {
if (conn == null) {
throw new IllegalStateException("Not connected");
}
return TransactionExecutor.executeUpdateStatements(
conn,
statements,
schema,
this::setSchemaSQL,
this::resetSchemaSQL
);
return unchecked(() -> {
if (!conn.getAutoCommit()) {
throw new IllegalStateException("Cannot start a one-shot transaction while a manual transaction is open");
}
return TransactionExecutor.executeUpdateStatements(
conn,
statements,
schema,
this::setSchemaSQL,
this::resetSchemaSQL
);
});
}
/** Starts an interactive transaction on the current JDBC connection. */
default Map<String, Object> beginManualTransaction(String schema) {
Connection conn = getConnection();
if (conn == null) {
throw new IllegalStateException("Not connected");
}
return unchecked(() -> {
if (!conn.getAutoCommit()) {
throw new IllegalStateException("Manual transaction already open");
}
if (schema != null && !schema.trim().isEmpty()) {
JdbcSchemaSwitcher.apply(conn, schema, this::setSchemaSQL, this::resetSchemaSQL);
}
conn.setAutoCommit(false);
JdbcSchemaSwitcher.preserve(conn);
return Collections.singletonMap("ok", (Object) true);
});
}
/** Commits an interactive transaction on the current JDBC connection. */
default Map<String, Object> commitManualTransaction() {
Connection conn = getConnection();
if (conn == null) {
throw new IllegalStateException("Not connected");
}
return unchecked(() -> {
if (conn.getAutoCommit()) {
throw new IllegalStateException("No manual transaction open");
}
conn.commit();
conn.setAutoCommit(true);
JdbcSchemaSwitcher.releasePreserved(conn);
return Collections.singletonMap("ok", (Object) true);
});
}
/** Rolls back an interactive transaction on the current JDBC connection. */
default Map<String, Object> rollbackManualTransaction() {
Connection conn = getConnection();
if (conn == null) {
throw new IllegalStateException("Not connected");
}
return unchecked(() -> {
if (conn.getAutoCommit()) {
throw new IllegalStateException("No manual transaction open");
}
conn.rollback();
conn.setAutoCommit(true);
JdbcSchemaSwitcher.releasePreserved(conn);
return Collections.singletonMap("ok", (Object) true);
});
}
default QueryResult executeBatch(List<String> statements, String schema) {
agents/common/src/main/java/com/dbx/agent/JdbcSchemaSwitcher.java
+18
View File
@@ -12,6 +12,8 @@ import java.util.function.Supplier;
final class JdbcSchemaSwitcher {
private static final Map<Connection, ResetState> RESET_STATE_BY_CONNECTION =
Collections.synchronizedMap(new WeakHashMap<>());
private static final Map<Connection, Boolean> PRESERVED_BY_CONNECTION =
Collections.synchronizedMap(new WeakHashMap<>());
private JdbcSchemaSwitcher() {
}
@@ -27,6 +29,9 @@ final class JdbcSchemaSwitcher {
Supplier<String> resetSchemaSql
) throws Exception {
if (schema == null || schema.trim().isEmpty()) {
if (isPreserved(conn)) {
return;
}
if (!restore(conn, true)) {
throw new SQLException("Failed to restore the original JDBC schema context");
}
@@ -70,7 +75,20 @@ final class JdbcSchemaSwitcher {
RESET_STATE_BY_CONNECTION.remove(conn);
}
static void preserve(Connection conn) {
PRESERVED_BY_CONNECTION.put(conn, Boolean.TRUE);
}
static void releasePreserved(Connection conn) {
PRESERVED_BY_CONNECTION.remove(conn);
}
private static boolean isPreserved(Connection conn) {
return PRESERVED_BY_CONNECTION.containsKey(conn);
}
private static boolean restore(Connection conn, boolean preserveOnFailure) {
PRESERVED_BY_CONNECTION.remove(conn);
ResetState resetState = RESET_STATE_BY_CONNECTION.remove(conn);
if (resetState == null) {
return true;
agents/common/src/main/java/com/dbx/agent/JsonRpcServer.java
+32 -1
View File
@@ -111,7 +111,9 @@ public final class JsonRpcServer {
jdbcExecutor,
succeeded,
requiresSessionAffinity(method, params),
evictAfterRequest(method)
evictAfterRequest(method),
endsSessionAffinity(method),
preservesSchemaContext()
);
}
}
@@ -315,6 +317,15 @@ public final class JsonRpcServer {
List<String> statements = gson.fromJson(params.get("statements"), statementsType);
return agent.executeTransaction(statements, stringOrNull(params, "schema"));
}
if (AgentProtocol.METHOD_BEGIN_MANUAL_TRANSACTION.equals(method)) {
return agent.beginManualTransaction(stringOrNull(params, "schema"));
}
if (AgentProtocol.METHOD_COMMIT_MANUAL_TRANSACTION.equals(method)) {
return agent.commitManualTransaction();
}
if (AgentProtocol.METHOD_ROLLBACK_MANUAL_TRANSACTION.equals(method)) {
return agent.rollbackManualTransaction();
}
if (AgentProtocol.METHOD_EXECUTE_BATCH.equals(method)) {
Type statementsType = new TypeToken<List<String>>() {}.getType();
List<String> statements = gson.fromJson(params.get("statements"), statementsType);
@@ -422,6 +433,9 @@ public final class JsonRpcServer {
}
private boolean requiresSessionAffinity(String method, JsonObject params) {
if (AgentProtocol.METHOD_BEGIN_MANUAL_TRANSACTION.equals(method)) {
return true;
}
try {
if (AgentProtocol.METHOD_EXECUTE_QUERY.equals(method)
|| AgentProtocol.METHOD_EXECUTE_QUERY_PAGE.equals(method)
@@ -447,6 +461,23 @@ public final class JsonRpcServer {
return false;
}
private static boolean endsSessionAffinity(String method) {
return AgentProtocol.METHOD_COMMIT_MANUAL_TRANSACTION.equals(method)
|| AgentProtocol.METHOD_ROLLBACK_MANUAL_TRANSACTION.equals(method);
}
private boolean preservesSchemaContext() {
Connection connection = agent.getConnection();
if (connection == null) {
return false;
}
try {
return !connection.getAutoCommit();
} catch (Exception ignored) {
return false;
}
}
private static boolean evictAfterRequest(String method) {
return AgentProtocol.METHOD_GET_EXPLAIN_INFO.equals(method);
}
agents/common/src/test/java/com/dbx/agent/AbstractJdbcAgentTest.java
+16
View File
@@ -300,6 +300,22 @@ class AbstractJdbcAgentTest {
);
}
@Test
void rejectsOneShotTransactionWhenManualTransactionIsOpen() {
TrackingConnection tracking = new TrackingConnection();
TestAgent agent = new TestAgent(tracking);
agent.connect(new ConnectParams());
agent.beginManualTransaction(null);
IllegalStateException error = assertThrows(
IllegalStateException.class,
() -> agent.executeTransaction(Collections.singletonList("UPDATE A SET ID = 1"), null)
);
assertEquals("Cannot start a one-shot transaction while a manual transaction is open", error.getMessage());
assertEquals(Collections.singletonList("setAutoCommit:false"), tracking.calls);
}
@Test
void preservesPlSqlBlockTerminatorDuringTransactionExecution() {
TrackingConnection tracking = new TrackingConnection();
agents/common/src/test/java/com/dbx/agent/CommonJavaCompatibilityTest.java
+36
View File
@@ -150,6 +150,42 @@ class CommonJavaCompatibilityTest {
assertEquals("Example JDBC", result.getAsJsonObject("databaseInfo").get("driverName").getAsString());
}
@Test
void jsonRpcServerDispatchesInteractiveTransactionMethods() {
List<String> calls = new ArrayList<>();
MinimalAgent agent = new MinimalAgent() {
@Override
public Map<String, Object> beginManualTransaction(String schema) {
calls.add("begin:" + schema);
return Collections.singletonMap("ok", (Object) true);
}
@Override
public Map<String, Object> commitManualTransaction() {
calls.add("commit");
return Collections.singletonMap("ok", (Object) true);
}
@Override
public Map<String, Object> rollbackManualTransaction() {
calls.add("rollback");
return Collections.singletonMap("ok", (Object) true);
}
};
JsonRpcServer server = new JsonRpcServer(agent);
assertTrue(JsonParser.parseString(server.handleRequest(
"{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"begin_manual_transaction\",\"params\":{\"schema\":\"APP\"}}"
)).getAsJsonObject().has("result"));
assertTrue(JsonParser.parseString(server.handleRequest(
"{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"commit_manual_transaction\",\"params\":{}}"
)).getAsJsonObject().has("result"));
assertTrue(JsonParser.parseString(server.handleRequest(
"{\"jsonrpc\":\"2.0\",\"id\":3,\"method\":\"rollback_manual_transaction\",\"params\":{}}"
)).getAsJsonObject().has("result"));
assertEquals(List.of("begin:APP", "commit", "rollback"), calls);
}
@Test
void multiSessionServerCreatesAndClosesIndependentAgents() throws Exception {
java.util.List<TrackingAgent> created = new java.util.ArrayList<>();
agents/common/src/test/java/com/dbx/agent/JdbcConnectionPoolingTest.java
+59
View File
@@ -1869,6 +1869,65 @@ class JdbcConnectionPoolingTest {
}
}
@Test
void interactiveTransactionPinsConnectionUntilCommit() {
AtomicInteger physicalOpens = new AtomicInteger();
AtomicInteger requestIds = new AtomicInteger();
String url = h2Url("manual_transaction");
try (MultiSessionJsonRpcServer server = server(url, physicalOpens, 2)) {
openSession(server, requestIds, "transaction");
query(server, requestIds, "transaction", "CREATE TABLE ITEMS (ID INT)", null);
result(request(
server,
requestIds,
AgentProtocol.METHOD_BEGIN_MANUAL_TRANSACTION,
sessionParams("transaction")
));
query(server, requestIds, "transaction", "INSERT INTO ITEMS VALUES (1)", null);
openSession(server, requestIds, "observer");
JsonObject beforeCommit = query(server, requestIds, "observer", "SELECT COUNT(*) FROM ITEMS", null);
assertEquals(0, beforeCommit.getAsJsonArray("rows").get(0).getAsJsonArray().get(0).getAsInt());
result(request(
server,
requestIds,
AgentProtocol.METHOD_COMMIT_MANUAL_TRANSACTION,
sessionParams("transaction")
));
JsonObject afterCommit = query(server, requestIds, "observer", "SELECT COUNT(*) FROM ITEMS", null);
assertEquals(1, afterCommit.getAsJsonArray("rows").get(0).getAsJsonArray().get(0).getAsInt());
}
}
@Test
void interactiveTransactionRestoresSchemaBeforeReturningConnection() {
AtomicInteger physicalOpens = new AtomicInteger();
AtomicInteger requestIds = new AtomicInteger();
String url = h2Url("manual_transaction_schema");
try (MultiSessionJsonRpcServer server = server(url, physicalOpens, 2)) {
openSession(server, requestIds, "transaction-schema");
query(server, requestIds, "transaction-schema", "CREATE SCHEMA IF NOT EXISTS APP", null);
JsonObject beginParams = sessionParams("transaction-schema");
beginParams.addProperty("schema", "APP");
result(request(server, requestIds, AgentProtocol.METHOD_BEGIN_MANUAL_TRANSACTION, beginParams));
JsonObject inTransaction = query(server, requestIds, "transaction-schema", "SELECT CURRENT_SCHEMA", null);
assertEquals("APP", inTransaction.getAsJsonArray("rows").get(0).getAsJsonArray().get(0).getAsString());
result(request(
server,
requestIds,
AgentProtocol.METHOD_COMMIT_MANUAL_TRANSACTION,
sessionParams("transaction-schema")
));
openSession(server, requestIds, "schema-observer");
JsonObject afterCommit = query(server, requestIds, "schema-observer", "SELECT CURRENT_SCHEMA", null);
assertEquals("PUBLIC", afterCommit.getAsJsonArray("rows").get(0).getAsJsonArray().get(0).getAsString());
}
}
@Test
void affinityDetectionCoversCommonSessionStateStatements() {
assertTrue(JdbcConnectionAffinity.requiresSessionAffinity("SET search_path TO app"));