Refactor DriverExecutor

apache · May 27, 2024 · e08063b · e08063b
1 parent 1482d9d
commit e08063b
Showing 1 changed file with 8 additions and 6 deletions.
diff --git a/jdbc/src/main/java/org/apache/shardingsphere/driver/executor/DriverExecutor.java b/jdbc/src/main/java/org/apache/shardingsphere/driver/executor/DriverExecutor.java
@@ -93,10 +93,10 @@ public Optional<ResultSet> executeAdvanceQuery(final ShardingSphereMetaData meta
                                                    final DriverExecutionPrepareEngine<JDBCExecutionUnit, Connection> prepareEngine) throws SQLException {
         Optional<String> trafficInstanceId = connection.getTrafficInstanceId(metaData.getGlobalRuleMetaData().getSingleRule(TrafficRule.class), queryContext);
         if (trafficInstanceId.isPresent()) {
-            TrafficExecutorCallback<ResultSet> callback = JDBCDriverType.STATEMENT.equals(prepareEngine.getType())
+            TrafficExecutorCallback<ResultSet> trafficCallback = JDBCDriverType.STATEMENT.equals(prepareEngine.getType())
                     ? Statement::executeQuery
                     : ((statement, sql) -> ((PreparedStatement) statement).executeQuery());
-            return Optional.of(trafficExecutor.execute(connection.getProcessId(), database.getName(), trafficInstanceId.get(), queryContext, prepareEngine, callback));
+            return Optional.of(trafficExecutor.execute(connection.getProcessId(), database.getName(), trafficInstanceId.get(), queryContext, prepareEngine, trafficCallback));
         }
         if (sqlFederationEngine.decide(queryContext.getSqlStatementContext(), queryContext.getParameters(), database, metaData.getGlobalRuleMetaData())) {
             ExecuteQueryCallback sqlFederationCallback = JDBCDriverType.STATEMENT.equals(prepareEngine.getType())
@@ -171,10 +171,12 @@ public Optional<Boolean> executeAdvance(final ShardingSphereMetaData metaData, f
         }
         if (sqlFederationEngine.decide(queryContext.getSqlStatementContext(), queryContext.getParameters(), database, metaData.getGlobalRuleMetaData())) {
             executeType = ExecuteType.FEDERATION;
-            ResultSet resultSet = sqlFederationEngine.executeQuery(prepareEngine,
-                    new StatementExecuteQueryCallback(database.getProtocolType(), database.getResourceMetaData(),
-                            queryContext.getSqlStatementContext().getSqlStatement(), SQLExecutorExceptionHandler.isExceptionThrown()),
-                    new SQLFederationContext(false, queryContext, metaData, connection.getProcessId()));
+            ExecuteQueryCallback sqlFederationCallback = JDBCDriverType.STATEMENT.equals(prepareEngine.getType())
+                    ? new StatementExecuteQueryCallback(database.getProtocolType(), database.getResourceMetaData(),
+                            queryContext.getSqlStatementContext().getSqlStatement(), SQLExecutorExceptionHandler.isExceptionThrown())
+                    : new PreparedStatementExecuteQueryCallback(database.getProtocolType(), database.getResourceMetaData(),
+                            queryContext.getSqlStatementContext().getSqlStatement(), SQLExecutorExceptionHandler.isExceptionThrown());
+            ResultSet resultSet = sqlFederationEngine.executeQuery(prepareEngine, sqlFederationCallback, new SQLFederationContext(false, queryContext, metaData, connection.getProcessId()));
             return Optional.of(null != resultSet);
         }
         return Optional.empty();