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.rewrite.aware.DatabaseTypeAware;
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.encrypt.rule.column.EncryptColumn;
26 import org.apache.shardingsphere.infra.binder.context.segment.select.orderby.OrderByItem;
27 import org.apache.shardingsphere.infra.binder.context.segment.select.projection.Projection;
28 import org.apache.shardingsphere.infra.binder.context.segment.select.projection.impl.ColumnProjection;
29 import org.apache.shardingsphere.infra.binder.context.statement.SQLStatementContext;
30 import org.apache.shardingsphere.infra.binder.context.statement.dml.SelectStatementContext;
31 import org.apache.shardingsphere.infra.database.core.metadata.database.enums.QuoteCharacter;
32 import org.apache.shardingsphere.infra.database.core.type.DatabaseType;
33 import org.apache.shardingsphere.infra.database.core.type.DatabaseTypeRegistry;
34 import org.apache.shardingsphere.infra.metadata.database.schema.model.ShardingSphereSchema;
35 import org.apache.shardingsphere.infra.rewrite.sql.token.generator.CollectionSQLTokenGenerator;
36 import org.apache.shardingsphere.infra.rewrite.sql.token.generator.aware.SchemaMetaDataAware;
37 import org.apache.shardingsphere.infra.rewrite.sql.token.pojo.SQLToken;
38 import org.apache.shardingsphere.infra.rewrite.sql.token.pojo.generic.SubstitutableColumnNameToken;
39 import org.apache.shardingsphere.sql.parser.sql.common.segment.dml.column.ColumnSegment;
40 import org.apache.shardingsphere.sql.parser.sql.common.segment.dml.order.item.ColumnOrderByItemSegment;
41 import org.apache.shardingsphere.sql.parser.sql.common.value.identifier.IdentifierValue;
42
43 import java.util.Collection;
44 import java.util.Collections;
45 import java.util.LinkedHashSet;
46 import java.util.LinkedList;
47 import java.util.Map;
48 import java.util.Optional;
49
50
51
52
53 @Setter
54 public final class EncryptGroupByItemTokenGenerator implements CollectionSQLTokenGenerator<SQLStatementContext>, SchemaMetaDataAware, EncryptRuleAware, DatabaseTypeAware {
55
56 private String databaseName;
57
58 private Map<String, ShardingSphereSchema> schemas;
59
60 private EncryptRule encryptRule;
61
62 private DatabaseType databaseType;
63
64 @Override
65 public boolean isGenerateSQLToken(final SQLStatementContext sqlStatementContext) {
66 return sqlStatementContext instanceof SelectStatementContext && containsGroupByItem(sqlStatementContext);
67 }
68
69 @Override
70 public Collection<SQLToken> generateSQLTokens(final SQLStatementContext sqlStatementContext) {
71 Collection<SQLToken> result = new LinkedHashSet<>();
72 String defaultSchema = new DatabaseTypeRegistry(sqlStatementContext.getDatabaseType()).getDefaultSchemaName(databaseName);
73 ShardingSphereSchema schema = sqlStatementContext.getTablesContext().getSchemaName().map(schemas::get).orElseGet(() -> schemas.get(defaultSchema));
74 for (OrderByItem each : getGroupByItems(sqlStatementContext)) {
75 if (each.getSegment() instanceof ColumnOrderByItemSegment) {
76 ColumnSegment columnSegment = ((ColumnOrderByItemSegment) each.getSegment()).getColumn();
77 Map<String, String> columnTableNames = sqlStatementContext.getTablesContext().findTableNamesByColumnSegment(Collections.singleton(columnSegment), schema);
78 result.addAll(generateSQLTokensWithColumnSegments(Collections.singleton(columnSegment), columnTableNames));
79 }
80 }
81 return result;
82 }
83
84 private Collection<SubstitutableColumnNameToken> generateSQLTokensWithColumnSegments(final Collection<ColumnSegment> columnSegments, final Map<String, String> columnTableNames) {
85 Collection<SubstitutableColumnNameToken> result = new LinkedList<>();
86 for (ColumnSegment each : columnSegments) {
87 String tableName = columnTableNames.getOrDefault(each.getExpression(), "");
88 Optional<EncryptTable> encryptTable = encryptRule.findEncryptTable(tableName);
89 String columnName = each.getIdentifier().getValue();
90 if (!encryptTable.isPresent() || !encryptTable.get().isEncryptColumn(columnName)) {
91 continue;
92 }
93 int startIndex = each.getOwner().isPresent() ? each.getOwner().get().getStopIndex() + 2 : each.getStartIndex();
94 int stopIndex = each.getStopIndex();
95 EncryptColumn encryptColumn = encryptTable.get().getEncryptColumn(columnName);
96 SubstitutableColumnNameToken encryptColumnNameToken = encryptColumn.getAssistedQuery()
97 .map(optional -> new SubstitutableColumnNameToken(startIndex, stopIndex, createColumnProjections(optional.getName(), each.getIdentifier().getQuoteCharacter()), databaseType))
98 .orElseGet(() -> new SubstitutableColumnNameToken(startIndex, stopIndex, createColumnProjections(encryptColumn.getCipher().getName(), each.getIdentifier().getQuoteCharacter()),
99 databaseType));
100 result.add(encryptColumnNameToken);
101 }
102 return result;
103 }
104
105 private Collection<OrderByItem> getGroupByItems(final SQLStatementContext sqlStatementContext) {
106 if (!(sqlStatementContext instanceof SelectStatementContext)) {
107 return Collections.emptyList();
108 }
109 SelectStatementContext statementContext = (SelectStatementContext) sqlStatementContext;
110 Collection<OrderByItem> result = new LinkedList<>(statementContext.getGroupByContext().getItems());
111 for (SelectStatementContext each : statementContext.getSubqueryContexts().values()) {
112 result.addAll(getGroupByItems(each));
113 }
114 return result;
115 }
116
117 private boolean containsGroupByItem(final SQLStatementContext sqlStatementContext) {
118 if (!(sqlStatementContext instanceof SelectStatementContext)) {
119 return false;
120 }
121 SelectStatementContext statementContext = (SelectStatementContext) sqlStatementContext;
122 if (!statementContext.getGroupByContext().getItems().isEmpty()) {
123 return true;
124 }
125 for (SelectStatementContext each : statementContext.getSubqueryContexts().values()) {
126 if (containsGroupByItem(each)) {
127 return true;
128 }
129 }
130 return false;
131 }
132
133 private Collection<Projection> createColumnProjections(final String columnName, final QuoteCharacter quoteCharacter) {
134 return Collections.singleton(new ColumnProjection(null, new IdentifierValue(columnName, quoteCharacter), null, databaseType));
135 }
136 }