This is an automated email from the ASF dual-hosted git repository.

starocean999 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/doris.git


The following commit(s) were added to refs/heads/master by this push:
     new 7e62c3c2de4 [fix](Nereids) store user variable in connect context 
(#26655)
7e62c3c2de4 is described below

commit 7e62c3c2de4e466b1b0cad6351b9626d7c5a6702
Author: 谢健 <jianx...@gmail.com>
AuthorDate: Mon Nov 13 12:25:08 2023 +0800

    [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
---
 .../org/apache/doris/analysis/VariableExpr.java    |  3 +-
 .../doris/nereids/rules/analysis/SlotBinder.java   |  2 +-
 .../java/org/apache/doris/qe/ConnectContext.java   | 75 ++++++++++++++++++++++
 .../main/java/org/apache/doris/qe/SetExecutor.java |  2 +-
 .../main/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 093a7861173..85204952e55 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.catalog.Type;
 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 class VariableExpr extends Expr {
     @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 a504e138909..c098b207e1b 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 class SlotBinder extends SubExprAnalyzer {
         } 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 81c8f7c53e5..dcadf2795b1 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.MysqlCommand;
 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.List;
 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 class ConnectContext {
     protected volatile UserIdentity currentUserIdentity;
     // Variables belong to this session.
     protected volatile SessionVariable sessionVariable;
+    // Store user variable in this connection
+    private Map<String, LiteralExpr> userVars = new HashMap<>();
     // Scheduler this connection belongs to
     protected volatile ConnectScheduler connectScheduler;
     // Executor
@@ -184,6 +198,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);
 
@@ -297,6 +312,7 @@ public class ConnectContext {
         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 class ConnectContext {
         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 9dbe9ca8a98..673f01562ce 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 @@ public class SetExecutor {
             // 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 14afc7a1bcc..03833c01dc0 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 @@ public class VariableMgr {
         }
     }
 
-    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 @@ public class VariableMgr {
         }
     }
 
-    // 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 @@ public class VariableMgr {
         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 41aebfe022d..3a1c1760567 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 af2d2caadc9..e653fe3e801 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


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscr...@doris.apache.org
For additional commands, e-mail: commits-h...@doris.apache.org

Reply via email to