1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18 package org.apache.shardingsphere.encrypt.rewrite.token.generator;
19
20 import lombok.Setter;
21 import org.apache.shardingsphere.encrypt.exception.syntax.UnsupportedEncryptSQLException;
22 import org.apache.shardingsphere.encrypt.rewrite.aware.EncryptRuleAware;
23 import org.apache.shardingsphere.encrypt.rule.EncryptRule;
24 import org.apache.shardingsphere.encrypt.rule.EncryptTable;
25 import org.apache.shardingsphere.infra.binder.context.segment.select.orderby.OrderByItem;
26 import org.apache.shardingsphere.infra.binder.context.statement.SQLStatementContext;
27 import org.apache.shardingsphere.infra.binder.context.statement.dml.SelectStatementContext;
28 import org.apache.shardingsphere.infra.database.core.type.DatabaseTypeRegistry;
29 import org.apache.shardingsphere.infra.exception.core.ShardingSpherePreconditions;
30 import org.apache.shardingsphere.infra.metadata.database.schema.model.ShardingSphereSchema;
31 import org.apache.shardingsphere.infra.rewrite.sql.token.generator.CollectionSQLTokenGenerator;
32 import org.apache.shardingsphere.infra.rewrite.sql.token.generator.aware.SchemaMetaDataAware;
33 import org.apache.shardingsphere.infra.rewrite.sql.token.pojo.SQLToken;
34 import org.apache.shardingsphere.infra.rewrite.sql.token.pojo.generic.SubstitutableColumnNameToken;
35 import org.apache.shardingsphere.sql.parser.sql.common.segment.dml.column.ColumnSegment;
36 import org.apache.shardingsphere.sql.parser.sql.common.segment.dml.order.item.ColumnOrderByItemSegment;
37
38 import java.util.Collection;
39 import java.util.Collections;
40 import java.util.LinkedHashSet;
41 import java.util.LinkedList;
42 import java.util.Map;
43 import java.util.Optional;
44
45
46
47
48 @Setter
49 public final class EncryptOrderByItemTokenGenerator implements CollectionSQLTokenGenerator<SQLStatementContext>, SchemaMetaDataAware, EncryptRuleAware {
50
51 private String databaseName;
52
53 private Map<String, ShardingSphereSchema> schemas;
54
55 private EncryptRule encryptRule;
56
57 @Override
58 public boolean isGenerateSQLToken(final SQLStatementContext sqlStatementContext) {
59 return sqlStatementContext instanceof SelectStatementContext && containsOrderByItem(sqlStatementContext);
60 }
61
62 @Override
63 public Collection<SQLToken> generateSQLTokens(final SQLStatementContext sqlStatementContext) {
64 Collection<SQLToken> result = new LinkedHashSet<>();
65 String defaultSchema = new DatabaseTypeRegistry(sqlStatementContext.getDatabaseType()).getDefaultSchemaName(databaseName);
66 ShardingSphereSchema schema = sqlStatementContext.getTablesContext().getSchemaName().map(schemas::get).orElseGet(() -> schemas.get(defaultSchema));
67 for (OrderByItem each : getOrderByItems(sqlStatementContext)) {
68 if (each.getSegment() instanceof ColumnOrderByItemSegment) {
69 ColumnSegment columnSegment = ((ColumnOrderByItemSegment) each.getSegment()).getColumn();
70 Map<String, String> columnTableNames = sqlStatementContext.getTablesContext().findTableNamesByColumnSegment(Collections.singleton(columnSegment), schema);
71 result.addAll(generateSQLTokensWithColumnSegments(Collections.singleton(columnSegment), columnTableNames));
72 }
73 }
74 return result;
75 }
76
77 private Collection<SubstitutableColumnNameToken> generateSQLTokensWithColumnSegments(final Collection<ColumnSegment> columnSegments, final Map<String, String> columnTableNames) {
78 Collection<SubstitutableColumnNameToken> result = new LinkedList<>();
79 for (ColumnSegment each : columnSegments) {
80 String tableName = columnTableNames.getOrDefault(each.getExpression(), "");
81 Optional<EncryptTable> encryptTable = encryptRule.findEncryptTable(tableName);
82 String columnName = each.getIdentifier().getValue();
83 ShardingSpherePreconditions.checkState(!encryptTable.isPresent() || !encryptTable.get().isEncryptColumn(columnName), () -> new UnsupportedEncryptSQLException("ORDER BY"));
84 }
85 return result;
86 }
87
88 private Collection<OrderByItem> getOrderByItems(final SQLStatementContext sqlStatementContext) {
89 if (!(sqlStatementContext instanceof SelectStatementContext)) {
90 return Collections.emptyList();
91 }
92 Collection<OrderByItem> result = new LinkedList<>();
93 SelectStatementContext statementContext = (SelectStatementContext) sqlStatementContext;
94 if (!statementContext.getOrderByContext().isGenerated()) {
95 result.addAll(statementContext.getOrderByContext().getItems());
96 }
97 for (SelectStatementContext each : statementContext.getSubqueryContexts().values()) {
98 result.addAll(getOrderByItems(each));
99 }
100 return result;
101 }
102
103 private boolean containsOrderByItem(final SQLStatementContext sqlStatementContext) {
104 if (!(sqlStatementContext instanceof SelectStatementContext)) {
105 return false;
106 }
107 SelectStatementContext statementContext = (SelectStatementContext) sqlStatementContext;
108 if (!statementContext.getOrderByContext().getItems().isEmpty() && !statementContext.getOrderByContext().isGenerated()) {
109 return true;
110 }
111 for (SelectStatementContext each : statementContext.getSubqueryContexts().values()) {
112 if (containsOrderByItem(each)) {
113 return true;
114 }
115 }
116 return false;
117 }
118 }