Refactor ShardingSpherePreparedStatement for support multi executionC…

…ontext.
apache · Oct 19, 2023 · f50b9ce · f50b9ce
1 parent 43bba09
commit f50b9ce
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 11 deletions.
diff --git a/...src/main/java/org/apache/shardingsphere/driver/jdbc/adapter/AbstractStatementAdapter.java b/...src/main/java/org/apache/shardingsphere/driver/jdbc/adapter/AbstractStatementAdapter.java
@@ -68,12 +68,17 @@ protected final boolean isNeedImplicitCommitTransaction(final ShardingSphereConn
         }
         if (1 == executionContexts.size()) {
             SQLStatement sqlStatement = executionContexts.iterator().next().getSqlStatementContext().getSqlStatement();
-            return isModifiedSQL(sqlStatement) && executionContexts.iterator().next().getExecutionUnits().size() > 1;
+            return isWriteDMLStatement(sqlStatement) && executionContexts.iterator().next().getExecutionUnits().size() > 1;
         }
-        return executionContexts.stream().anyMatch(each -> isModifiedSQL(each.getSqlStatementContext().getSqlStatement()));
+        for (ExecutionContext each : executionContexts) {
+            if (isWriteDMLStatement(each.getSqlStatementContext().getSqlStatement())) {
+                return true;
+            }
+        }
+        return false;
     }
 
-    private boolean isModifiedSQL(final SQLStatement sqlStatement) {
+    private boolean isWriteDMLStatement(final SQLStatement sqlStatement) {
         return sqlStatement instanceof DMLStatement && !(sqlStatement instanceof SelectStatement);
     }
 

diff --git a/...in/java/org/apache/shardingsphere/driver/jdbc/core/statement/ShardingSphereStatement.java b/...in/java/org/apache/shardingsphere/driver/jdbc/core/statement/ShardingSphereStatement.java
@@ -17,6 +17,7 @@
 
 package org.apache.shardingsphere.driver.jdbc.core.statement;
 
+import com.google.common.base.Preconditions;
 import com.google.common.base.Strings;
 import lombok.AccessLevel;
 import lombok.Getter;
@@ -195,7 +196,9 @@ private ShardingSphereResultSet doExecuteQuery(final Collection<ExecutionContext
             MergedResult mergedResult = mergeQuery(queryResults, each.getSqlStatementContext());
             boolean selectContainsEnhancedTable =
                     each.getSqlStatementContext() instanceof SelectStatementContext && ((SelectStatementContext) each.getSqlStatementContext()).isContainsEnhancedTable();
-            result = new ShardingSphereResultSet(getResultSets(), mergedResult, this, selectContainsEnhancedTable, each);
+            if (null == result) {
+                result = new ShardingSphereResultSet(getResultSets(), mergedResult, this, selectContainsEnhancedTable, each);
+            }
         }
         return result;
     }
@@ -358,13 +361,17 @@ private int executeUpdateWithImplicitCommitTransaction(final ExecuteUpdateCallba
 
     private int useDriverToExecuteUpdate(final ExecuteUpdateCallback updateCallback, final SQLStatementContext sqlStatementContext,
                                          final Collection<ExecutionContext> executionContexts) throws SQLException {
-        int result = 0;
+        Integer result = null;
+        Preconditions.checkArgument(!executionContexts.isEmpty());
         for (ExecutionContext each : executionContexts) {
             ExecutionGroupContext<JDBCExecutionUnit> executionGroupContext = createExecutionGroupContext(each);
             cacheStatements(executionGroupContext.getInputGroups());
             JDBCExecutorCallback<Integer> callback = createExecuteUpdateCallback(updateCallback, sqlStatementContext);
-            result = executor.getRegularExecutor().executeUpdate(executionGroupContext,
+            int effectedCount = executor.getRegularExecutor().executeUpdate(executionGroupContext,
                     each.getQueryContext(), each.getRouteContext().getRouteUnits(), callback);
+            if (null == result) {
+                result = effectedCount;
+            }
         }
         return result;
     }
@@ -567,13 +574,17 @@ private boolean executeWithImplicitCommitTransaction(final ExecuteCallback callb
     }
 
     private boolean useDriverToExecute(final ExecuteCallback callback, final Collection<ExecutionContext> executionContexts) throws SQLException {
-        boolean result = false;
+        Boolean result = null;
+        Preconditions.checkArgument(!executionContexts.isEmpty());
         for (ExecutionContext each : executionContexts) {
             ExecutionGroupContext<JDBCExecutionUnit> executionGroupContext = createExecutionGroupContext(each);
             cacheStatements(executionGroupContext.getInputGroups());
             JDBCExecutorCallback<Boolean> jdbcExecutorCallback = createExecuteCallback(callback, each.getSqlStatementContext().getSqlStatement());
-            result = executor.getRegularExecutor().execute(executionGroupContext,
+            boolean isWrite = executor.getRegularExecutor().execute(executionGroupContext,
                     each.getQueryContext(), each.getRouteContext().getRouteUnits(), jdbcExecutorCallback);
+            if (null == result) {
+                result = isWrite;
+            }
         }
         return result;
     }

diff --git a/...re/src/main/java/org/apache/shardingsphere/proxy/backend/connector/DatabaseConnector.java b/...re/src/main/java/org/apache/shardingsphere/proxy/backend/connector/DatabaseConnector.java
@@ -18,7 +18,6 @@
 package org.apache.shardingsphere.proxy.backend.connector;
 
 import com.google.common.base.Preconditions;
-import org.apache.shardingsphere.infra.exception.dialect.SQLExceptionTransformEngine;
 import org.apache.shardingsphere.infra.binder.context.aware.CursorDefinitionAware;
 import org.apache.shardingsphere.infra.binder.context.segment.insert.keygen.GeneratedKeyContext;
 import org.apache.shardingsphere.infra.binder.context.statement.SQLStatementContext;
@@ -31,6 +30,8 @@
 import org.apache.shardingsphere.infra.connection.kernel.KernelProcessor;
 import org.apache.shardingsphere.infra.connection.refresher.MetaDataRefreshEngine;
 import org.apache.shardingsphere.infra.database.core.type.DatabaseType;
+import org.apache.shardingsphere.infra.exception.core.ShardingSpherePreconditions;
+import org.apache.shardingsphere.infra.exception.dialect.SQLExceptionTransformEngine;
 import org.apache.shardingsphere.infra.executor.sql.context.ExecutionContext;
 import org.apache.shardingsphere.infra.executor.sql.execute.engine.SQLExecutorExceptionHandler;
 import org.apache.shardingsphere.infra.executor.sql.execute.engine.driver.jdbc.JDBCExecutionUnit;
@@ -46,7 +47,6 @@
 import org.apache.shardingsphere.infra.metadata.database.schema.util.SystemSchemaUtils;
 import org.apache.shardingsphere.infra.rule.identifier.type.DataNodeContainedRule;
 import org.apache.shardingsphere.infra.session.query.QueryContext;
-import org.apache.shardingsphere.infra.exception.core.ShardingSpherePreconditions;
 import org.apache.shardingsphere.mode.manager.ContextManager;
 import org.apache.shardingsphere.mode.metadata.MetaDataContexts;
 import org.apache.shardingsphere.proxy.backend.connector.jdbc.executor.callback.ProxyJDBCExecutorCallback;
@@ -191,7 +191,12 @@ private boolean isNeedImplicitCommitTransaction(final Collection<ExecutionContex
             SQLStatement sqlStatement = executionContexts.iterator().next().getSqlStatementContext().getSqlStatement();
             return isWriteDMLStatement(sqlStatement) && executionContexts.iterator().next().getExecutionUnits().size() > 1;
         }
-        return executionContexts.stream().anyMatch(each -> isWriteDMLStatement(each.getSqlStatementContext().getSqlStatement()));
+        for (ExecutionContext each : executionContexts) {
+            if (isWriteDMLStatement(each.getSqlStatementContext().getSqlStatement())) {
+                return true;
+            }
+        }
+        return false;
     }
 
     private boolean isWriteDMLStatement(final SQLStatement sqlStatement) {