1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18 package org.apache.shardingsphere.shadow.route.engine.dml;
19
20 import lombok.AccessLevel;
21 import lombok.Getter;
22 import lombok.RequiredArgsConstructor;
23 import org.apache.shardingsphere.infra.binder.context.statement.SQLStatementContext;
24 import org.apache.shardingsphere.infra.binder.context.type.TableAvailable;
25 import org.apache.shardingsphere.infra.route.context.RouteContext;
26 import org.apache.shardingsphere.shadow.api.shadow.ShadowOperationType;
27 import org.apache.shardingsphere.shadow.api.shadow.column.ColumnShadowAlgorithm;
28 import org.apache.shardingsphere.shadow.api.shadow.hint.HintShadowAlgorithm;
29 import org.apache.shardingsphere.shadow.condition.ShadowColumnCondition;
30 import org.apache.shardingsphere.shadow.condition.ShadowDetermineCondition;
31 import org.apache.shardingsphere.shadow.route.engine.ShadowRouteEngine;
32 import org.apache.shardingsphere.shadow.route.engine.determiner.ColumnShadowAlgorithmDeterminer;
33 import org.apache.shardingsphere.shadow.route.engine.determiner.HintShadowAlgorithmDeterminer;
34 import org.apache.shardingsphere.shadow.rule.ShadowRule;
35 import org.apache.shardingsphere.shadow.spi.ShadowAlgorithm;
36 import org.apache.shardingsphere.sql.parser.sql.common.segment.generic.CommentSegment;
37 import org.apache.shardingsphere.sql.parser.sql.common.segment.generic.table.SimpleTableSegment;
38 import org.apache.shardingsphere.sql.parser.sql.common.statement.AbstractSQLStatement;
39 import org.apache.shardingsphere.sql.parser.sql.common.statement.SQLStatement;
40
41 import java.util.Collection;
42 import java.util.Collections;
43 import java.util.LinkedHashMap;
44 import java.util.Map;
45 import java.util.Optional;
46 import java.util.stream.Collectors;
47
48
49
50
51 @RequiredArgsConstructor(access = AccessLevel.PROTECTED)
52 @Getter
53 public abstract class AbstractShadowDMLStatementRouteEngine implements ShadowRouteEngine {
54
55 private final SQLStatementContext sqlStatementContext;
56
57 private final ShadowOperationType operationType;
58
59 private final Map<String, String> tableAliasNameMappings = new LinkedHashMap<>();
60
61 @Override
62 public final void route(final RouteContext routeContext, final ShadowRule rule) {
63 tableAliasNameMappings.putAll(getTableAliasNameMappings(((TableAvailable) sqlStatementContext).getAllTables()));
64 decorateRouteContext(routeContext, rule, findShadowDataSourceMappings(rule));
65 }
66
67 private Map<String, String> getTableAliasNameMappings(final Collection<SimpleTableSegment> tableSegments) {
68 Map<String, String> result = new LinkedHashMap<>(tableSegments.size(), 1F);
69 for (SimpleTableSegment each : tableSegments) {
70 String tableName = each.getTableName().getIdentifier().getValue();
71 String alias = each.getAliasName().isPresent() ? each.getAliasName().get() : tableName;
72 result.put(alias, tableName);
73 }
74 return result;
75 }
76
77 private Map<String, String> findShadowDataSourceMappings(final ShadowRule rule) {
78 Collection<String> relatedShadowTables = rule.getRelatedShadowTables(tableAliasNameMappings.values());
79 Collection<String> sqlComments = getSQLComments();
80 if (relatedShadowTables.isEmpty() && isMatchDefaultAlgorithm(rule, sqlComments)) {
81 return rule.getAllShadowDataSourceMappings();
82 }
83 Map<String, String> result = findBySQLComments(rule, sqlComments, relatedShadowTables);
84 return result.isEmpty() ? findByShadowColumn(rule, relatedShadowTables) : result;
85 }
86
87 private Collection<String> getSQLComments() {
88 SQLStatement sqlStatement = sqlStatementContext.getSqlStatement();
89 return ((AbstractSQLStatement) sqlStatement).getCommentSegments().stream().map(CommentSegment::getText).collect(Collectors.toList());
90 }
91
92 @SuppressWarnings("unchecked")
93 private boolean isMatchDefaultAlgorithm(final ShadowRule rule, final Collection<String> sqlComments) {
94 Optional<ShadowAlgorithm> defaultAlgorithm = rule.getDefaultShadowAlgorithm();
95 if (defaultAlgorithm.isPresent() && defaultAlgorithm.get() instanceof HintShadowAlgorithm<?>) {
96 ShadowDetermineCondition determineCondition = new ShadowDetermineCondition("", ShadowOperationType.HINT_MATCH);
97 return HintShadowAlgorithmDeterminer.isShadow((HintShadowAlgorithm<Comparable<?>>) defaultAlgorithm.get(), determineCondition.initSQLComments(sqlComments), rule);
98 }
99 return false;
100 }
101
102 private Map<String, String> findBySQLComments(final ShadowRule rule, final Collection<String> sqlComments, final Collection<String> relatedShadowTables) {
103 Map<String, String> result = new LinkedHashMap<>();
104 for (String each : relatedShadowTables) {
105 if (isContainsShadowInSQLComments(rule, each, sqlComments, new ShadowDetermineCondition(each, operationType))) {
106 result.putAll(rule.getRelatedShadowDataSourceMappings(each));
107 return result;
108 }
109 }
110 return result;
111 }
112
113 private boolean isContainsShadowInSQLComments(final ShadowRule rule, final String tableName, final Collection<String> sqlComments, final ShadowDetermineCondition shadowCondition) {
114 ShadowDetermineCondition shadowConditionWithComments = shadowCondition.initSQLComments(sqlComments);
115 for (HintShadowAlgorithm<Comparable<?>> each : rule.getRelatedHintShadowAlgorithms(tableName)) {
116 if (HintShadowAlgorithmDeterminer.isShadow(each, shadowConditionWithComments, rule)) {
117 return true;
118 }
119 }
120 return false;
121 }
122
123 private Map<String, String> findByShadowColumn(final ShadowRule rule, final Collection<String> relatedShadowTables) {
124 for (String each : relatedShadowTables) {
125 Collection<String> relatedShadowColumnNames = rule.getRelatedShadowColumnNames(operationType, each);
126 if (!relatedShadowColumnNames.isEmpty() && isMatchAnyColumnShadowAlgorithms(rule, each, relatedShadowColumnNames)) {
127 return rule.getRelatedShadowDataSourceMappings(each);
128 }
129 }
130 return Collections.emptyMap();
131 }
132
133 private boolean isMatchAnyColumnShadowAlgorithms(final ShadowRule rule, final String shadowTable, final Collection<String> shadowColumnNames) {
134 for (String each : shadowColumnNames) {
135 if (isMatchAnyColumnShadowAlgorithms(rule, shadowTable, each)) {
136 return true;
137 }
138 }
139 return false;
140 }
141
142 private boolean isMatchAnyColumnShadowAlgorithms(final ShadowRule rule, final String shadowTable, final String shadowColumn) {
143 Collection<ColumnShadowAlgorithm<Comparable<?>>> columnShadowAlgorithms = rule.getRelatedColumnShadowAlgorithms(operationType, shadowTable, shadowColumn);
144 if (columnShadowAlgorithms.isEmpty()) {
145 return false;
146 }
147 for (ShadowColumnCondition each : getShadowColumnConditions(shadowColumn)) {
148 if (isMatchColumnShadowAlgorithm(shadowTable, columnShadowAlgorithms, each)) {
149 return true;
150 }
151 }
152 return false;
153 }
154
155 private boolean isMatchColumnShadowAlgorithm(final String shadowTable, final Collection<ColumnShadowAlgorithm<Comparable<?>>> algorithms, final ShadowColumnCondition condition) {
156 for (ColumnShadowAlgorithm<Comparable<?>> each : algorithms) {
157 if (ColumnShadowAlgorithmDeterminer.isShadow(each, new ShadowDetermineCondition(shadowTable, operationType).initShadowColumnCondition(condition))) {
158 return true;
159 }
160 }
161 return false;
162 }
163
164 protected abstract Collection<ShadowColumnCondition> getShadowColumnConditions(String shadowColumnName);
165
166 protected final String getSingleTableName() {
167 return tableAliasNameMappings.entrySet().iterator().next().getValue();
168 }
169 }