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

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

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

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

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

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