From 7e62c3c2de4e466b1b0cad6351b9626d7c5a6702 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 | 75 +++++++++++++++++++ .../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, 90 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 093a7861173c0a..85204952e557b7 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 a504e138909ed9..c098b207e1bc6e 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 @@ -97,7 +97,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 81c8f7c53e5cd5..dcadf2795b1f21 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,12 +17,22 @@ package org.apache.doris.qe; +import org.apache.doris.analysis.BoolLiteral; +import org.apache.doris.analysis.DecimalLiteral; import org.apache.doris.analysis.Expr; +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.util.DebugUtil; @@ -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.service.arrowflight.results.FlightSqlChannel; @@ -66,6 +77,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 @@ -136,6 +148,8 @@ public enum ConnectType { 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 @@ -184,6 +198,7 @@ public enum ConnectType { private SessionContext sessionContext; + // This context is used for SSL connection between server and mysql client. private final MysqlSslContext mysqlSslContext = new MysqlSslContext(SSL_PROTOCOL); @@ -297,6 +312,7 @@ public void init() { returnRows = 0; isKilled = false; sessionVariable = VariableMgr.newSessionVariable(); + userVars = new HashMap<>(); command = MysqlCommand.COM_SLEEP; if (Config.use_fuzzy_session_variable) { sessionVariable.initFuzzyModeVariables(); @@ -446,6 +462,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 14afc7a1bcc24b..03833c01dc098c 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; @@ -271,10 +265,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 @@ -542,36 +532,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) { @@ -643,30 +603,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