Skip to content

Commit

Permalink
Optimize ShardingSpherePreparedStatement and ShardingSphereStatement …
Browse files Browse the repository at this point in the history
…for multi execution (#28927)

* Optimize ShardingSpherePreparedStatement and ShardingSphereStatement for support multi executionContext.

* Optimize ShardingSpherePreparedStatement and ShardingSphereStatement for support multi executionContext.
  • Loading branch information
tuichenchuxin authored Nov 3, 2023
1 parent c7df22d commit 5f57f08
Show file tree
Hide file tree
Showing 2 changed files with 35 additions and 20 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -369,7 +369,7 @@ public int executeUpdate() throws SQLException {
}
return accumulate(results);
}
return isNeedImplicitCommitTransaction(connection, executionContexts) ? executeUpdateWithImplicitCommitTransaction() : useDriverToExecuteUpdate();
return executeUpdateWithExecutionContexts(executionContexts);
// CHECKSTYLE:OFF
} catch (final RuntimeException ex) {
// CHECKSTYLE:ON
Expand All @@ -380,7 +380,7 @@ public int executeUpdate() throws SQLException {
}
}

private int useDriverToExecuteUpdate() throws SQLException {
private int useDriverToExecuteUpdate(final Collection<ExecutionContext> executionContexts) throws SQLException {
Integer result = null;
Preconditions.checkArgument(!executionContexts.isEmpty());
// TODO support multi execution context, currently executionContexts.size() always equals 1
Expand Down Expand Up @@ -448,7 +448,7 @@ public boolean execute() throws SQLException {
}
return results.iterator().next() instanceof QueryResult;
}
return isNeedImplicitCommitTransaction(connection, executionContexts) ? executeWithImplicitCommitTransaction() : useDriverToExecute();
return executeWithExecutionContexts(executionContexts);
// CHECKSTYLE:OFF
} catch (final RuntimeException ex) {
// CHECKSTYLE:ON
Expand All @@ -474,11 +474,15 @@ private ExecutionGroupContext<RawSQLExecutionUnit> createRawExecutionGroupContex
.prepare(executionContext.getRouteContext(), executionContext.getExecutionUnits(), new ExecutionGroupReportContext(databaseName));
}

private boolean executeWithImplicitCommitTransaction() throws SQLException {
private boolean executeWithExecutionContexts(final Collection<ExecutionContext> executionContexts) throws SQLException {
return isNeedImplicitCommitTransaction(connection, executionContexts) ? executeWithImplicitCommitTransaction(executionContexts) : useDriverToExecute(executionContexts);
}

private boolean executeWithImplicitCommitTransaction(final Collection<ExecutionContext> executionContexts) throws SQLException {
boolean result;
try {
connection.setAutoCommit(false);
result = useDriverToExecute();
result = useDriverToExecute(executionContexts);
connection.commit();
// CHECKSTYLE:OFF
} catch (final Exception ex) {
Expand All @@ -491,11 +495,15 @@ private boolean executeWithImplicitCommitTransaction() throws SQLException {
return result;
}

private int executeUpdateWithImplicitCommitTransaction() throws SQLException {
private int executeUpdateWithExecutionContexts(final Collection<ExecutionContext> executionContexts) throws SQLException {
return isNeedImplicitCommitTransaction(connection, executionContexts) ? executeUpdateWithImplicitCommitTransaction(executionContexts) : useDriverToExecuteUpdate(executionContexts);
}

private int executeUpdateWithImplicitCommitTransaction(final Collection<ExecutionContext> executionContexts) throws SQLException {
int result;
try {
connection.setAutoCommit(false);
result = useDriverToExecuteUpdate();
result = useDriverToExecuteUpdate(executionContexts);
connection.commit();
// CHECKSTYLE:OFF
} catch (final RuntimeException ex) {
Expand All @@ -508,7 +516,7 @@ private int executeUpdateWithImplicitCommitTransaction() throws SQLException {
return result;
}

private boolean useDriverToExecute() throws SQLException {
private boolean useDriverToExecute(final Collection<ExecutionContext> executionContexts) throws SQLException {
Boolean result = null;
Preconditions.checkArgument(!executionContexts.isEmpty());
// TODO support multi execution context, currently executionContexts.size() always equals 1
Expand Down Expand Up @@ -554,18 +562,19 @@ public ResultSet getResultSet() throws SQLException {
if (useFederation) {
return executor.getSqlFederationEngine().getResultSet();
}
if (executionContexts.iterator().next().getSqlStatementContext() instanceof SelectStatementContext
|| executionContexts.iterator().next().getSqlStatementContext().getSqlStatement() instanceof DALStatement) {
ExecutionContext executionContext = executionContexts.iterator().next();
if (executionContext.getSqlStatementContext() instanceof SelectStatementContext
|| executionContext.getSqlStatementContext().getSqlStatement() instanceof DALStatement) {
List<ResultSet> resultSets = getResultSets();
if (resultSets.isEmpty()) {
return currentResultSet;
}
SQLStatementContext sqlStatementContext = executionContexts.iterator().next().getSqlStatementContext();
SQLStatementContext sqlStatementContext = executionContext.getSqlStatementContext();
MergedResult mergedResult = mergeQuery(getQueryResults(resultSets), sqlStatementContext);
if (null == columnLabelAndIndexMap) {
columnLabelAndIndexMap = ShardingSphereResultSetUtils.createColumnLabelAndIndexMap(sqlStatementContext, selectContainsEnhancedTable, resultSets.get(0).getMetaData());
}
currentResultSet = new ShardingSphereResultSet(resultSets, mergedResult, this, selectContainsEnhancedTable, executionContexts.iterator().next(), columnLabelAndIndexMap);
currentResultSet = new ShardingSphereResultSet(resultSets, mergedResult, this, selectContainsEnhancedTable, executionContext, columnLabelAndIndexMap);
}
return currentResultSet;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -477,8 +477,7 @@ private boolean execute0(final String sql, final ExecuteCallback executeCallback
}
return results.iterator().next() instanceof QueryResult;
}
return isNeedImplicitCommitTransaction(connection, executionContexts) ? executeWithImplicitCommitTransaction(executeCallback, executionContexts)
: useDriverToExecute(executeCallback, executionContexts);
return executeWithExecutionContexts(executeCallback, executionContexts);
} finally {
currentResultSet = null;
}
Expand Down Expand Up @@ -554,6 +553,11 @@ private ExecutionGroupContext<RawSQLExecutionUnit> createRawExecutionContext(fin
.prepare(executionContext.getRouteContext(), executionContext.getExecutionUnits(), new ExecutionGroupReportContext(databaseName));
}

private boolean executeWithExecutionContexts(final ExecuteCallback executeCallback, final Collection<ExecutionContext> executionContexts) throws SQLException {
return isNeedImplicitCommitTransaction(connection, executionContexts) ? executeWithImplicitCommitTransaction(executeCallback, executionContexts)
: useDriverToExecute(executeCallback, executionContexts);
}

private boolean executeWithImplicitCommitTransaction(final ExecuteCallback callback, final Collection<ExecutionContext> executionContexts) throws SQLException {
boolean result;
try {
Expand Down Expand Up @@ -626,16 +630,17 @@ public ResultSet getResultSet() throws SQLException {
if (useFederation) {
return executor.getSqlFederationEngine().getResultSet();
}
if (executionContexts.iterator().next().getSqlStatementContext() instanceof SelectStatementContext
|| executionContexts.iterator().next().getSqlStatementContext().getSqlStatement() instanceof DALStatement) {
ExecutionContext executionContext = executionContexts.iterator().next();
if (executionContext.getSqlStatementContext() instanceof SelectStatementContext
|| executionContext.getSqlStatementContext().getSqlStatement() instanceof DALStatement) {
List<ResultSet> resultSets = getResultSets();
if (resultSets.isEmpty()) {
return currentResultSet;
}
SQLStatementContext sqlStatementContext = executionContexts.iterator().next().getSqlStatementContext();
SQLStatementContext sqlStatementContext = executionContext.getSqlStatementContext();
MergedResult mergedResult = mergeQuery(getQueryResults(resultSets), sqlStatementContext);
boolean selectContainsEnhancedTable = sqlStatementContext instanceof SelectStatementContext && ((SelectStatementContext) sqlStatementContext).isContainsEnhancedTable();
currentResultSet = new ShardingSphereResultSet(resultSets, mergedResult, this, selectContainsEnhancedTable, executionContexts.iterator().next());
currentResultSet = new ShardingSphereResultSet(resultSets, mergedResult, this, selectContainsEnhancedTable, executionContext);
}
return currentResultSet;
}
Expand Down Expand Up @@ -712,8 +717,9 @@ public ResultSet getGeneratedKeys() throws SQLException {
}

private Optional<GeneratedKeyContext> findGeneratedKey() {
return executionContexts.iterator().next().getSqlStatementContext() instanceof InsertStatementContext
? ((InsertStatementContext) executionContexts.iterator().next().getSqlStatementContext()).getGeneratedKeyContext()
ExecutionContext executionContext = executionContexts.iterator().next();
return executionContext.getSqlStatementContext() instanceof InsertStatementContext
? ((InsertStatementContext) executionContext.getSqlStatementContext()).getGeneratedKeyContext()
: Optional.empty();
}

Expand Down

0 comments on commit 5f57f08

Please sign in to comment.