[pick](Nereids) store user variable in connect context (#26655)

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)) {

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());

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<String, LiteralExpr> 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;
     }

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);
         }

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()) {

regression-test/data/query_p0/set/test_user_var.out

@@ -14,3 +14,9 @@ true	false
 -- !select5 --
 \N	\N
 
+-- !select6 --
+2
+
+-- !select7 --
+1
+

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'
 }