Move ShardingSpherePreparedStatement.executionContext into DriverExec…

…uteBatchExecutor (#31595)
apache · Jun 5, 2024 · 156475b · 156475b
1 parent 5cc7b71
commit 156475b
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 10 deletions.
diff --git a/jdbc/src/main/java/org/apache/shardingsphere/driver/executor/DriverExecuteBatchExecutor.java b/jdbc/src/main/java/org/apache/shardingsphere/driver/executor/DriverExecuteBatchExecutor.java
@@ -65,6 +65,8 @@ public final class DriverExecuteBatchExecutor {
     @Getter
     private final BatchPreparedStatementExecutor batchPreparedStatementExecutor;
 
+    private ExecutionContext executionContext;
+
     public DriverExecuteBatchExecutor(final ShardingSphereConnection connection, final ShardingSphereMetaData metaData, final ShardingSphereDatabase database, final JDBCExecutor jdbcExecutor) {
         this.connection = connection;
         this.metaData = metaData;
@@ -76,13 +78,12 @@ public DriverExecuteBatchExecutor(final ShardingSphereConnection connection, fin
      * 
      * @param queryContext query context
      * @param database database
-     * @return execution context
      */
-    public ExecutionContext addBatch(final QueryContext queryContext, final ShardingSphereDatabase database) {
+    public void addBatch(final QueryContext queryContext, final ShardingSphereDatabase database) {
         Optional<String> trafficInstanceId = connection.getTrafficInstanceId(metaData.getGlobalRuleMetaData().getSingleRule(TrafficRule.class), queryContext);
         ExecutionContext result = trafficInstanceId.map(optional -> createExecutionContext(queryContext, optional)).orElseGet(() -> createExecutionContext(queryContext, database));
         batchPreparedStatementExecutor.addBatchForExecutionUnits(result.getExecutionUnits());
-        return result;
+        executionContext = result;
     }
 
     private ExecutionContext createExecutionContext(final QueryContext queryContext, final String trafficInstanceId) {
@@ -103,7 +104,6 @@ private ExecutionContext createExecutionContext(final QueryContext queryContext,
      * @param generatedValues generated values
      * @param statementOption statement option
      * @param prepareEngine prepare engine
-     * @param executionContext execution context
      * @param addCallback statement add callback
      * @param replayCallback prepared statement parameters replay callback
      * @param generatedKeyCallback generated key callback
@@ -113,7 +113,7 @@ private ExecutionContext createExecutionContext(final QueryContext queryContext,
     @SuppressWarnings("rawtypes")
     public int[] executeBatch(final ShardingSphereDatabase database, final SQLStatementContext sqlStatementContext, final Collection<Comparable<?>> generatedValues,
                               final StatementOption statementOption, final DriverExecutionPrepareEngine<JDBCExecutionUnit, Connection> prepareEngine,
-                              final ExecutionContext executionContext, final StatementAddCallback addCallback, final PreparedStatementParametersReplayCallback replayCallback,
+                              final StatementAddCallback addCallback, final PreparedStatementParametersReplayCallback replayCallback,
                               final GeneratedKeyCallback generatedKeyCallback) throws SQLException {
         if (null == executionContext) {
             return new int[0];

diff --git a/...org/apache/shardingsphere/driver/jdbc/core/statement/ShardingSpherePreparedStatement.java b/...org/apache/shardingsphere/driver/jdbc/core/statement/ShardingSpherePreparedStatement.java
@@ -40,7 +40,6 @@
 import org.apache.shardingsphere.infra.exception.core.ShardingSpherePreconditions;
 import org.apache.shardingsphere.infra.exception.dialect.SQLExceptionTransformEngine;
 import org.apache.shardingsphere.infra.exception.kernel.syntax.EmptySQLException;
-import org.apache.shardingsphere.infra.executor.sql.context.ExecutionContext;
 import org.apache.shardingsphere.infra.executor.sql.execute.engine.driver.jdbc.JDBCExecutionUnit;
 import org.apache.shardingsphere.infra.executor.sql.execute.result.query.QueryResult;
 import org.apache.shardingsphere.infra.executor.sql.execute.result.query.impl.driver.jdbc.type.stream.JDBCStreamQueryResult;
@@ -116,8 +115,6 @@ public final class ShardingSpherePreparedStatement extends AbstractPreparedState
     @Getter
     private final boolean selectContainsEnhancedTable;
 
-    private ExecutionContext executionContext;
-
     private Map<String, Integer> columnLabelAndIndexMap;
 
     private ResultSet currentResultSet;
@@ -383,7 +380,7 @@ private String getGeneratedKeysColumnName(final String columnName) {
     public void addBatch() {
         currentResultSet = null;
         QueryContext queryContext = createQueryContext();
-        executionContext = driverExecutorFacade.getExecuteBatchExecutor().addBatch(queryContext, metaData.getDatabase(databaseName));
+        driverExecutorFacade.getExecuteBatchExecutor().addBatch(queryContext, metaData.getDatabase(databaseName));
         findGeneratedKey().ifPresent(optional -> generatedValues.addAll(optional.getGeneratedValues()));
         clearParameters();
     }
@@ -393,7 +390,7 @@ public int[] executeBatch() throws SQLException {
         ShardingSphereDatabase database = metaData.getDatabase(databaseName);
         try {
             return driverExecutorFacade.getExecuteBatchExecutor().executeBatch(database, sqlStatementContext, generatedValues, statementOption,
-                    createDriverExecutionPrepareEngine(database), executionContext, (StatementAddCallback<PreparedStatement>) (statements, parameterSets) -> this.statements.addAll(statements),
+                    createDriverExecutionPrepareEngine(database), (StatementAddCallback<PreparedStatement>) (statements, parameterSets) -> this.statements.addAll(statements),
                     this::replaySetParameter,
                     () -> {
                         currentBatchGeneratedKeysResultSet = getGeneratedKeys();