database/doris
2
0
Fork 0
Code Issues Pull Requests Actions 3 Packages Projects Releases Wiki Activity

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

Browse Source 
1.user variable should be case insensitive
2.user variable should be cleared after the connection reset
This commit is contained in:
谢健
2023-11-13 12:25:08 +08:00
committed by GitHub
parent fa3c7d98c8
commit 7e62c3c2de
7 changed files with 90 additions and 67 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
fe/fe-core/src/main/java/org/apache/doris/analysis/VariableExpr.java
View File

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

2
fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/SlotBinder.java
View File

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

75
fe/fe-core/src/main/java/org/apache/doris/qe/ConnectContext.java
View File

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

2
fe/fe-core/src/main/java/org/apache/doris/qe/SetExecutor.java
View File

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

64
fe/fe-core/src/main/java/org/apache/doris/qe/VariableMgr.java
View File

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