From 4aceb0fff311f0073529dda4b46c7237d5b72c6c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=B0=A2=E5=81=A5?= Date: Mon, 13 Nov 2023 12:25:08 +0800 Subject: [PATCH] [fix](Nereids) store user variable in connect context (#26655) 1.user variable should be case insensitive 2.user variable should be cleared after the connection reset --- .../apache/doris/analysis/VariableExpr.java | 3 +- .../nereids/rules/analysis/SlotBinder.java | 2 +- .../org/apache/doris/qe/ConnectContext.java | 86 +++++++++++++++++++ .../java/org/apache/doris/qe/SetExecutor.java | 2 +- .../java/org/apache/doris/qe/VariableMgr.java | 64 -------------- .../data/query_p0/set/test_user_var.out | 6 ++ .../suites/query_p0/set/test_user_var.groovy | 5 ++ 7 files changed, 101 insertions(+), 67 deletions(-) diff --git a/fe/fe-core/src/main/java/org/apache/doris/analysis/VariableExpr.java b/fe/fe-core/src/main/java/org/apache/doris/analysis/VariableExpr.java index 8e657d6418dacf..b9b9e5d3d7a566 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/analysis/VariableExpr.java +++ b/fe/fe-core/src/main/java/org/apache/doris/analysis/VariableExpr.java @@ -21,6 +21,7 @@ import org.apache.doris.common.AnalysisException; import org.apache.doris.common.DdlException; import org.apache.doris.common.ErrorReport; +import org.apache.doris.qe.ConnectContext; import org.apache.doris.qe.VariableMgr; import org.apache.doris.qe.VariableVarConverters; import org.apache.doris.thrift.TBoolLiteral; @@ -77,7 +78,7 @@ public Expr clone() { @Override public void analyzeImpl(Analyzer analyzer) throws AnalysisException { if (setType == SetType.USER) { - VariableMgr.fillValueForUserDefinedVar(this); + ConnectContext.get().fillValueForUserDefinedVar(this); } else { VariableMgr.fillValue(analyzer.getContext().getSessionVariable(), this); if (!Strings.isNullOrEmpty(name) && VariableVarConverters.hasConverter(name)) { diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/SlotBinder.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/SlotBinder.java index 4b5b977b0aa2f8..57d84610cfd143 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/SlotBinder.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/SlotBinder.java @@ -94,7 +94,7 @@ public Expression visitUnboundVariable(UnboundVariable unboundVariable, Cascades } else if (unboundVariable.getType() == VariableType.GLOBAL) { literal = VariableMgr.getLiteral(sessionVariable, name, SetType.GLOBAL); } else if (unboundVariable.getType() == VariableType.USER) { - literal = VariableMgr.getLiteralForUserVar(name); + literal = ConnectContext.get().getLiteralForUserVar(name); } if (literal == null) { throw new AnalysisException("Unsupported system variable: " + unboundVariable.getName()); diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectContext.java b/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectContext.java index cfb41d3b819a2b..d6c0aaf820b32f 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectContext.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectContext.java @@ -17,11 +17,21 @@ package org.apache.doris.qe; +import org.apache.doris.analysis.BoolLiteral; +import org.apache.doris.analysis.DecimalLiteral; +import org.apache.doris.analysis.FloatLiteral; +import org.apache.doris.analysis.IntLiteral; +import org.apache.doris.analysis.LiteralExpr; +import org.apache.doris.analysis.NullLiteral; +import org.apache.doris.analysis.SetVar; +import org.apache.doris.analysis.StringLiteral; import org.apache.doris.analysis.UserIdentity; +import org.apache.doris.analysis.VariableExpr; import org.apache.doris.catalog.DatabaseIf; import org.apache.doris.catalog.Env; import org.apache.doris.catalog.FunctionRegistry; import org.apache.doris.catalog.TableIf; +import org.apache.doris.catalog.Type; import org.apache.doris.cluster.ClusterNamespace; import org.apache.doris.common.Config; import org.apache.doris.common.telemetry.Telemetry; @@ -37,6 +47,7 @@ import org.apache.doris.mysql.MysqlSslContext; import org.apache.doris.nereids.StatementContext; import org.apache.doris.nereids.stats.StatsErrorEstimator; +import org.apache.doris.nereids.trees.expressions.literal.Literal; import org.apache.doris.plugin.AuditEvent.AuditEventBuilder; import org.apache.doris.resource.Tag; import org.apache.doris.statistics.ColumnStatistic; @@ -61,6 +72,7 @@ import java.util.Map; import java.util.Optional; import java.util.Set; +import javax.annotation.Nullable; // When one client connect in, we create a connect context for it. // We store session information here. Meanwhile ConnectScheduler all @@ -112,6 +124,8 @@ public class ConnectContext { protected volatile UserIdentity currentUserIdentity; // Variables belong to this session. protected volatile SessionVariable sessionVariable; + // Store user variable in this connection + private Map userVars = new HashMap<>(); // Scheduler this connection belongs to protected volatile ConnectScheduler connectScheduler; // Executor @@ -162,6 +176,7 @@ public class ConnectContext { private SessionContext sessionContext; + // This context is used for SSL connection between server and mysql client. private final MysqlSslContext mysqlSslContext = new MysqlSslContext(SSL_PROTOCOL); @@ -253,6 +268,18 @@ public boolean notEvalNondeterministicFunction() { return notEvalNondeterministicFunction; } + public void init() { + state = new QueryState(); + returnRows = 0; + isKilled = false; + sessionVariable = VariableMgr.newSessionVariable(); + userVars = new HashMap<>(); + command = MysqlCommand.COM_SLEEP; + if (Config.use_fuzzy_session_variable) { + sessionVariable.initFuzzyModeVariables(); + } + } + public ConnectContext() { this(null); } @@ -378,6 +405,65 @@ public void setEnv(Env env) { defaultCatalog = env.getInternalCatalog().getName(); } + public void setUserVar(SetVar setVar) { + userVars.put(setVar.getVariable().toLowerCase(), setVar.getResult()); + } + + public @Nullable Literal getLiteralForUserVar(String varName) { + varName = varName.toLowerCase(); + if (userVars.containsKey(varName)) { + LiteralExpr literalExpr = userVars.get(varName); + if (literalExpr instanceof BoolLiteral) { + return Literal.of(((BoolLiteral) literalExpr).getValue()); + } else if (literalExpr instanceof IntLiteral) { + return Literal.of(((IntLiteral) literalExpr).getValue()); + } else if (literalExpr instanceof FloatLiteral) { + return Literal.of(((FloatLiteral) literalExpr).getValue()); + } else if (literalExpr instanceof DecimalLiteral) { + return Literal.of(((DecimalLiteral) literalExpr).getValue()); + } else if (literalExpr instanceof StringLiteral) { + return Literal.of(((StringLiteral) literalExpr).getValue()); + } else if (literalExpr instanceof NullLiteral) { + return Literal.of(null); + } else { + return Literal.of(""); + } + } else { + // If there are no such user defined var, just return the NULL value. + return Literal.of(null); + } + } + + // Get variable value through variable name, used to satisfy statement like `SELECT @@comment_version` + public void fillValueForUserDefinedVar(VariableExpr desc) { + String varName = desc.getName().toLowerCase(); + if (userVars.containsKey(varName)) { + LiteralExpr literalExpr = userVars.get(varName); + desc.setType(literalExpr.getType()); + if (literalExpr instanceof BoolLiteral) { + desc.setBoolValue(((BoolLiteral) literalExpr).getValue()); + } else if (literalExpr instanceof IntLiteral) { + desc.setIntValue(((IntLiteral) literalExpr).getValue()); + } else if (literalExpr instanceof FloatLiteral) { + desc.setFloatValue(((FloatLiteral) literalExpr).getValue()); + } else if (literalExpr instanceof DecimalLiteral) { + desc.setDecimalValue(((DecimalLiteral) literalExpr).getValue()); + } else if (literalExpr instanceof StringLiteral) { + desc.setStringValue(((StringLiteral) literalExpr).getValue()); + } else if (literalExpr instanceof NullLiteral) { + desc.setType(Type.NULL); + desc.setIsNull(); + } else { + desc.setType(Type.VARCHAR); + desc.setStringValue(""); + } + } else { + // If there are no such user defined var, just fill the NULL value. + desc.setType(Type.NULL); + desc.setIsNull(); + } + } + public Env getEnv() { return env; } diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/SetExecutor.java b/fe/fe-core/src/main/java/org/apache/doris/qe/SetExecutor.java index 9dbe9ca8a988aa..673f01562cef5e 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/SetExecutor.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/SetExecutor.java @@ -56,7 +56,7 @@ private void setVariable(SetVar var) throws DdlException { // do nothing return; } else if (var instanceof SetUserDefinedVar) { - VariableMgr.setUserVar(var); + ConnectContext.get().setUserVar(var); } else { VariableMgr.setVar(ctx.getSessionVariable(), var); } diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/VariableMgr.java b/fe/fe-core/src/main/java/org/apache/doris/qe/VariableMgr.java index 80cf288a32af1e..c729f0a444c038 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/VariableMgr.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/VariableMgr.java @@ -17,15 +17,9 @@ package org.apache.doris.qe; -import org.apache.doris.analysis.BoolLiteral; -import org.apache.doris.analysis.DecimalLiteral; -import org.apache.doris.analysis.FloatLiteral; -import org.apache.doris.analysis.IntLiteral; import org.apache.doris.analysis.LiteralExpr; -import org.apache.doris.analysis.NullLiteral; import org.apache.doris.analysis.SetType; import org.apache.doris.analysis.SetVar; -import org.apache.doris.analysis.StringLiteral; import org.apache.doris.analysis.VariableExpr; import org.apache.doris.catalog.Env; import org.apache.doris.catalog.Type; @@ -272,10 +266,6 @@ private static void checkUpdate(SetVar setVar, int flag) throws DdlException { } } - public static void setUserVar(SetVar setVar) { - userVars.put(setVar.getVariable(), setVar.getResult()); - } - // Entry of handling SetVarStmt // Input: // sessionVariable: the variable of current session @@ -517,36 +507,6 @@ private static void fillValue(Object obj, Field field, VariableExpr desc) { } } - // Get variable value through variable name, used to satisfy statement like `SELECT @@comment_version` - public static void fillValueForUserDefinedVar(VariableExpr desc) { - String varName = desc.getName(); - if (userVars.containsKey(varName)) { - LiteralExpr literalExpr = userVars.get(varName); - desc.setType(literalExpr.getType()); - if (literalExpr instanceof BoolLiteral) { - desc.setBoolValue(((BoolLiteral) literalExpr).getValue()); - } else if (literalExpr instanceof IntLiteral) { - desc.setIntValue(((IntLiteral) literalExpr).getValue()); - } else if (literalExpr instanceof FloatLiteral) { - desc.setFloatValue(((FloatLiteral) literalExpr).getValue()); - } else if (literalExpr instanceof DecimalLiteral) { - desc.setDecimalValue(((DecimalLiteral) literalExpr).getValue()); - } else if (literalExpr instanceof StringLiteral) { - desc.setStringValue(((StringLiteral) literalExpr).getValue()); - } else if (literalExpr instanceof NullLiteral) { - desc.setType(Type.NULL); - desc.setIsNull(); - } else { - desc.setType(Type.VARCHAR); - desc.setStringValue(""); - } - } else { - // If there are no such user defined var, just fill the NULL value. - desc.setType(Type.NULL); - desc.setIsNull(); - } - } - private static String getValue(SessionVariable var, String name, SetType setType) throws AnalysisException { VarContext ctx = ctxByVarName.get(name); if (ctx == null) { @@ -618,30 +578,6 @@ private static Literal getLiteral(Object obj, Field field) { return Literal.of(""); } - public static @Nullable Literal getLiteralForUserVar(String varName) { - if (userVars.containsKey(varName)) { - LiteralExpr literalExpr = userVars.get(varName); - if (literalExpr instanceof BoolLiteral) { - return Literal.of(((BoolLiteral) literalExpr).getValue()); - } else if (literalExpr instanceof IntLiteral) { - return Literal.of(((IntLiteral) literalExpr).getValue()); - } else if (literalExpr instanceof FloatLiteral) { - return Literal.of(((FloatLiteral) literalExpr).getValue()); - } else if (literalExpr instanceof DecimalLiteral) { - return Literal.of(((DecimalLiteral) literalExpr).getValue()); - } else if (literalExpr instanceof StringLiteral) { - return Literal.of(((StringLiteral) literalExpr).getValue()); - } else if (literalExpr instanceof NullLiteral) { - return Literal.of(null); - } else { - return Literal.of(""); - } - } else { - // If there are no such user defined var, just return the NULL value. - return Literal.of(null); - } - } - private static String getValue(Object obj, Field field) { try { switch (field.getType().getSimpleName()) { diff --git a/regression-test/data/query_p0/set/test_user_var.out b/regression-test/data/query_p0/set/test_user_var.out index 41aebfe022d892..3a1c176056779c 100644 --- a/regression-test/data/query_p0/set/test_user_var.out +++ b/regression-test/data/query_p0/set/test_user_var.out @@ -14,3 +14,9 @@ true false -- !select5 -- \N \N +-- !select6 -- +2 + +-- !select7 -- +1 + diff --git a/regression-test/suites/query_p0/set/test_user_var.groovy b/regression-test/suites/query_p0/set/test_user_var.groovy index af2d2caadc912e..e653fe3e801b46 100644 --- a/regression-test/suites/query_p0/set/test_user_var.groovy +++ b/regression-test/suites/query_p0/set/test_user_var.groovy @@ -27,4 +27,9 @@ suite("test_user_var") { qt_select3 'select @c1, @c2;' qt_select4 'select @d1, @d2;' qt_select5 'select @f1, @f2;' + + sql "SET @A1=2" + qt_select6 'select @a1' + sql "SET @a1 = 1" + qt_select7 'select @A1' } \ No newline at end of file