This is an automated email from the ASF dual-hosted git repository.
zhonghongsheng pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/shardingsphere.git
The following commit(s) were added to refs/heads/master by this push:
new 0f04f334230 Support ConnectionContextAware for DistSQLQueryExecutor
(#37469)
0f04f334230 is described below
commit 0f04f334230847693d9c1b9bcb3aefda100d8319
Author: Hongsheng Zhong <[email protected]>
AuthorDate: Tue Dec 23 16:05:00 2025 +0800
Support ConnectionContextAware for DistSQLQueryExecutor (#37469)
* Add QueryContext param in DistSQLUpdateProxyBackendHandler constructor
* Refactor DistSQLUpdateExecuteEngine constructor, replace
ConnectionContext to DistSQLConnectionContext
* Refactor DistSQLExecutorAwareSetter set method, remove ConnectionContext
param
* Update unit test
---
.../handler/aware/DistSQLExecutorAwareSetter.java | 11 ++++-----
.../engine/query/DistSQLQueryExecuteEngine.java | 2 +-
.../engine/update/DistSQLUpdateExecuteEngine.java | 11 ++++-----
.../distsql/DistSQLProxyBackendHandlerFactory.java | 4 ++--
.../distsql/DistSQLUpdateProxyBackendHandler.java | 9 ++++++--
.../DistSQLProxyBackendHandlerFactoryTest.java | 26 +++++++++++-----------
.../DistSQLUpdateProxyBackendHandlerTest.java | 14 +++++++-----
.../DistSQLGlobalRuleDefinitionExecutorAssert.java | 4 ++--
8 files changed, 43 insertions(+), 38 deletions(-)
diff --git
a/infra/distsql-handler/src/main/java/org/apache/shardingsphere/distsql/handler/aware/DistSQLExecutorAwareSetter.java
b/infra/distsql-handler/src/main/java/org/apache/shardingsphere/distsql/handler/aware/DistSQLExecutorAwareSetter.java
index 328ba678cd5..2039bfb5a4a 100644
---
a/infra/distsql-handler/src/main/java/org/apache/shardingsphere/distsql/handler/aware/DistSQLExecutorAwareSetter.java
+++
b/infra/distsql-handler/src/main/java/org/apache/shardingsphere/distsql/handler/aware/DistSQLExecutorAwareSetter.java
@@ -28,7 +28,6 @@ import
org.apache.shardingsphere.infra.exception.generic.UnsupportedSQLOperation
import
org.apache.shardingsphere.infra.metadata.database.ShardingSphereDatabase;
import
org.apache.shardingsphere.infra.rewrite.sql.token.common.generator.aware.ConnectionContextAware;
import org.apache.shardingsphere.infra.rule.ShardingSphereRule;
-import org.apache.shardingsphere.infra.session.connection.ConnectionContext;
import org.apache.shardingsphere.mode.manager.ContextManager;
import java.util.Optional;
@@ -46,13 +45,11 @@ public final class DistSQLExecutorAwareSetter {
*
* @param contextManager context manager
* @param database database
- * @param distSQLConnectionContext DistSQL connection context
- * @param connectionContext connection context
+ * @param distsqlConnectionContext DistSQL connection context
* @param sqlStatement DistSQL statement
*/
@SuppressWarnings("rawtypes")
- public void set(final ContextManager contextManager, final
ShardingSphereDatabase database, final DistSQLConnectionContext
distSQLConnectionContext,
- final ConnectionContext connectionContext, final
DistSQLStatement sqlStatement) {
+ public void set(final ContextManager contextManager, final
ShardingSphereDatabase database, final DistSQLConnectionContext
distsqlConnectionContext, final DistSQLStatement sqlStatement) {
if (executor instanceof DistSQLExecutorDatabaseAware) {
ShardingSpherePreconditions.checkNotNull(database,
NoDatabaseSelectedException::new);
((DistSQLExecutorDatabaseAware) executor).setDatabase(database);
@@ -64,10 +61,10 @@ public final class DistSQLExecutorAwareSetter {
setRule((DistSQLExecutorRuleAware) executor, contextManager,
database);
}
if (executor instanceof DistSQLExecutorConnectionContextAware) {
- ((DistSQLExecutorConnectionContextAware)
executor).setConnectionContext(distSQLConnectionContext);
+ ((DistSQLExecutorConnectionContextAware)
executor).setConnectionContext(distsqlConnectionContext);
}
if (executor instanceof ConnectionContextAware) {
- ((ConnectionContextAware)
executor).setConnectionContext(connectionContext);
+ ((ConnectionContextAware)
executor).setConnectionContext(distsqlConnectionContext.getQueryContext().getConnectionContext());
}
}
diff --git
a/infra/distsql-handler/src/main/java/org/apache/shardingsphere/distsql/handler/engine/query/DistSQLQueryExecuteEngine.java
b/infra/distsql-handler/src/main/java/org/apache/shardingsphere/distsql/handler/engine/query/DistSQLQueryExecuteEngine.java
index e5ae6b46749..a4209e54d81 100644
---
a/infra/distsql-handler/src/main/java/org/apache/shardingsphere/distsql/handler/engine/query/DistSQLQueryExecuteEngine.java
+++
b/infra/distsql-handler/src/main/java/org/apache/shardingsphere/distsql/handler/engine/query/DistSQLQueryExecuteEngine.java
@@ -67,7 +67,7 @@ public final class DistSQLQueryExecuteEngine {
public void executeQuery() throws SQLException {
DistSQLQueryExecutor<DistSQLStatement> executor =
TypedSPILoader.getService(DistSQLQueryExecutor.class, sqlStatement.getClass());
try {
- new DistSQLExecutorAwareSetter(executor).set(contextManager, null
== databaseName ? null : contextManager.getDatabase(databaseName),
distsqlConnectionContext, null, sqlStatement);
+ new DistSQLExecutorAwareSetter(executor).set(contextManager, null
== databaseName ? null : contextManager.getDatabase(databaseName),
distsqlConnectionContext, sqlStatement);
} catch (final UnsupportedSQLOperationException ignored) {
columnNames = executor.getColumnNames(sqlStatement);
rows = Collections.emptyList();
diff --git
a/infra/distsql-handler/src/main/java/org/apache/shardingsphere/distsql/handler/engine/update/DistSQLUpdateExecuteEngine.java
b/infra/distsql-handler/src/main/java/org/apache/shardingsphere/distsql/handler/engine/update/DistSQLUpdateExecuteEngine.java
index 5bd4750a84f..e90e28f4b9e 100644
---
a/infra/distsql-handler/src/main/java/org/apache/shardingsphere/distsql/handler/engine/update/DistSQLUpdateExecuteEngine.java
+++
b/infra/distsql-handler/src/main/java/org/apache/shardingsphere/distsql/handler/engine/update/DistSQLUpdateExecuteEngine.java
@@ -18,6 +18,7 @@
package org.apache.shardingsphere.distsql.handler.engine.update;
import
org.apache.shardingsphere.distsql.handler.aware.DistSQLExecutorAwareSetter;
+import
org.apache.shardingsphere.distsql.handler.engine.DistSQLConnectionContext;
import
org.apache.shardingsphere.distsql.handler.engine.update.rdl.rule.engine.database.DatabaseRuleDefinitionExecuteEngine;
import
org.apache.shardingsphere.distsql.handler.engine.update.rdl.rule.engine.global.GlobalRuleDefinitionExecuteEngine;
import
org.apache.shardingsphere.distsql.handler.engine.update.rdl.rule.spi.database.DatabaseRuleDefinitionExecutor;
@@ -31,7 +32,6 @@ import
org.apache.shardingsphere.distsql.statement.type.rdl.rule.RuleDefinitionS
import
org.apache.shardingsphere.distsql.statement.type.rdl.rule.database.DatabaseRuleDefinitionStatement;
import
org.apache.shardingsphere.distsql.statement.type.rdl.rule.global.GlobalRuleDefinitionStatement;
import
org.apache.shardingsphere.infra.metadata.database.ShardingSphereDatabase;
-import org.apache.shardingsphere.infra.session.connection.ConnectionContext;
import org.apache.shardingsphere.infra.spi.type.typed.TypedSPILoader;
import org.apache.shardingsphere.mode.manager.ContextManager;
@@ -49,13 +49,14 @@ public final class DistSQLUpdateExecuteEngine {
private final String databaseName;
- private final ConnectionContext connectionContext;
+ private final DistSQLConnectionContext distsqlConnectionContext;
- public DistSQLUpdateExecuteEngine(final DistSQLStatement sqlStatement,
final String currentDatabaseName, final ContextManager contextManager, final
ConnectionContext connectionContext) {
+ public DistSQLUpdateExecuteEngine(final DistSQLStatement sqlStatement,
final String currentDatabaseName,
+ final ContextManager contextManager,
final DistSQLConnectionContext distsqlConnectionContext) {
this.sqlStatement = sqlStatement;
this.contextManager = contextManager;
databaseName = DatabaseNameUtils.getDatabaseName(sqlStatement,
currentDatabaseName);
- this.connectionContext = connectionContext;
+ this.distsqlConnectionContext = distsqlConnectionContext;
}
/**
@@ -90,7 +91,7 @@ public final class DistSQLUpdateExecuteEngine {
Optional<AdvancedDistSQLUpdateExecutor> advancedExecutor =
TypedSPILoader.findService(AdvancedDistSQLUpdateExecutor.class,
sqlStatement.getClass());
DistSQLUpdateExecutor executor = advancedExecutor.isPresent() ?
advancedExecutor.get() : TypedSPILoader.getService(DistSQLUpdateExecutor.class,
sqlStatement.getClass());
ShardingSphereDatabase database = null == databaseName ? null :
contextManager.getDatabase(databaseName);
- new DistSQLExecutorAwareSetter(executor).set(contextManager, database,
null, connectionContext, sqlStatement);
+ new DistSQLExecutorAwareSetter(executor).set(contextManager, database,
distsqlConnectionContext, sqlStatement);
new DistSQLExecutorRequiredChecker(executor).check(sqlStatement,
contextManager, database);
executor.executeUpdate(sqlStatement, contextManager);
}
diff --git
a/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/handler/distsql/DistSQLProxyBackendHandlerFactory.java
b/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/handler/distsql/DistSQLProxyBackendHandlerFactory.java
index 9d6ac67bb08..6976e5ba3d3 100644
---
a/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/handler/distsql/DistSQLProxyBackendHandlerFactory.java
+++
b/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/handler/distsql/DistSQLProxyBackendHandlerFactory.java
@@ -53,12 +53,12 @@ public final class DistSQLProxyBackendHandlerFactory {
return new DistSQLQueryProxyBackendHandler(sqlStatement,
queryContext, connectionSession, contextManager);
}
if (sqlStatement instanceof RDLStatement) {
- return new DistSQLUpdateProxyBackendHandler(sqlStatement,
connectionSession, contextManager);
+ return new DistSQLUpdateProxyBackendHandler(sqlStatement,
queryContext, connectionSession, contextManager);
}
if (sqlStatement instanceof RALStatement) {
return sqlStatement instanceof QueryableRALStatement
? new DistSQLQueryProxyBackendHandler(sqlStatement,
queryContext, connectionSession, contextManager)
- : new DistSQLUpdateProxyBackendHandler(sqlStatement,
connectionSession, contextManager);
+ : new DistSQLUpdateProxyBackendHandler(sqlStatement,
queryContext, connectionSession, contextManager);
}
throw new
UnsupportedSQLOperationException(sqlStatement.getClass().getName());
}
diff --git
a/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/handler/distsql/DistSQLUpdateProxyBackendHandler.java
b/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/handler/distsql/DistSQLUpdateProxyBackendHandler.java
index 3d45d34f630..f643e7f2b6c 100644
---
a/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/handler/distsql/DistSQLUpdateProxyBackendHandler.java
+++
b/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/handler/distsql/DistSQLUpdateProxyBackendHandler.java
@@ -17,8 +17,10 @@
package org.apache.shardingsphere.proxy.backend.handler.distsql;
+import
org.apache.shardingsphere.distsql.handler.engine.DistSQLConnectionContext;
import
org.apache.shardingsphere.distsql.handler.engine.update.DistSQLUpdateExecuteEngine;
import org.apache.shardingsphere.distsql.statement.DistSQLStatement;
+import org.apache.shardingsphere.infra.session.query.QueryContext;
import org.apache.shardingsphere.mode.manager.ContextManager;
import org.apache.shardingsphere.proxy.backend.response.header.ResponseHeader;
import
org.apache.shardingsphere.proxy.backend.response.header.update.UpdateResponseHeader;
@@ -35,9 +37,12 @@ public final class DistSQLUpdateProxyBackendHandler
implements DistSQLProxyBacke
private final DistSQLUpdateExecuteEngine engine;
- public DistSQLUpdateProxyBackendHandler(final DistSQLStatement
sqlStatement, final ConnectionSession connectionSession, final ContextManager
contextManager) {
+ public DistSQLUpdateProxyBackendHandler(final DistSQLStatement
sqlStatement, final QueryContext queryContext, final ConnectionSession
connectionSession, final ContextManager contextManager) {
this.sqlStatement = sqlStatement;
- engine = new DistSQLUpdateExecuteEngine(sqlStatement,
connectionSession.getUsedDatabaseName(), contextManager,
connectionSession.getConnectionContext());
+ DistSQLConnectionContext distsqlConnectionContext = new
DistSQLConnectionContext(queryContext,
+
connectionSession.getDatabaseConnectionManager().getConnectionSize(),
connectionSession.getProtocolType(),
+ connectionSession.getDatabaseConnectionManager(),
connectionSession.getStatementManager());
+ engine = new DistSQLUpdateExecuteEngine(sqlStatement,
connectionSession.getUsedDatabaseName(), contextManager,
distsqlConnectionContext);
}
@Override
diff --git
a/proxy/backend/core/src/test/java/org/apache/shardingsphere/proxy/backend/handler/distsql/DistSQLProxyBackendHandlerFactoryTest.java
b/proxy/backend/core/src/test/java/org/apache/shardingsphere/proxy/backend/handler/distsql/DistSQLProxyBackendHandlerFactoryTest.java
index cae0b4d9483..7e24f21eabc 100644
---
a/proxy/backend/core/src/test/java/org/apache/shardingsphere/proxy/backend/handler/distsql/DistSQLProxyBackendHandlerFactoryTest.java
+++
b/proxy/backend/core/src/test/java/org/apache/shardingsphere/proxy/backend/handler/distsql/DistSQLProxyBackendHandlerFactoryTest.java
@@ -114,7 +114,7 @@ class DistSQLProxyBackendHandlerFactoryTest {
void assertExecuteDataSourcesContext() throws SQLException {
RegisterStorageUnitStatement sqlStatement =
mock(RegisterStorageUnitStatement.class);
when(sqlStatement.getAttributes()).thenReturn(new
SQLStatementAttributes());
- assertThat(new DistSQLUpdateProxyBackendHandler(sqlStatement,
connectionSession, contextManager).execute(), isA(UpdateResponseHeader.class));
+ assertThat(new DistSQLUpdateProxyBackendHandler(sqlStatement, mock(),
connectionSession, contextManager).execute(), isA(UpdateResponseHeader.class));
}
@Test
@@ -122,21 +122,21 @@ class DistSQLProxyBackendHandlerFactoryTest {
when(contextManager.getDatabase("foo_db").getRuleMetaData()).thenReturn(new
RuleMetaData(Collections.emptyList()));
CreateShardingTableRuleStatement sqlStatement =
mock(CreateShardingTableRuleStatement.class);
when(sqlStatement.getAttributes()).thenReturn(new
SQLStatementAttributes());
- assertThat(new DistSQLUpdateProxyBackendHandler(sqlStatement,
connectionSession, contextManager).execute(), isA(UpdateResponseHeader.class));
+ assertThat(new DistSQLUpdateProxyBackendHandler(sqlStatement, mock(),
connectionSession, contextManager).execute(), isA(UpdateResponseHeader.class));
}
@Test
void assertExecuteAddResourceContext() throws SQLException {
RegisterStorageUnitStatement sqlStatement =
mock(RegisterStorageUnitStatement.class);
when(sqlStatement.getAttributes()).thenReturn(new
SQLStatementAttributes());
- assertThat(new DistSQLUpdateProxyBackendHandler(sqlStatement,
connectionSession, contextManager).execute(), isA(UpdateResponseHeader.class));
+ assertThat(new DistSQLUpdateProxyBackendHandler(sqlStatement, mock(),
connectionSession, contextManager).execute(), isA(UpdateResponseHeader.class));
}
@Test
void assertExecuteAlterResourceContext() throws SQLException {
AlterStorageUnitStatement sqlStatement =
mock(AlterStorageUnitStatement.class);
when(sqlStatement.getAttributes()).thenReturn(new
SQLStatementAttributes());
- assertThat(new DistSQLUpdateProxyBackendHandler(sqlStatement,
connectionSession, contextManager).execute(), isA(UpdateResponseHeader.class));
+ assertThat(new DistSQLUpdateProxyBackendHandler(sqlStatement, mock(),
connectionSession, contextManager).execute(), isA(UpdateResponseHeader.class));
}
@Test
@@ -145,7 +145,7 @@ class DistSQLProxyBackendHandlerFactoryTest {
when(contextManager.getDatabase("foo_db")).thenReturn(database);
AlterShadowRuleStatement sqlStatement =
mock(AlterShadowRuleStatement.class);
when(sqlStatement.getAttributes()).thenReturn(new
SQLStatementAttributes());
- assertThat(new DistSQLUpdateProxyBackendHandler(sqlStatement,
connectionSession, contextManager).execute(), isA(UpdateResponseHeader.class));
+ assertThat(new DistSQLUpdateProxyBackendHandler(sqlStatement, mock(),
connectionSession, contextManager).execute(), isA(UpdateResponseHeader.class));
}
@Test
@@ -154,7 +154,7 @@ class DistSQLProxyBackendHandlerFactoryTest {
when(contextManager.getDatabase("foo_db")).thenReturn(database);
CreateShadowRuleStatement sqlStatement =
mock(CreateShadowRuleStatement.class);
when(sqlStatement.getAttributes()).thenReturn(new
SQLStatementAttributes());
- assertThat(new DistSQLUpdateProxyBackendHandler(sqlStatement,
connectionSession, contextManager).execute(), isA(UpdateResponseHeader.class));
+ assertThat(new DistSQLUpdateProxyBackendHandler(sqlStatement, mock(),
connectionSession, contextManager).execute(), isA(UpdateResponseHeader.class));
}
@Test
@@ -163,7 +163,7 @@ class DistSQLProxyBackendHandlerFactoryTest {
when(contextManager.getDatabase("foo_db")).thenReturn(database);
DropShadowRuleStatement sqlStatement =
mock(DropShadowRuleStatement.class);
when(sqlStatement.getAttributes()).thenReturn(new
SQLStatementAttributes());
- assertThat(new DistSQLUpdateProxyBackendHandler(sqlStatement,
connectionSession, contextManager).execute(), isA(UpdateResponseHeader.class));
+ assertThat(new DistSQLUpdateProxyBackendHandler(sqlStatement, mock(),
connectionSession, contextManager).execute(), isA(UpdateResponseHeader.class));
}
@Test
@@ -172,7 +172,7 @@ class DistSQLProxyBackendHandlerFactoryTest {
when(contextManager.getDatabase("foo_db")).thenReturn(database);
AlterDefaultShadowAlgorithmStatement statement = new
AlterDefaultShadowAlgorithmStatement(
new ShadowAlgorithmSegment("foo", new
AlgorithmSegment("SQL_HINT", PropertiesBuilder.build(new Property("type",
"value")))));
- assertThat(new DistSQLUpdateProxyBackendHandler(statement,
connectionSession, contextManager).execute(), isA(UpdateResponseHeader.class));
+ assertThat(new DistSQLUpdateProxyBackendHandler(statement, mock(),
connectionSession, contextManager).execute(), isA(UpdateResponseHeader.class));
}
@Test
@@ -205,33 +205,33 @@ class DistSQLProxyBackendHandlerFactoryTest {
when(contextManager.getDatabase("foo_db")).thenReturn(database);
DropShadowAlgorithmStatement sqlStatement =
mock(DropShadowAlgorithmStatement.class);
when(sqlStatement.getAttributes()).thenReturn(new
SQLStatementAttributes());
- assertThat(new DistSQLUpdateProxyBackendHandler(sqlStatement,
connectionSession, contextManager).execute(), isA(UpdateResponseHeader.class));
+ assertThat(new DistSQLUpdateProxyBackendHandler(sqlStatement, mock(),
connectionSession, contextManager).execute(), isA(UpdateResponseHeader.class));
}
@Test
void assertExecuteDropResourceContext() throws SQLException {
UnregisterStorageUnitStatement sqlStatement =
mock(UnregisterStorageUnitStatement.class);
when(sqlStatement.getAttributes()).thenReturn(new
SQLStatementAttributes());
- assertThat(new DistSQLUpdateProxyBackendHandler(sqlStatement,
connectionSession, contextManager).execute(), isA(UpdateResponseHeader.class));
+ assertThat(new DistSQLUpdateProxyBackendHandler(sqlStatement, mock(),
connectionSession, contextManager).execute(), isA(UpdateResponseHeader.class));
}
@Test
void assertExecuteDropReadwriteSplittingRuleContext() {
assertThrows(MissingRequiredRuleException.class,
- () -> new
DistSQLUpdateProxyBackendHandler(mock(DropReadwriteSplittingRuleStatement.class,
RETURNS_DEEP_STUBS), connectionSession, contextManager).execute());
+ () -> new
DistSQLUpdateProxyBackendHandler(mock(DropReadwriteSplittingRuleStatement.class,
RETURNS_DEEP_STUBS), mock(), connectionSession, contextManager).execute());
}
@Test
void assertExecuteCreateReadwriteSplittingRuleContext() throws
SQLException {
CreateReadwriteSplittingRuleStatement sqlStatement =
mock(CreateReadwriteSplittingRuleStatement.class);
when(sqlStatement.getAttributes()).thenReturn(new
SQLStatementAttributes());
- assertThat(new DistSQLUpdateProxyBackendHandler(sqlStatement,
connectionSession, contextManager).execute(), isA(UpdateResponseHeader.class));
+ assertThat(new DistSQLUpdateProxyBackendHandler(sqlStatement, mock(),
connectionSession, contextManager).execute(), isA(UpdateResponseHeader.class));
}
@Test
void assertExecuteAlterReadwriteSplittingRuleContext() {
assertThrows(MissingRequiredRuleException.class,
- () -> new
DistSQLUpdateProxyBackendHandler(mock(AlterReadwriteSplittingRuleStatement.class,
RETURNS_DEEP_STUBS), connectionSession, contextManager).execute());
+ () -> new
DistSQLUpdateProxyBackendHandler(mock(AlterReadwriteSplittingRuleStatement.class,
RETURNS_DEEP_STUBS), mock(), connectionSession, contextManager).execute());
}
private ShardingSphereDatabase mockDatabaseWithRule() {
diff --git
a/proxy/backend/core/src/test/java/org/apache/shardingsphere/proxy/backend/handler/distsql/DistSQLUpdateProxyBackendHandlerTest.java
b/proxy/backend/core/src/test/java/org/apache/shardingsphere/proxy/backend/handler/distsql/DistSQLUpdateProxyBackendHandlerTest.java
index dc5630f22e8..c95417aef29 100644
---
a/proxy/backend/core/src/test/java/org/apache/shardingsphere/proxy/backend/handler/distsql/DistSQLUpdateProxyBackendHandlerTest.java
+++
b/proxy/backend/core/src/test/java/org/apache/shardingsphere/proxy/backend/handler/distsql/DistSQLUpdateProxyBackendHandlerTest.java
@@ -45,6 +45,7 @@ import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.isA;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.RETURNS_DEEP_STUBS;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
@@ -59,7 +60,7 @@ class DistSQLUpdateProxyBackendHandlerTest {
@Test
void assertEmptyStorageUnit() {
when(contextManager.getDatabase("foo_db")).thenReturn(new
ShardingSphereDatabase("foo_db", databaseType, mock(), mock(),
Collections.emptyList()));
- DistSQLUpdateProxyBackendHandler backendHandler = new
DistSQLUpdateProxyBackendHandler(new RefreshTableMetaDataStatement(),
mockConnectionSession("foo_db"), contextManager);
+ DistSQLUpdateProxyBackendHandler backendHandler = new
DistSQLUpdateProxyBackendHandler(new RefreshTableMetaDataStatement(), mock(),
mockConnectionSession("foo_db"), contextManager);
assertThrows(EmptyStorageUnitException.class, backendHandler::execute);
}
@@ -69,7 +70,7 @@ class DistSQLUpdateProxyBackendHandlerTest {
when(resourceMetaData.getStorageUnits()).thenReturn(Collections.singletonMap("ds_0",
mock(StorageUnit.class)));
when(contextManager.getDatabase("foo_db")).thenReturn(new
ShardingSphereDatabase("foo_db", databaseType, resourceMetaData, mock(),
Collections.emptyList()));
DistSQLUpdateProxyBackendHandler backendHandler = new
DistSQLUpdateProxyBackendHandler(
- new RefreshTableMetaDataStatement("t_order", "ds_1", null),
mockConnectionSession("foo_db"), contextManager);
+ new RefreshTableMetaDataStatement("t_order", "ds_1", null),
mock(), mockConnectionSession("foo_db"), contextManager);
assertThrows(MissingRequiredStorageUnitsException.class,
backendHandler::execute);
}
@@ -79,7 +80,7 @@ class DistSQLUpdateProxyBackendHandlerTest {
when(resourceMetaData.getStorageUnits()).thenReturn(Collections.singletonMap("ds_0",
mock(StorageUnit.class)));
when(contextManager.getDatabase("foo_db")).thenReturn(new
ShardingSphereDatabase("foo_db", databaseType, resourceMetaData, mock(),
Collections.emptyList()));
DistSQLUpdateProxyBackendHandler backendHandler = new
DistSQLUpdateProxyBackendHandler(
- new RefreshTableMetaDataStatement("t_order", "ds_0",
"bar_db"), mockConnectionSession("foo_db"), contextManager);
+ new RefreshTableMetaDataStatement("t_order", "ds_0",
"bar_db"), mock(), mockConnectionSession("foo_db"), contextManager);
assertThrows(SchemaNotFoundException.class, backendHandler::execute);
}
@@ -94,7 +95,7 @@ class DistSQLUpdateProxyBackendHandlerTest {
when(database.getResourceMetaData()).thenReturn(resourceMetaData);
when(contextManager.getDatabase("foo_db")).thenReturn(database);
DistSQLUpdateProxyBackendHandler backendHandler = new
DistSQLUpdateProxyBackendHandler(
- new RefreshTableMetaDataStatement("t_order", "ds_0",
"foo_db"), mockConnectionSession("foo_db"), contextManager);
+ new RefreshTableMetaDataStatement("t_order", "ds_0",
"foo_db"), mock(), mockConnectionSession("foo_db"), contextManager);
assertThrows(TableNotFoundException.class, backendHandler::execute);
}
@@ -107,13 +108,14 @@ class DistSQLUpdateProxyBackendHandlerTest {
when(database.getProtocolType()).thenReturn(databaseType);
when(database.getResourceMetaData()).thenReturn(resourceMetaData);
when(contextManager.getDatabase("foo_db")).thenReturn(database);
- ResponseHeader actual = new DistSQLUpdateProxyBackendHandler(new
RefreshTableMetaDataStatement(), mockConnectionSession("foo_db"),
contextManager).execute();
+ ResponseHeader actual = new DistSQLUpdateProxyBackendHandler(new
RefreshTableMetaDataStatement(), mock(), mockConnectionSession("foo_db"),
contextManager).execute();
assertThat(actual, isA(UpdateResponseHeader.class));
}
private ConnectionSession mockConnectionSession(final String databaseName)
{
- ConnectionSession result = mock(ConnectionSession.class);
+ ConnectionSession result = mock(ConnectionSession.class,
RETURNS_DEEP_STUBS);
when(result.getUsedDatabaseName()).thenReturn(databaseName);
+
when(result.getDatabaseConnectionManager().getConnectionSize()).thenReturn(1);
return result;
}
}
diff --git
a/test/it/distsql/src/main/java/org/apache/shardingsphere/test/it/distsql/handler/engine/update/DistSQLGlobalRuleDefinitionExecutorAssert.java
b/test/it/distsql/src/main/java/org/apache/shardingsphere/test/it/distsql/handler/engine/update/DistSQLGlobalRuleDefinitionExecutorAssert.java
index 00548ba1dca..4b80aeb6132 100644
---
a/test/it/distsql/src/main/java/org/apache/shardingsphere/test/it/distsql/handler/engine/update/DistSQLGlobalRuleDefinitionExecutorAssert.java
+++
b/test/it/distsql/src/main/java/org/apache/shardingsphere/test/it/distsql/handler/engine/update/DistSQLGlobalRuleDefinitionExecutorAssert.java
@@ -18,13 +18,13 @@
package org.apache.shardingsphere.test.it.distsql.handler.engine.update;
import lombok.RequiredArgsConstructor;
+import
org.apache.shardingsphere.distsql.handler.engine.DistSQLConnectionContext;
import
org.apache.shardingsphere.distsql.handler.engine.update.DistSQLUpdateExecuteEngine;
import org.apache.shardingsphere.distsql.statement.DistSQLStatement;
import org.apache.shardingsphere.infra.config.rule.RuleConfiguration;
import
org.apache.shardingsphere.infra.config.rule.scope.GlobalRuleConfiguration;
import org.apache.shardingsphere.infra.metadata.database.rule.RuleMetaData;
import org.apache.shardingsphere.infra.rule.scope.GlobalRule;
-import org.apache.shardingsphere.infra.session.connection.ConnectionContext;
import org.apache.shardingsphere.mode.manager.ContextManager;
import
org.apache.shardingsphere.mode.persist.service.MetaDataManagerPersistService;
@@ -58,7 +58,7 @@ public final class DistSQLGlobalRuleDefinitionExecutorAssert {
public void assertExecuteUpdate(final GlobalRuleConfiguration ruleConfig,
final DistSQLStatement sqlStatement, final
RuleConfiguration matchedRuleConfig, final Class<? extends Exception>
expectedException) throws SQLException {
ContextManager contextManager = mockContextManager(ruleConfig);
- DistSQLUpdateExecuteEngine engine = new
DistSQLUpdateExecuteEngine(sqlStatement, null, contextManager,
mock(ConnectionContext.class));
+ DistSQLUpdateExecuteEngine engine = new
DistSQLUpdateExecuteEngine(sqlStatement, null, contextManager,
mock(DistSQLConnectionContext.class));
if (null != expectedException) {
assertThrows(expectedException, engine::executeUpdate);
return;
