You are viewing a plain text version of this content. The canonical link for it is here.
Posted to notifications@shardingsphere.apache.org by me...@apache.org on 2022/01/05 08:11:15 UTC
[shardingsphere] branch master updated: Integrate governance API for traffic (#14531)
This is an automated email from the ASF dual-hosted git repository.
menghaoran 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 063013f Integrate governance API for traffic (#14531)
063013f is described below
commit 063013f39c1f6ed31bf3d15adeaafc4d1292057f
Author: Zhengqiang Duan <du...@apache.org>
AuthorDate: Wed Jan 5 16:10:20 2022 +0800
Integrate governance API for traffic (#14531)
* integrate governance API for traffic
* add more unit test for TrafficRuleTest
* modify governance api
---
.../driver/executor/DriverExecutor.java | 2 +-
.../statement/ShardingSpherePreparedStatement.java | 2 +-
.../core/statement/ShardingSphereStatement.java | 2 +-
.../traffic/engine/TrafficEngine.java | 22 +++--
.../traffic/executor/TrafficExecutorFactory.java | 9 +-
.../traffic/executor/jdbc/JDBCTrafficExecutor.java | 73 +++++++++++++---
.../shardingsphere/traffic/rule/TrafficRule.java | 13 +++
.../traffic/rule/TrafficRuleTest.java | 99 +++++++++++++++++++++-
8 files changed, 196 insertions(+), 26 deletions(-)
diff --git a/shardingsphere-jdbc/shardingsphere-jdbc-core/src/main/java/org/apache/shardingsphere/driver/executor/DriverExecutor.java b/shardingsphere-jdbc/shardingsphere-jdbc-core/src/main/java/org/apache/shardingsphere/driver/executor/DriverExecutor.java
index 9ef2e79..b3015a8 100644
--- a/shardingsphere-jdbc/shardingsphere-jdbc-core/src/main/java/org/apache/shardingsphere/driver/executor/DriverExecutor.java
+++ b/shardingsphere-jdbc/shardingsphere-jdbc-core/src/main/java/org/apache/shardingsphere/driver/executor/DriverExecutor.java
@@ -49,7 +49,7 @@ public final class DriverExecutor implements AutoCloseable {
regularExecutor = new DriverJDBCExecutor(connection.getSchema(), metaDataContexts, jdbcExecutor);
rawExecutor = new RawExecutor(metaDataContexts.getExecutorEngine(), connection.isHoldTransaction(), metaDataContexts.getProps());
federationExecutor = FederationExecutorFactory.newInstance(connection.getSchema(), metaDataContexts.getOptimizerContext(), metaDataContexts.getProps(), jdbcExecutor);
- trafficExecutor = TrafficExecutorFactory.newInstance(connection.getContextManager());
+ trafficExecutor = TrafficExecutorFactory.newInstance(connection.getSchema(), metaDataContexts);
}
/**
diff --git a/shardingsphere-jdbc/shardingsphere-jdbc-core/src/main/java/org/apache/shardingsphere/driver/jdbc/core/statement/ShardingSpherePreparedStatement.java b/shardingsphere-jdbc/shardingsphere-jdbc-core/src/main/java/org/apache/shardingsphere/driver/jdbc/core/statement/ShardingSpherePreparedStatement.java
index c9b5297..9e79414 100644
--- a/shardingsphere-jdbc/shardingsphere-jdbc-core/src/main/java/org/apache/shardingsphere/driver/jdbc/core/statement/ShardingSpherePreparedStatement.java
+++ b/shardingsphere-jdbc/shardingsphere-jdbc-core/src/main/java/org/apache/shardingsphere/driver/jdbc/core/statement/ShardingSpherePreparedStatement.java
@@ -208,7 +208,7 @@ public final class ShardingSpherePreparedStatement extends AbstractPreparedState
private TrafficContext createTrafficContext(final LogicSQL logicSQL) {
Optional<TrafficRule> trafficRule = metaDataContexts.getGlobalRuleMetaData().findSingleRule(TrafficRule.class);
- return trafficRule.map(optional -> new TrafficEngine(optional).dispatch(logicSQL)).orElse(new TrafficContext());
+ return trafficRule.map(optional -> new TrafficEngine(optional, metaDataContexts).dispatch(logicSQL)).orElse(new TrafficContext());
}
private void resetParameters() throws SQLException {
diff --git a/shardingsphere-jdbc/shardingsphere-jdbc-core/src/main/java/org/apache/shardingsphere/driver/jdbc/core/statement/ShardingSphereStatement.java b/shardingsphere-jdbc/shardingsphere-jdbc-core/src/main/java/org/apache/shardingsphere/driver/jdbc/core/statement/ShardingSphereStatement.java
index 3e8c681..acefef6 100644
--- a/shardingsphere-jdbc/shardingsphere-jdbc-core/src/main/java/org/apache/shardingsphere/driver/jdbc/core/statement/ShardingSphereStatement.java
+++ b/shardingsphere-jdbc/shardingsphere-jdbc-core/src/main/java/org/apache/shardingsphere/driver/jdbc/core/statement/ShardingSphereStatement.java
@@ -165,7 +165,7 @@ public final class ShardingSphereStatement extends AbstractStatementAdapter {
private TrafficContext createTrafficContext(final LogicSQL logicSQL) {
Optional<TrafficRule> trafficRule = metaDataContexts.getGlobalRuleMetaData().findSingleRule(TrafficRule.class);
- return trafficRule.map(optional -> new TrafficEngine(optional).dispatch(logicSQL)).orElse(new TrafficContext());
+ return trafficRule.map(optional -> new TrafficEngine(optional, metaDataContexts).dispatch(logicSQL)).orElse(new TrafficContext());
}
private List<ResultSet> getShardingSphereResultSets() {
diff --git a/shardingsphere-kernel/shardingsphere-traffic/shardingsphere-traffic-core/src/main/java/org/apache/shardingsphere/traffic/engine/TrafficEngine.java b/shardingsphere-kernel/shardingsphere-traffic/shardingsphere-traffic-core/src/main/java/org/apache/shardingsphere/traffic/engine/TrafficEngine.java
index b7bc172..7226ae9 100644
--- a/shardingsphere-kernel/shardingsphere-traffic/shardingsphere-traffic-core/src/main/java/org/apache/shardingsphere/traffic/engine/TrafficEngine.java
+++ b/shardingsphere-kernel/shardingsphere-traffic/shardingsphere-traffic-core/src/main/java/org/apache/shardingsphere/traffic/engine/TrafficEngine.java
@@ -19,13 +19,16 @@ package org.apache.shardingsphere.traffic.engine;
import lombok.RequiredArgsConstructor;
import org.apache.shardingsphere.infra.binder.LogicSQL;
+import org.apache.shardingsphere.infra.instance.ComputeNodeInstance;
+import org.apache.shardingsphere.infra.instance.InstanceType;
+import org.apache.shardingsphere.mode.metadata.MetaDataContexts;
import org.apache.shardingsphere.traffic.context.TrafficContext;
import org.apache.shardingsphere.traffic.rule.TrafficRule;
import org.apache.shardingsphere.traffic.rule.TrafficStrategyRule;
import org.apache.shardingsphere.traffic.spi.TrafficLoadBalanceAlgorithm;
+import java.util.ArrayList;
import java.util.Collection;
-import java.util.Collections;
import java.util.List;
import java.util.Optional;
@@ -37,6 +40,8 @@ public final class TrafficEngine {
private final TrafficRule trafficRule;
+ private final MetaDataContexts metaDataContexts;
+
/**
* Dispatch.
*
@@ -50,13 +55,20 @@ public final class TrafficEngine {
return result;
}
List<String> dataSourceNames = getDataSourceNamesByLabels(strategyRule.get().getLabels());
- TrafficLoadBalanceAlgorithm loadBalancer = trafficRule.findLoadBalancer(strategyRule.get().getLoadBalancerName());
- result.setDataSourceName(loadBalancer.getDataSourceName(dataSourceNames));
+ if (!dataSourceNames.isEmpty()) {
+ TrafficLoadBalanceAlgorithm loadBalancer = trafficRule.findLoadBalancer(strategyRule.get().getLoadBalancerName());
+ result.setDataSourceName(loadBalancer.getDataSourceName(dataSourceNames));
+ }
return result;
}
private List<String> getDataSourceNamesByLabels(final Collection<String> labels) {
- // TODO implements this logic when provide mode api to get dataSource configs by labels
- return Collections.emptyList();
+ List<String> result = new ArrayList<>();
+ if (metaDataContexts.getMetaDataPersistService().isPresent()) {
+ for (ComputeNodeInstance each : metaDataContexts.getMetaDataPersistService().get().loadComputeNodeInstances(InstanceType.PROXY, labels)) {
+ result.add(each.getIp() + "@" + each.getPort());
+ }
+ }
+ return result;
}
}
diff --git a/shardingsphere-kernel/shardingsphere-traffic/shardingsphere-traffic-core/src/main/java/org/apache/shardingsphere/traffic/executor/TrafficExecutorFactory.java b/shardingsphere-kernel/shardingsphere-traffic/shardingsphere-traffic-core/src/main/java/org/apache/shardingsphere/traffic/executor/TrafficExecutorFactory.java
index 740fcef..0850775 100644
--- a/shardingsphere-kernel/shardingsphere-traffic/shardingsphere-traffic-core/src/main/java/org/apache/shardingsphere/traffic/executor/TrafficExecutorFactory.java
+++ b/shardingsphere-kernel/shardingsphere-traffic/shardingsphere-traffic-core/src/main/java/org/apache/shardingsphere/traffic/executor/TrafficExecutorFactory.java
@@ -19,7 +19,7 @@ package org.apache.shardingsphere.traffic.executor;
import lombok.AccessLevel;
import lombok.NoArgsConstructor;
-import org.apache.shardingsphere.mode.manager.ContextManager;
+import org.apache.shardingsphere.mode.metadata.MetaDataContexts;
import org.apache.shardingsphere.traffic.executor.jdbc.JDBCTrafficExecutor;
/**
@@ -31,10 +31,11 @@ public final class TrafficExecutorFactory {
/**
* Create new instance of traffic executor factory.
*
- * @param contextManager context manager
+ * @param schema schema
+ * @param metaDataContexts meta data contexts
* @return new instance of traffic executor
*/
- public static TrafficExecutor newInstance(final ContextManager contextManager) {
- return new JDBCTrafficExecutor(contextManager);
+ public static TrafficExecutor newInstance(final String schema, final MetaDataContexts metaDataContexts) {
+ return new JDBCTrafficExecutor(schema, metaDataContexts);
}
}
diff --git a/shardingsphere-kernel/shardingsphere-traffic/shardingsphere-traffic-core/src/main/java/org/apache/shardingsphere/traffic/executor/jdbc/JDBCTrafficExecutor.java b/shardingsphere-kernel/shardingsphere-traffic/shardingsphere-traffic-core/src/main/java/org/apache/shardingsphere/traffic/executor/jdbc/JDBCTrafficExecutor.java
index 3fd1f86..96fe1b7 100644
--- a/shardingsphere-kernel/shardingsphere-traffic/shardingsphere-traffic-core/src/main/java/org/apache/shardingsphere/traffic/executor/jdbc/JDBCTrafficExecutor.java
+++ b/shardingsphere-kernel/shardingsphere-traffic/shardingsphere-traffic-core/src/main/java/org/apache/shardingsphere/traffic/executor/jdbc/JDBCTrafficExecutor.java
@@ -17,17 +17,24 @@
package org.apache.shardingsphere.traffic.executor.jdbc;
+import com.zaxxer.hikari.HikariDataSource;
import org.apache.shardingsphere.infra.binder.LogicSQL;
import org.apache.shardingsphere.infra.config.datasource.DataSourceConfiguration;
import org.apache.shardingsphere.infra.config.datasource.pool.creator.DataSourcePoolCreatorUtil;
+import org.apache.shardingsphere.infra.database.metadata.DataSourceMetaData;
+import org.apache.shardingsphere.infra.database.type.DatabaseTypeRegistry;
import org.apache.shardingsphere.infra.exception.ShardingSphereException;
-import org.apache.shardingsphere.mode.manager.ContextManager;
+import org.apache.shardingsphere.infra.instance.ComputeNodeInstance;
+import org.apache.shardingsphere.infra.instance.InstanceType;
+import org.apache.shardingsphere.infra.metadata.user.ShardingSphereUser;
+import org.apache.shardingsphere.mode.metadata.MetaDataContexts;
import org.apache.shardingsphere.spi.ShardingSphereServiceLoader;
import org.apache.shardingsphere.spi.typed.TypedSPIRegistry;
import org.apache.shardingsphere.traffic.executor.TrafficExecutor;
import org.apache.shardingsphere.traffic.executor.TrafficExecutorCallback;
import org.apache.shardingsphere.traffic.executor.context.TrafficExecutorContext;
import org.apache.shardingsphere.traffic.executor.context.builder.TrafficExecutorContextBuilder;
+import org.apache.shardingsphere.traffic.rule.TrafficRule;
import javax.sql.DataSource;
import java.sql.Connection;
@@ -35,10 +42,11 @@ import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;
-import java.util.Collections;
+import java.util.Collection;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
+import java.util.Optional;
import java.util.Properties;
import java.util.concurrent.ConcurrentHashMap;
@@ -47,7 +55,13 @@ import java.util.concurrent.ConcurrentHashMap;
*/
public final class JDBCTrafficExecutor implements TrafficExecutor {
- private static final Map<String, TrafficExecutorContextBuilder> TYPE_CONTEXT_BUILDERS = new ConcurrentHashMap<>();
+ private static final Map<String, TrafficExecutorContextBuilder<?>> TYPE_CONTEXT_BUILDERS = new ConcurrentHashMap<>();
+
+ private static final String JDBC_URL = "jdbcUrl";
+
+ private static final String USER_NAME = "username";
+
+ private static final String PASSWORD = "password";
private final Map<String, DataSource> dataSources = new LinkedHashMap<>();
@@ -57,29 +71,64 @@ public final class JDBCTrafficExecutor implements TrafficExecutor {
ShardingSphereServiceLoader.register(TrafficExecutorContextBuilder.class);
}
- public JDBCTrafficExecutor(final ContextManager contextManager) {
- dataSources.putAll(DataSourcePoolCreatorUtil.getDataSourceMap(createDataSourceConfigs()));
+ public JDBCTrafficExecutor(final String schema, final MetaDataContexts metaDataContexts) {
+ Optional<TrafficRule> trafficRule = metaDataContexts.getGlobalRuleMetaData().findSingleRule(TrafficRule.class);
+ if (trafficRule.isPresent() && metaDataContexts.getMetaDataPersistService().isPresent()) {
+ Map<String, DataSourceConfiguration> dataSourceConfigs = metaDataContexts.getMetaDataPersistService().get().getDataSourceService().load(schema);
+ if (dataSourceConfigs.isEmpty()) {
+ throw new ShardingSphereException("Can not get dataSource configs from meta data.");
+ }
+ DataSourceConfiguration dataSourceConfigSample = dataSourceConfigs.values().iterator().next();
+ Collection<ComputeNodeInstance> instances = metaDataContexts.getMetaDataPersistService().get().loadComputeNodeInstances(InstanceType.PROXY, trafficRule.get().getLabels());
+ dataSources.putAll(DataSourcePoolCreatorUtil.getDataSourceMap(createDataSourceConfigs(instances, dataSourceConfigSample, schema)));
+ }
+ }
+
+ private Map<String, DataSourceConfiguration> createDataSourceConfigs(final Collection<ComputeNodeInstance> instances,
+ final DataSourceConfiguration dataSourceConfigSample, final String schema) {
+ Map<String, DataSourceConfiguration> result = new LinkedHashMap<>();
+ for (ComputeNodeInstance each : instances) {
+ result.put(each.getIp() + "@" + each.getPort(), createDataSourceConfig(each, dataSourceConfigSample, schema));
+ }
+ return result;
+ }
+
+ private DataSourceConfiguration createDataSourceConfig(final ComputeNodeInstance instance,
+ final DataSourceConfiguration dataSourceConfigSample, final String schema) {
+ Map<String, Object> props = dataSourceConfigSample.getProps();
+ props.put(JDBC_URL, createJdbcUrl(instance, schema, props));
+ if (instance.getUsers().isEmpty()) {
+ throw new ShardingSphereException("Can not get users from meta data.");
+ }
+ ShardingSphereUser user = instance.getUsers().iterator().next();
+ props.put(USER_NAME, user.getGrantee().getUsername());
+ props.put(PASSWORD, user.getPassword());
+ DataSourceConfiguration result = new DataSourceConfiguration(HikariDataSource.class.getName());
+ result.getProps().putAll(props);
+ return result;
}
- private Map<String, DataSourceConfiguration> createDataSourceConfigs() {
- // TODO Use governance API to create data source configuration
- return Collections.emptyMap();
+ private String createJdbcUrl(final ComputeNodeInstance instance, final String schema, final Map<String, Object> props) {
+ String jdbcUrl = String.valueOf(props.get(JDBC_URL));
+ String username = String.valueOf(props.get(USER_NAME));
+ DataSourceMetaData dataSourceMetaData = DatabaseTypeRegistry.getDatabaseTypeByURL(jdbcUrl).getDataSourceMetaData(jdbcUrl, username);
+ return jdbcUrl.replace(dataSourceMetaData.getHostname(), instance.getIp())
+ .replace(String.valueOf(dataSourceMetaData.getPort()), instance.getPort()).replace(dataSourceMetaData.getCatalog(), schema);
}
@SuppressWarnings({"unchecked", "rawtypes"})
@Override
public TrafficExecutorContext<Statement> prepare(final LogicSQL logicSQL, final String dataSourceName, final String type) throws SQLException {
if (!dataSources.containsKey(dataSourceName)) {
- throw new ShardingSphereException("Can not get dataSource by %.", dataSourceName);
+ throw new ShardingSphereException("Can not get dataSource of %.", dataSourceName);
}
DataSource dataSource = dataSources.get(dataSourceName);
TrafficExecutorContextBuilder builder = getCachedTrafficExecutorContextBuilder(type);
return builder.build(logicSQL, dataSource.getConnection());
}
- @SuppressWarnings("rawtypes")
- private TrafficExecutorContextBuilder getCachedTrafficExecutorContextBuilder(final String type) {
- TrafficExecutorContextBuilder result;
+ private TrafficExecutorContextBuilder<?> getCachedTrafficExecutorContextBuilder(final String type) {
+ TrafficExecutorContextBuilder<?> result;
if (null == (result = TYPE_CONTEXT_BUILDERS.get(type))) {
result = TYPE_CONTEXT_BUILDERS.computeIfAbsent(type, key -> TypedSPIRegistry.getRegisteredService(TrafficExecutorContextBuilder.class, key, new Properties()));
}
diff --git a/shardingsphere-kernel/shardingsphere-traffic/shardingsphere-traffic-core/src/main/java/org/apache/shardingsphere/traffic/rule/TrafficRule.java b/shardingsphere-kernel/shardingsphere-traffic/shardingsphere-traffic-core/src/main/java/org/apache/shardingsphere/traffic/rule/TrafficRule.java
index b5e7f9e..f80c3ce 100644
--- a/shardingsphere-kernel/shardingsphere-traffic/shardingsphere-traffic-core/src/main/java/org/apache/shardingsphere/traffic/rule/TrafficRule.java
+++ b/shardingsphere-kernel/shardingsphere-traffic/shardingsphere-traffic-core/src/main/java/org/apache/shardingsphere/traffic/rule/TrafficRule.java
@@ -122,4 +122,17 @@ public final class TrafficRule implements GlobalRule {
Preconditions.checkState(null != loadBalanceAlgorithm, "Traffic load balance algorithm can not be null.");
return loadBalanceAlgorithm;
}
+
+ /**
+ * Get label collection.
+ *
+ * @return label collection
+ */
+ public Collection<String> getLabels() {
+ Collection<String> result = new LinkedList<>();
+ for (TrafficStrategyRule each : trafficStrategyRules) {
+ result.addAll(each.getLabels());
+ }
+ return result;
+ }
}
diff --git a/shardingsphere-kernel/shardingsphere-traffic/shardingsphere-traffic-core/src/test/java/org/apache/shardingsphere/traffic/rule/TrafficRuleTest.java b/shardingsphere-kernel/shardingsphere-traffic/shardingsphere-traffic-core/src/test/java/org/apache/shardingsphere/traffic/rule/TrafficRuleTest.java
index 0dd2a14..3c6954c 100644
--- a/shardingsphere-kernel/shardingsphere-traffic/shardingsphere-traffic-core/src/test/java/org/apache/shardingsphere/traffic/rule/TrafficRuleTest.java
+++ b/shardingsphere-kernel/shardingsphere-traffic/shardingsphere-traffic-core/src/test/java/org/apache/shardingsphere/traffic/rule/TrafficRuleTest.java
@@ -17,18 +17,113 @@
package org.apache.shardingsphere.traffic.rule;
+import org.apache.shardingsphere.infra.binder.LogicSQL;
+import org.apache.shardingsphere.infra.binder.statement.SQLStatementContext;
+import org.apache.shardingsphere.infra.binder.statement.dml.SelectStatementContext;
+import org.apache.shardingsphere.infra.config.algorithm.ShardingSphereAlgorithmConfiguration;
+import org.apache.shardingsphere.infra.metadata.ShardingSphereMetaData;
+import org.apache.shardingsphere.sql.parser.sql.common.segment.dml.item.ProjectionsSegment;
+import org.apache.shardingsphere.sql.parser.sql.common.segment.generic.CommentSegment;
+import org.apache.shardingsphere.sql.parser.sql.dialect.statement.mysql.dml.MySQLSelectStatement;
import org.apache.shardingsphere.traffic.api.config.TrafficRuleConfiguration;
+import org.apache.shardingsphere.traffic.api.config.TrafficStrategyConfiguration;
+import org.apache.shardingsphere.traffic.spi.TrafficLoadBalanceAlgorithm;
import org.junit.Test;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Properties;
+
import static org.hamcrest.CoreMatchers.is;
+import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertThat;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
public final class TrafficRuleTest {
@Test
public void assertGetRuleType() {
- TrafficRuleConfiguration ruleConfig = new TrafficRuleConfiguration();
- TrafficRule authorityRule = new TrafficRule(ruleConfig);
+ TrafficRule authorityRule = new TrafficRule(new TrafficRuleConfiguration());
assertThat(authorityRule.getType(), is(TrafficRule.class.getSimpleName()));
}
+
+ @Test
+ public void assertFindMatchedStrategyRule() {
+ TrafficRule trafficRule = new TrafficRule(createTrafficRuleConfig());
+ Optional<TrafficStrategyRule> actual = trafficRule.findMatchedStrategyRule(createLogicSQL(true));
+ assertTrue(actual.isPresent());
+ assertThat(actual.get().getName(), is("sql_hint_traffic"));
+ assertThat(actual.get().getLabels(), is(Arrays.asList("OLTP", "OLAP")));
+ assertThat(actual.get().getAlgorithmName(), is("sql_hint_match"));
+ assertThat(actual.get().getLoadBalancerName(), is("random"));
+ }
+
+ @Test
+ public void assertFindMatchedStrategyRuleWhenSQLHintNotMatch() {
+ TrafficRule trafficRule = new TrafficRule(createTrafficRuleConfig());
+ Optional<TrafficStrategyRule> actual = trafficRule.findMatchedStrategyRule(createLogicSQL(false));
+ assertFalse(actual.isPresent());
+ }
+
+ @Test
+ public void assertFindLoadBalancer() {
+ TrafficRule trafficRule = new TrafficRule(createTrafficRuleConfig());
+ TrafficLoadBalanceAlgorithm actual = trafficRule.findLoadBalancer("random");
+ assertThat(actual.getType(), is("RANDOM"));
+ }
+
+ @Test
+ public void assertGetLabels() {
+ TrafficRule trafficRule = new TrafficRule(createTrafficRuleConfig());
+ Collection<String> actual = trafficRule.getLabels();
+ assertThat(actual, is(Arrays.asList("OLTP", "OLAP")));
+ }
+
+ @SuppressWarnings({"rawtypes", "unchecked"})
+ private LogicSQL createLogicSQL(final boolean includeComments) {
+ LogicSQL result = mock(LogicSQL.class);
+ MySQLSelectStatement sqlStatement = mock(MySQLSelectStatement.class);
+ Collection<CommentSegment> comments = includeComments ? Collections.singletonList(
+ new CommentSegment("/* ShardingSphere hint: traffic=true */", 0, 0)) : Collections.emptyList();
+ when(sqlStatement.getCommentSegments()).thenReturn(comments);
+ when(sqlStatement.getProjections()).thenReturn(new ProjectionsSegment(0, 0));
+ SQLStatementContext statementContext = new SelectStatementContext(createMetaDataMap(), Collections.emptyList(), sqlStatement, "sharding_db");
+ when(result.getSqlStatementContext()).thenReturn(statementContext);
+ return result;
+ }
+
+ private Map<String, ShardingSphereMetaData> createMetaDataMap() {
+ Map<String, ShardingSphereMetaData> result = new HashMap<>(1, 1);
+ result.put("sharding_db", mock(ShardingSphereMetaData.class));
+ return result;
+ }
+
+ private TrafficRuleConfiguration createTrafficRuleConfig() {
+ TrafficRuleConfiguration result = new TrafficRuleConfiguration();
+ result.getTrafficStrategies().add(new TrafficStrategyConfiguration("sql_hint_traffic", Arrays.asList("OLTP", "OLAP"), "sql_hint_match", "random"));
+ result.getTrafficAlgorithms().put("sql_hint_match", createTrafficAlgorithm());
+ result.getLoadBalancers().put("random", createLoadBalancer());
+ return result;
+ }
+
+ private ShardingSphereAlgorithmConfiguration createTrafficAlgorithm() {
+ ShardingSphereAlgorithmConfiguration result = mock(ShardingSphereAlgorithmConfiguration.class);
+ when(result.getType()).thenReturn("SQL_HINT");
+ Properties props = new Properties();
+ props.put("traffic", true);
+ when(result.getProps()).thenReturn(props);
+ return result;
+ }
+
+ private ShardingSphereAlgorithmConfiguration createLoadBalancer() {
+ ShardingSphereAlgorithmConfiguration result = mock(ShardingSphereAlgorithmConfiguration.class);
+ when(result.getType()).thenReturn("RANDOM");
+ return result;
+ }
}