diff --git a/fe/fe-core/src/main/antlr4/org/apache/doris/nereids/DorisLexer.g4 b/fe/fe-core/src/main/antlr4/org/apache/doris/nereids/DorisLexer.g4 index 95c21ff35224f6..98bcf27b8f4f1b 100644 --- a/fe/fe-core/src/main/antlr4/org/apache/doris/nereids/DorisLexer.g4 +++ b/fe/fe-core/src/main/antlr4/org/apache/doris/nereids/DorisLexer.g4 @@ -225,6 +225,7 @@ IS_NULL_PRED: 'IS_NULL_PRED'; IS_NOT_NULL_PRED: 'IS_NOT_NULL_PRED'; ITEMS: 'ITEMS'; JOIN: 'JOIN'; +KEY: 'KEY'; KEYS: 'KEYS'; LABEL: 'LABEL'; LAST: 'LAST'; diff --git a/fe/fe-core/src/main/antlr4/org/apache/doris/nereids/DorisParser.g4 b/fe/fe-core/src/main/antlr4/org/apache/doris/nereids/DorisParser.g4 index cd178cd22039f6..5f4af17fbd785d 100644 --- a/fe/fe-core/src/main/antlr4/org/apache/doris/nereids/DorisParser.g4 +++ b/fe/fe-core/src/main/antlr4/org/apache/doris/nereids/DorisParser.g4 @@ -385,6 +385,7 @@ primaryExpression | identifier #columnReference | base=primaryExpression DOT fieldName=identifier #dereference | LEFT_PAREN expression RIGHT_PAREN #parenthesizedExpression + | KEY (dbName=identifier DOT)? keyName=identifier #encryptKey | EXTRACT LEFT_PAREN field=identifier FROM (DATE | TIMESTAMP)? source=valueExpression RIGHT_PAREN #extract ; diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/parser/LogicalPlanBuilder.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/parser/LogicalPlanBuilder.java index f0a270bdb61b53..95528f9e1e7331 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/parser/LogicalPlanBuilder.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/parser/LogicalPlanBuilder.java @@ -182,6 +182,7 @@ import org.apache.doris.nereids.trees.expressions.functions.scalar.DaysDiff; import org.apache.doris.nereids.trees.expressions.functions.scalar.DaysSub; import org.apache.doris.nereids.trees.expressions.functions.scalar.ElementAt; +import org.apache.doris.nereids.trees.expressions.functions.scalar.EncryptKeyRef; import org.apache.doris.nereids.trees.expressions.functions.scalar.HoursAdd; import org.apache.doris.nereids.trees.expressions.functions.scalar.HoursDiff; import org.apache.doris.nereids.trees.expressions.functions.scalar.HoursSub; @@ -214,6 +215,7 @@ import org.apache.doris.nereids.trees.expressions.literal.Literal; import org.apache.doris.nereids.trees.expressions.literal.NullLiteral; import org.apache.doris.nereids.trees.expressions.literal.SmallIntLiteral; +import org.apache.doris.nereids.trees.expressions.literal.StringLiteral; import org.apache.doris.nereids.trees.expressions.literal.TinyIntLiteral; import org.apache.doris.nereids.trees.expressions.literal.VarcharLiteral; import org.apache.doris.nereids.trees.plans.JoinHint; @@ -1109,6 +1111,13 @@ public UnboundFunction visitExtract(DorisParser.ExtractContext ctx) { }); } + @Override + public Expression visitEncryptKey(DorisParser.EncryptKeyContext ctx) { + String db = ctx.dbName == null ? "" : ctx.dbName.getText(); + String key = ctx.keyName.getText(); + return new EncryptKeyRef(new StringLiteral(db), new StringLiteral(key)); + } + @Override public Expression visitFunctionCall(DorisParser.FunctionCallContext ctx) { return ParserUtils.withOrigin(ctx, () -> { diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/expression/rules/FoldConstantRuleOnFE.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/expression/rules/FoldConstantRuleOnFE.java index 42f62db904143c..cef81acd04620d 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/expression/rules/FoldConstantRuleOnFE.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/expression/rules/FoldConstantRuleOnFE.java @@ -17,7 +17,10 @@ package org.apache.doris.nereids.rules.expression.rules; +import org.apache.doris.catalog.EncryptKey; +import org.apache.doris.catalog.Env; import org.apache.doris.cluster.ClusterNamespace; +import org.apache.doris.nereids.exceptions.AnalysisException; import org.apache.doris.nereids.rules.expression.AbstractExpressionRewriteRule; import org.apache.doris.nereids.rules.expression.ExpressionRewriteContext; import org.apache.doris.nereids.trees.expressions.AggregateExpression; @@ -51,6 +54,7 @@ import org.apache.doris.nereids.trees.expressions.functions.scalar.CurrentUser; import org.apache.doris.nereids.trees.expressions.functions.scalar.Database; import org.apache.doris.nereids.trees.expressions.functions.scalar.Date; +import org.apache.doris.nereids.trees.expressions.functions.scalar.EncryptKeyRef; import org.apache.doris.nereids.trees.expressions.functions.scalar.If; import org.apache.doris.nereids.trees.expressions.functions.scalar.Password; import org.apache.doris.nereids.trees.expressions.functions.scalar.User; @@ -70,9 +74,11 @@ import org.apache.doris.nereids.types.BooleanType; import org.apache.doris.nereids.types.DataType; import org.apache.doris.nereids.util.ExpressionUtils; +import org.apache.doris.qe.ConnectContext; import org.apache.doris.qe.GlobalVariable; import com.google.common.base.Preconditions; +import com.google.common.base.Strings; import com.google.common.collect.Lists; import org.apache.commons.codec.digest.DigestUtils; @@ -111,6 +117,29 @@ public Expression visitLiteral(Literal literal, ExpressionRewriteContext context return literal; } + @Override + public Expression visitEncryptKeyRef(EncryptKeyRef encryptKeyRef, ExpressionRewriteContext context) { + String dbName = encryptKeyRef.getDbName(); + ConnectContext connectContext = ConnectContext.get(); + if (Strings.isNullOrEmpty(dbName)) { + dbName = connectContext.getDatabase(); + } + if ("".equals(dbName)) { + throw new AnalysisException("DB " + dbName + "not found"); + } + dbName = ClusterNamespace.getFullName(connectContext.getClusterName(), dbName); + org.apache.doris.catalog.Database database = + Env.getCurrentEnv().getInternalCatalog().getDbNullable(dbName); + if (database == null) { + throw new AnalysisException("DB " + dbName + "not found"); + } + EncryptKey encryptKey = database.getEncryptKey(encryptKeyRef.getEncryptKeyName()); + if (encryptKey == null) { + throw new AnalysisException("Can not found encryptKey" + encryptKeyRef.getEncryptKeyName()); + } + return new StringLiteral(encryptKey.getKeyString()); + } + @Override public Expression visitEqualTo(EqualTo equalTo, ExpressionRewriteContext context) { equalTo = rewriteChildren(equalTo, context); diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/EncryptKeyRef.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/EncryptKeyRef.java new file mode 100644 index 00000000000000..7bbe59cfaf0709 --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/EncryptKeyRef.java @@ -0,0 +1,84 @@ +// 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.doris.nereids.trees.expressions.functions.scalar; + +import org.apache.doris.catalog.FunctionSignature; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.functions.AlwaysNotNullable; +import org.apache.doris.nereids.trees.expressions.functions.ExplicitlyCastableSignature; +import org.apache.doris.nereids.trees.expressions.literal.StringLikeLiteral; +import org.apache.doris.nereids.trees.expressions.shape.BinaryExpression; +import org.apache.doris.nereids.trees.expressions.visitor.ExpressionVisitor; +import org.apache.doris.nereids.types.StringType; + +import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableList; +import org.apache.hadoop.util.Lists; + +import java.util.List; + +/** + * ScalarFunction 'EncryptKeyRef'. + */ +public class EncryptKeyRef extends ScalarFunction + implements BinaryExpression, ExplicitlyCastableSignature, AlwaysNotNullable { + + public static final List SIGNATURES = ImmutableList.of( + FunctionSignature.ret(StringType.INSTANCE).args(StringType.INSTANCE, StringType.INSTANCE) + ); + + /** + * constructor with 2 arguments. + */ + public EncryptKeyRef(Expression arg0, Expression arg1) { + super("encryptKeyRef", Lists.newArrayList(arg0, arg1)); + } + + public EncryptKeyRef(List args) { + super("encryptKeyRef", args); + } + + public String getDbName() { + Preconditions.checkArgument(children.get(0) instanceof StringLikeLiteral); + return ((StringLikeLiteral) children.get(0)).value; + } + + public String getEncryptKeyName() { + Preconditions.checkArgument(children.get(1) instanceof StringLikeLiteral); + return ((StringLikeLiteral) children.get(1)).value; + } + + /** + * withChildren. + */ + @Override + public EncryptKeyRef withChildren(List children) { + Preconditions.checkArgument(children.size() == 2); + return new EncryptKeyRef(children); + } + + @Override + public List getSignatures() { + return SIGNATURES; + } + + @Override + public R accept(ExpressionVisitor visitor, C context) { + return visitor.visitEncryptKeyRef(this, context); + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/visitor/ScalarFunctionVisitor.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/visitor/ScalarFunctionVisitor.java index c228430f35d29c..f77d241898fe63 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/visitor/ScalarFunctionVisitor.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/visitor/ScalarFunctionVisitor.java @@ -131,6 +131,7 @@ import org.apache.doris.nereids.trees.expressions.functions.scalar.E; import org.apache.doris.nereids.trees.expressions.functions.scalar.ElementAt; import org.apache.doris.nereids.trees.expressions.functions.scalar.Elt; +import org.apache.doris.nereids.trees.expressions.functions.scalar.EncryptKeyRef; import org.apache.doris.nereids.trees.expressions.functions.scalar.EndsWith; import org.apache.doris.nereids.trees.expressions.functions.scalar.EsQuery; import org.apache.doris.nereids.trees.expressions.functions.scalar.Exp; @@ -841,6 +842,10 @@ default R visitEndsWith(EndsWith endsWith, C context) { return visitScalarFunction(endsWith, context); } + default R visitEncryptKeyRef(EncryptKeyRef encryptKeyRef, C context) { + return visitScalarFunction(encryptKeyRef, context); + } + default R visitEsQuery(EsQuery esQuery, C context) { return visitScalarFunction(esQuery, context); } diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/parser/NereidsParserTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/parser/NereidsParserTest.java index fd4d1f2e1287f1..c7ecd44e920c72 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/parser/NereidsParserTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/parser/NereidsParserTest.java @@ -99,7 +99,7 @@ public void testParseCTE() { Assertions.assertEquals(PlanType.LOGICAL_CTE, logicalPlan.getType()); Assertions.assertEquals(((LogicalCTE) logicalPlan).getAliasQueries().size(), 2); - String cteSql3 = "with t1 (key, name) as (select s_suppkey, s_name from supplier) select * from t1"; + String cteSql3 = "with t1 (keyy, name) as (select s_suppkey, s_name from supplier) select * from t1"; logicalPlan = (LogicalPlan) nereidsParser.parseSingle(cteSql3).child(0); Assertions.assertEquals(PlanType.LOGICAL_CTE, logicalPlan.getType()); Assertions.assertEquals(((LogicalCTE) logicalPlan).getAliasQueries().size(), 1); @@ -290,43 +290,43 @@ public void parseSetOperation() { @Test public void testJoinHint() { // no hint - parsePlan("select * from t1 join t2 on t1.key=t2.key") + parsePlan("select * from t1 join t2 on t1.keyy=t2.keyy") .matches(logicalJoin().when(j -> j.getHint() == JoinHint.NONE)); // valid hint - parsePlan("select * from t1 join [shuffle] t2 on t1.key=t2.key") + parsePlan("select * from t1 join [shuffle] t2 on t1.keyy=t2.keyy") .matches(logicalJoin().when(j -> j.getHint() == JoinHint.SHUFFLE_RIGHT)); - parsePlan("select * from t1 join [ shuffle ] t2 on t1.key=t2.key") + parsePlan("select * from t1 join [ shuffle ] t2 on t1.keyy=t2.keyy") .matches(logicalJoin().when(j -> j.getHint() == JoinHint.SHUFFLE_RIGHT)); - parsePlan("select * from t1 join [broadcast] t2 on t1.key=t2.key") + parsePlan("select * from t1 join [broadcast] t2 on t1.keyy=t2.keyy") .matches(logicalJoin().when(j -> j.getHint() == JoinHint.BROADCAST_RIGHT)); - parsePlan("select * from t1 join /*+ broadcast */ t2 on t1.key=t2.key") + parsePlan("select * from t1 join /*+ broadcast */ t2 on t1.keyy=t2.keyy") .matches(logicalJoin().when(j -> j.getHint() == JoinHint.BROADCAST_RIGHT)); // invalid hint position - parsePlan("select * from [shuffle] t1 join t2 on t1.key=t2.key") + parsePlan("select * from [shuffle] t1 join t2 on t1.keyy=t2.keyy") .assertThrowsExactly(ParseException.class); - parsePlan("select * from /*+ shuffle */ t1 join t2 on t1.key=t2.key") + parsePlan("select * from /*+ shuffle */ t1 join t2 on t1.keyy=t2.keyy") .assertThrowsExactly(ParseException.class); // invalid hint content - parsePlan("select * from t1 join [bucket] t2 on t1.key=t2.key") + parsePlan("select * from t1 join [bucket] t2 on t1.keyy=t2.keyy") .assertThrowsExactly(ParseException.class) .assertMessageContains("Invalid join hint: bucket(line 1, pos 22)\n" + "\n" + "== SQL ==\n" - + "select * from t1 join [bucket] t2 on t1.key=t2.key\n" + + "select * from t1 join [bucket] t2 on t1.keyy=t2.keyy\n" + "----------------------^^^"); // invalid multiple hints - parsePlan("select * from t1 join /*+ shuffle , broadcast */ t2 on t1.key=t2.key") + parsePlan("select * from t1 join /*+ shuffle , broadcast */ t2 on t1.keyy=t2.keyy") .assertThrowsExactly(ParseException.class); - parsePlan("select * from t1 join [shuffle,broadcast] t2 on t1.key=t2.key") + parsePlan("select * from t1 join [shuffle,broadcast] t2 on t1.keyy=t2.keyy") .assertThrowsExactly(ParseException.class); } diff --git a/regression-test/data/nereids_function_p0/scalar_function/E.out b/regression-test/data/nereids_function_p0/scalar_function/E.out index ccaa576f128d6f..f543bd22d9be7d 100644 --- a/regression-test/data/nereids_function_p0/scalar_function/E.out +++ b/regression-test/data/nereids_function_p0/scalar_function/E.out @@ -173,3 +173,6 @@ true +-- !sql_encryptkey -- +ABCD123456789 ABCD123456789 + diff --git a/regression-test/suites/nereids_function_p0/scalar_function/E.groovy b/regression-test/suites/nereids_function_p0/scalar_function/E.groovy index e55df2d2b9c5c5..8962c1c795cfd4 100644 --- a/regression-test/suites/nereids_function_p0/scalar_function/E.groovy +++ b/regression-test/suites/nereids_function_p0/scalar_function/E.groovy @@ -19,6 +19,10 @@ suite("nereids_scalar_fn_E") { sql 'use regression_test_nereids_function_p0' sql 'set enable_nereids_planner=true' sql 'set enable_fallback_to_original_planner=false' + + sql """ + CREATE ENCRYPTKEY if not exists my_key AS "ABCD123456789"; + """ qt_sql_elt_Integer_Varchar "select elt(kint, kvchrs1) from fn_test order by kint, kvchrs1" qt_sql_elt_Integer_Varchar_notnull "select elt(kint, kvchrs1) from fn_test_not_nullable order by kint, kvchrs1" qt_sql_elt_Integer_String "select elt(kint, kstr) from fn_test order by kint, kstr" @@ -31,4 +35,6 @@ suite("nereids_scalar_fn_E") { qt_sql_exp_Double_notnull "select exp(kdbl) from fn_test_not_nullable order by kdbl" qt_sql_extract_url_parameter_Varchar_Varchar "select extract_url_parameter(kvchrs1, kvchrs1) from fn_test order by kvchrs1, kvchrs1" qt_sql_extract_url_parameter_Varchar_Varchar_notnull "select extract_url_parameter(kvchrs1, kvchrs1) from fn_test_not_nullable order by kvchrs1, kvchrs1" + + qt_sql_encryptkey "select key my_key, key regression_test_nereids_function_p0.my_key" } \ No newline at end of file diff --git a/regression-test/suites/nereids_p0/with/test_with_and_two_phase_agg.groovy b/regression-test/suites/nereids_p0/with/test_with_and_two_phase_agg.groovy index c7b5865219a0aa..6b80546bc70185 100644 --- a/regression-test/suites/nereids_p0/with/test_with_and_two_phase_agg.groovy +++ b/regression-test/suites/nereids_p0/with/test_with_and_two_phase_agg.groovy @@ -22,21 +22,21 @@ suite("test_with_and_two_phase_agg") { sql """ DROP TABLE IF EXISTS ${tableName} """ sql """ CREATE TABLE IF NOT EXISTS ${tableName}( - `key` int not null, + `key1` int not null, `key2` varchar(50) not null, `account` varchar(50) not null ) ENGINE = OLAP - UNIQUE KEY (`key`, `key2`) - DISTRIBUTED BY HASH(`key`) + UNIQUE KEY (`key1`, `key2`) + DISTRIBUTED BY HASH(`key1`) PROPERTIES("replication_num" = "1"); """ sql """ INSERT INTO ${tableName} VALUES (1, '1332050726', '1332050726'); """ qt_select """ - WITH t2 AS( SELECT sum(`key`) num, COUNT(DISTINCT `account`) unt + WITH t2 AS( SELECT sum(`key1`) num, COUNT(DISTINCT `account`) unt FROM ${tableName}) SELECT num FROM t2; """ qt_select2 """ - WITH t2 AS( SELECT `key2`, sum(`key`) num, COUNT(DISTINCT `account`) unt + WITH t2 AS( SELECT `key2`, sum(`key1`) num, COUNT(DISTINCT `account`) unt FROM ${tableName} GROUP BY `key2`) SELECT num FROM t2; """ }