Merge c4b719b into b8a5f4c

sharding-core/sharding-core-common/src/main/java/org/apache/shardingsphere/core/rule/EncryptRule.java

@@ -114,6 +114,19 @@ public String getLogicColumn(final String logicTable, final String cipherColumn)
         return tables.get(logicTable).getLogicColumn(cipherColumn);
     }
 
+    /**
+     * Get logic columns.
+     *
+     * @param logicTable logic table
+     * @return logic columns
+     */
+    public Collection<String> getLogicColumns(final String logicTable) {
+        if (!tables.containsKey(logicTable)) {
+            return Collections.emptyList();
+        }
+        return tables.get(logicTable).getLogicColumns();
+    }
+
     /**
      * Get plain column.
      *
@@ -183,7 +196,13 @@ public Optional<String> getAssistedQueryColumn(final String logicTable, final St
         return tables.get(logicTable).getAssistedQueryColumn(logicColumn);
     }
 
-    private Collection<String> getAssistedQueryColumns(final String logicTable) {
+    /**
+     * Get assisted query columns.
+     * 
+     * @param logicTable logic table
+     * @return assisted query columns
+     */
+    public Collection<String> getAssistedQueryColumns(final String logicTable) {
         if (!tables.containsKey(logicTable)) {
             return Collections.emptyList();
         }

sharding-core/sharding-core-common/src/main/java/org/apache/shardingsphere/core/strategy/encrypt/EncryptTable.java

@@ -64,6 +64,15 @@ public String getLogicColumn(final String cipherColumn) {
         throw new ShardingException("Can not find logic column by %s.", cipherColumn);
     }
 
+    /**
+     * Get logic columns.
+     *
+     * @return logic column
+     */
+    public Collection<String> getLogicColumns() {
+        return columns.keySet();        
+    }
+
     /**
      * Get plain column.
      * 

sharding-core/sharding-core-execute/src/main/java/org/apache/shardingsphere/core/execute/sql/execute/result/QueryResultMetaData.java

@@ -37,40 +37,40 @@
  */
 public final class QueryResultMetaData {
 
-    private final Map<String, Integer> columnLabelAndIndexes;
-
     private final ResultSetMetaData resultSetMetaData;
 
     private final ShardingRule shardingRule;
 
     private final EncryptRule encryptRule;
 
+    private final Map<String, Integer> columnLabelAndIndexes;
+
     @SneakyThrows
     public QueryResultMetaData(final ResultSetMetaData resultSetMetaData, final ShardingRule shardingRule) {
-        columnLabelAndIndexes = getColumnLabelAndIndexMap(resultSetMetaData);
         this.resultSetMetaData = resultSetMetaData;
         this.shardingRule = shardingRule;
         this.encryptRule = shardingRule.getEncryptRule();
+        columnLabelAndIndexes = getColumnLabelAndIndexMap();
     }
 
     @SneakyThrows
     public QueryResultMetaData(final ResultSetMetaData resultSetMetaData, final EncryptRule encryptRule) {
-        columnLabelAndIndexes = getColumnLabelAndIndexMap(resultSetMetaData);
         this.resultSetMetaData = resultSetMetaData;
         this.shardingRule = null;
         this.encryptRule = encryptRule;
+        columnLabelAndIndexes = getColumnLabelAndIndexMap();
     }
 
     @SneakyThrows
     public QueryResultMetaData(final ResultSetMetaData resultSetMetaData) {
-        columnLabelAndIndexes = getColumnLabelAndIndexMap(resultSetMetaData);
         this.resultSetMetaData = resultSetMetaData;
         this.shardingRule = null;
         this.encryptRule = new EncryptRule();
+        columnLabelAndIndexes = getColumnLabelAndIndexMap();
     }
 
     @SneakyThrows
-    private Map<String, Integer> getColumnLabelAndIndexMap(final ResultSetMetaData resultSetMetaData) {
+    private Map<String, Integer> getColumnLabelAndIndexMap() {
         Map<String, Integer> result = new TreeMap<>(String.CASE_INSENSITIVE_ORDER);
         for (int columnIndex = resultSetMetaData.getColumnCount(); columnIndex > 0; columnIndex--) {
             result.put(resultSetMetaData.getColumnLabel(columnIndex), columnIndex);

sharding-core/sharding-core-rewrite/src/main/java/org/apache/shardingsphere/core/rewrite/token/EncryptTokenGenerateEngine.java

@@ -17,12 +17,13 @@
 
 package org.apache.shardingsphere.core.rewrite.token;
 
-import org.apache.shardingsphere.core.rewrite.token.generator.InsertQueryAndPlainNamesTokenGenerator;
 import org.apache.shardingsphere.core.rewrite.token.generator.InsertCipherNameTokenGenerator;
-import org.apache.shardingsphere.core.rewrite.token.generator.InsertSetQueryAndPlainColumnsTokenGenerator;
+import org.apache.shardingsphere.core.rewrite.token.generator.InsertQueryAndPlainNamesTokenGenerator;
 import org.apache.shardingsphere.core.rewrite.token.generator.InsertSetCipherColumnTokenGenerator;
+import org.apache.shardingsphere.core.rewrite.token.generator.InsertSetQueryAndPlainColumnsTokenGenerator;
 import org.apache.shardingsphere.core.rewrite.token.generator.InsertValuesTokenGenerator;
 import org.apache.shardingsphere.core.rewrite.token.generator.SQLTokenGenerator;
+import org.apache.shardingsphere.core.rewrite.token.generator.SelectCipherItemTokenGenerator;
 import org.apache.shardingsphere.core.rewrite.token.generator.UpdateEncryptColumnTokenGenerator;
 import org.apache.shardingsphere.core.rewrite.token.generator.WhereEncryptColumnTokenGenerator;
 import org.apache.shardingsphere.core.rule.EncryptRule;
@@ -40,6 +41,7 @@ public final class EncryptTokenGenerateEngine extends SQLTokenGenerateEngine<Enc
     private static final Collection<SQLTokenGenerator> SQL_TOKEN_GENERATORS = new LinkedList<>();
 
     static {
+        SQL_TOKEN_GENERATORS.add(new SelectCipherItemTokenGenerator());
         SQL_TOKEN_GENERATORS.add(new UpdateEncryptColumnTokenGenerator());
         SQL_TOKEN_GENERATORS.add(new WhereEncryptColumnTokenGenerator());
         SQL_TOKEN_GENERATORS.add(new InsertCipherNameTokenGenerator());

sharding-core/sharding-core-rewrite/src/main/java/org/apache/shardingsphere/core/rewrite/token/generator/SelectCipherItemTokenGenerator.java

@@ -0,0 +1,93 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.shardingsphere.core.rewrite.token.generator;
+
+import com.google.common.base.Optional;
+import org.apache.shardingsphere.core.optimize.api.statement.OptimizedStatement;
+import org.apache.shardingsphere.core.parse.sql.segment.dml.item.ColumnSelectItemSegment;
+import org.apache.shardingsphere.core.parse.sql.segment.dml.item.SelectItemSegment;
+import org.apache.shardingsphere.core.parse.sql.segment.dml.item.SelectItemsSegment;
+import org.apache.shardingsphere.core.parse.sql.statement.dml.SelectStatement;
+import org.apache.shardingsphere.core.rewrite.builder.ParameterBuilder;
+import org.apache.shardingsphere.core.rewrite.token.pojo.SelectCipherItemToken;
+import org.apache.shardingsphere.core.rule.EncryptRule;
+
+import java.util.Collection;
+import java.util.Collections;
+import java.util.LinkedList;
+
+/**
+ * Select cipher item token generator.
+ *
+ * @author panjuan
+ */
+public final class SelectCipherItemTokenGenerator implements CollectionSQLTokenGenerator<EncryptRule> {
+
+    private EncryptRule encryptRule;
+
+    private OptimizedStatement optimizedStatement;
+
+    @Override
+    public Collection<SelectCipherItemToken> generateSQLTokens(final OptimizedStatement optimizedStatement, 
+                                                               final ParameterBuilder parameterBuilder, final EncryptRule rule, final boolean isQueryWithCipherColumn) {
+        if (!isNeedToGenerateSQLToken(optimizedStatement)) {
+            return Collections.emptyList();
+        }
+        initParameters(rule, optimizedStatement);
+        return createSelectCipherItemTokens();
+    }
+
+    private boolean isNeedToGenerateSQLToken(final OptimizedStatement optimizedStatement) {
+        if (!isSelectStatementWithTable(optimizedStatement)) {
+            return false;
+        }
+        Optional<SelectItemsSegment> selectItemsSegment = optimizedStatement.getSQLStatement().findSQLSegment(SelectItemsSegment.class);
+        return selectItemsSegment.isPresent() && !selectItemsSegment.get().getSelectItems().isEmpty();
+    }
+
+    private boolean isSelectStatementWithTable(final OptimizedStatement optimizedStatement) {
+        return optimizedStatement.getSQLStatement() instanceof SelectStatement && !optimizedStatement.getTables().isEmpty();
+    }
+
+    private void initParameters(final EncryptRule rule, final OptimizedStatement optimizedStatement) {
+        encryptRule = rule;
+        this.optimizedStatement = optimizedStatement;
+    }
+
+    private Collection<SelectCipherItemToken> createSelectCipherItemTokens() {
+        Collection<SelectCipherItemToken> result = new LinkedList<>();
+        SelectItemsSegment selectItemsSegment = optimizedStatement.getSQLStatement().findSQLSegment(SelectItemsSegment.class).get();
+        String tableName = optimizedStatement.getTables().getSingleTableName();
+        Collection<String> logicColumns = encryptRule.getLogicColumns(tableName);
+        for (SelectItemSegment each : selectItemsSegment.getSelectItems()) {
+            if (isLogicColumn(each, logicColumns)) {
+                result.add(createSelectCipherItemToken(each, tableName));
+            }
+        }
+        return result;
+    }
+
+    private boolean isLogicColumn(final SelectItemSegment each, final Collection<String> logicColumns) {
+        return each instanceof ColumnSelectItemSegment && logicColumns.contains(((ColumnSelectItemSegment) each).getName());
+    }
+
+    private SelectCipherItemToken createSelectCipherItemToken(final SelectItemSegment each, final String tableName) {
+        return new SelectCipherItemToken(each.getStartIndex(),
+                each.getStopIndex(), encryptRule.getCipherColumn(tableName, ((ColumnSelectItemSegment) each).getName()));
+    }
+}

sharding-core/sharding-core-rewrite/src/main/java/org/apache/shardingsphere/core/rewrite/token/pojo/InsertCipherNameToken.java

@@ -20,7 +20,7 @@
 import lombok.Getter;
 
 /**
- * Insert column token.
+ * Insert cipher name token.
  *
  * @author panjuan
  */

sharding-core/sharding-core-rewrite/src/main/java/org/apache/shardingsphere/core/rewrite/token/pojo/InsertGeneratedKeyNameToken.java

@@ -21,7 +21,7 @@
 import lombok.Getter;
 
 /**
- * Insert generated key columns token.
+ * Insert generated key name token.
  *
  * @author panjuan
  */

sharding-core/sharding-core-rewrite/src/main/java/org/apache/shardingsphere/core/rewrite/token/pojo/InsertQueryAndPlainNamesToken.java

@@ -24,7 +24,7 @@
 import java.util.Collection;
 
 /**
- * Insert assisted columns token.
+ * Insert query and plain names token.
  *
  * @author panjuan
  */

sharding-core/sharding-core-rewrite/src/main/java/org/apache/shardingsphere/core/rewrite/token/pojo/InsertRegularNamesToken.java

@@ -24,7 +24,7 @@
 import java.util.Collection;
 
 /**
- * Insert columns token.
+ * Insert regular names token.
  *
  * @author panjuan
  */

sharding-core/sharding-core-rewrite/src/main/java/org/apache/shardingsphere/core/rewrite/token/pojo/InsertSetCipherColumnToken.java

@@ -24,7 +24,7 @@
 import org.apache.shardingsphere.core.parse.sql.segment.dml.expr.simple.ParameterMarkerExpressionSegment;
 
 /**
- * Insert set encrypt value token.
+ * Insert set cipher column token.
  *
  * @author panjuan
  */

sharding-core/sharding-core-rewrite/src/main/java/org/apache/shardingsphere/core/rewrite/token/pojo/InsertSetGeneratedKeyColumnToken.java

@@ -24,7 +24,7 @@
 import org.apache.shardingsphere.core.parse.sql.segment.dml.expr.simple.ParameterMarkerExpressionSegment;
 
 /**
- * Insert set add generated key token.
+ * Insert set generated key column token.
  *
  * @author panjuan
  */

sharding-core/sharding-core-rewrite/src/main/java/org/apache/shardingsphere/core/rewrite/token/pojo/InsertSetQueryAndPlainColumnsToken.java

@@ -26,7 +26,7 @@
 import java.util.List;
 
 /**
- * Insert set add item token.
+ * Insert set query and plain columns token.
  *
  * @author panjuan
  */

sharding-core/sharding-core-rewrite/src/main/java/org/apache/shardingsphere/core/rewrite/token/pojo/SelectCipherItemToken.java

@@ -0,0 +1,44 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.shardingsphere.core.rewrite.token.pojo;
+
+import lombok.Getter;
+
+/**
+ * Insert cipher item token.
+ *
+ * @author panjuan
+ */
+@Getter
+public final class SelectCipherItemToken extends SQLToken implements Substitutable {
+
+    private final int stopIndex;
+
+    private final String selectItemName;
+
+    public SelectCipherItemToken(final int startIndex, final int stopIndex, final String selectItemName) {
+        super(startIndex);
+        this.stopIndex = stopIndex;
+        this.selectItemName = selectItemName;
+    }
+
+    @Override
+    public String toString() {
+        return selectItemName;
+    }
+}

sharding-core/sharding-core-rewrite/src/test/java/org/apache/shardingsphere/core/rewrite/rewriter/EncryptSQLRewriteEngineTest.java

@@ -127,6 +127,14 @@ public void assertSelectWithoutPlaceholderWithPlainEncrypt() {
         assertThat(actual.getParameters().size(), is(0));
     }
 
+    @Test
+    public void assertSelectWithoutPlaceholderWithPlainEncryptWithLogicColumn() {
+        String sql = "SELECT col3, col4 FROM t_plain_encrypt WHERE col3 = 1 or col4 = 2";
+        SQLUnit actual = getSQLUnit(sql, Collections.emptyList(), true);
+        assertThat(actual.getSql(), is("SELECT col1, col2 FROM t_plain_encrypt WHERE col1 = 'encryptValue' or col2 = 'encryptValue'"));
+        assertThat(actual.getParameters().size(), is(0));
+    }
+
     @Test
     public void assertSelectWithPlaceholderWithQueryEncrypt() {
         String sql = "SELECT * FROM t_plain_query WHERE col3 in (?, ?) and col4 in (?, ?)";
@@ -139,6 +147,18 @@ public void assertSelectWithPlaceholderWithQueryEncrypt() {
         assertThat(actual.getParameters().get(2), is((Object) "assistedEncryptValue"));
     }
 
+    @Test
+    public void assertSelectWithPlaceholderWithQueryEncryptWithLogicColumn() {
+        String sql = "SELECT col3 as alias FROM t_plain_query WHERE col3 in (?, ?) and col4 in (?, ?)";
+        SQLUnit actual = getSQLUnit(sql, parametersOfIn, true);
+        assertThat(actual.getSql(), is("SELECT col1 as alias FROM t_plain_query WHERE query1 IN (?, ?) and query2 IN (?, ?)"));
+        assertThat(actual.getParameters().size(), is(4));
+        assertThat(actual.getParameters().get(0), is((Object) "assistedEncryptValue"));
+        assertThat(actual.getParameters().get(1), is((Object) "assistedEncryptValue"));
+        assertThat(actual.getParameters().get(2), is((Object) "assistedEncryptValue"));
+        assertThat(actual.getParameters().get(2), is((Object) "assistedEncryptValue"));
+    }
+
     @Test
     public void assertSelectWithPlaceholderWithQueryPlainEncrypt() {
         String sql = "SELECT * FROM t_plain_query WHERE col3 = ? or col4 = ?";