Skip to content

Commit

Permalink
[fix](Nereids) store user variable in connect context (#26655) (#26920)
Browse files Browse the repository at this point in the history 
pick from master #26655

1. user variable should be case insensitive
2. user variable should be cleared after the connection reset
  • Loading branch information
@keanji-x
keanji-x authored Nov 16, 2023
1 parent bdf9c5f commit 8df819e
Show file tree
Hide file tree
Showing 7 changed files with 101 additions and 67 deletions.
3 changes: 2 additions & 1 deletion fe/fe-core/src/main/java/org/apache/doris/analysis/VariableExpr.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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)) {
Expand Down
2 changes: 1 addition & 1 deletion fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/SlotBinder.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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());
Expand Down
86 changes: 86 additions & 0 deletions fe/fe-core/src/main/java/org/apache/doris/qe/ConnectContext.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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);

Expand Down Expand Up @@ -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);
}
Expand Down Expand Up @@ -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;
}
Expand Down
2 changes: 1 addition & 1 deletion fe/fe-core/src/main/java/org/apache/doris/qe/SetExecutor.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
Expand Down
64 changes: 0 additions & 64 deletions fe/fe-core/src/main/java/org/apache/doris/qe/VariableMgr.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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()) {
Expand Down
6 changes: 6 additions & 0 deletions regression-test/data/query_p0/set/test_user_var.out
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,3 +14,9 @@ true false
-- !select5 --
\N \N

-- !select6 --
2

-- !select7 --
1

5 changes: 5 additions & 0 deletions regression-test/suites/query_p0/set/test_user_var.groovy
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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'
}

0 comments on commit 8df819e

Please sign in to comment.