[feature](cmd) add UNSET_VARIABLE statement to set back variables (#27552)

This commit is contained in:
Yulei-Yang
2023-11-27 20:30:04 +08:00
committed by GitHub
parent 36a528b6bc
commit 3d7d166355
7 changed files with 418 additions and 1 deletions

View File

@ -0,0 +1,97 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
package org.apache.doris.analysis;
import org.apache.doris.common.AnalysisException;
import org.apache.doris.common.UserException;
import com.amazonaws.util.StringUtils;
// Unset variables statement
public class UnsetVariableStmt extends StatementBase {
private SetType setType;
// variables to restore
private String variable = null;
private boolean applyToAll = false;
public UnsetVariableStmt(SetType setType, String varName) {
this.setType = setType;
this.variable = varName;
}
public UnsetVariableStmt(SetType setType, boolean applyToAll) {
this.setType = setType;
this.applyToAll = applyToAll;
}
public SetType getSetType() {
return setType;
}
public String getVariable() {
return variable;
}
public boolean isApplyToAll() {
return applyToAll;
}
// change type global to session avoid to write in non-master node.
public void modifySetVarsForExecute() {
if (setType == SetType.GLOBAL) {
setType = SetType.SESSION;
}
}
@Override
public void analyze(Analyzer analyzer) throws UserException {
if (StringUtils.isNullOrEmpty(variable) && !applyToAll) {
throw new AnalysisException("You should specific the unset variable.");
}
}
@Override
public String toSql() {
StringBuilder sb = new StringBuilder();
sb.append("UNSET ");
sb.append(setType).append(" VARIABLE ");
if (!StringUtils.isNullOrEmpty(variable)) {
sb.append(variable).append(" ");
} else if (applyToAll) {
sb.append("ALL");
}
return sb.toString();
}
@Override
public String toString() {
return toSql();
}
@Override
public RedirectStatus getRedirectStatus() {
if (setType == SetType.GLOBAL) {
return RedirectStatus.FORWARD_WITH_SYNC;
}
return RedirectStatus.NO_FORWARD;
}
}

View File

@ -57,6 +57,7 @@ import org.apache.doris.analysis.SelectStmt;
import org.apache.doris.analysis.SetOperationStmt;
import org.apache.doris.analysis.SetStmt;
import org.apache.doris.analysis.SetVar;
import org.apache.doris.analysis.SetVar.SetVarType;
import org.apache.doris.analysis.ShowStmt;
import org.apache.doris.analysis.SlotRef;
import org.apache.doris.analysis.SqlParser;
@ -72,6 +73,7 @@ import org.apache.doris.analysis.TransactionRollbackStmt;
import org.apache.doris.analysis.TransactionStmt;
import org.apache.doris.analysis.UnifiedLoadStmt;
import org.apache.doris.analysis.UnlockTablesStmt;
import org.apache.doris.analysis.UnsetVariableStmt;
import org.apache.doris.analysis.UnsupportedStmt;
import org.apache.doris.analysis.UpdateStmt;
import org.apache.doris.analysis.UseStmt;
@ -735,6 +737,8 @@ public class StmtExecutor {
handleQueryWithRetry(queryId);
} else if (parsedStmt instanceof SetStmt) {
handleSetStmt();
} else if (parsedStmt instanceof UnsetVariableStmt) {
handleUnsetVariableStmt();
} else if (parsedStmt instanceof SwitchStmt) {
handleSwitchStmt();
} else if (parsedStmt instanceof UseStmt) {
@ -1271,6 +1275,30 @@ public class StmtExecutor {
context.getState().setOk();
}
// Process unset variable statement.
private void handleUnsetVariableStmt() {
try {
UnsetVariableStmt unsetStmt = (UnsetVariableStmt) parsedStmt;
if (unsetStmt.isApplyToAll()) {
VariableMgr.setAllVarsToDefaultValue(context.getSessionVariable(), unsetStmt.getSetType());
} else {
String defaultValue = VariableMgr.getDefaultValue(unsetStmt.getVariable());
if (defaultValue == null) {
ErrorReport.reportDdlException(ErrorCode.ERR_UNKNOWN_SYSTEM_VARIABLE, unsetStmt.getVariable());
}
SetVar var = new SetVar(unsetStmt.getSetType(), unsetStmt.getVariable(),
new StringLiteral(defaultValue), SetVarType.SET_SESSION_VAR);
VariableMgr.setVar(context.getSessionVariable(), var);
}
} catch (DdlException e) {
LOG.warn("", e);
// Return error message to client.
context.getState().setError(ErrorCode.ERR_LOCAL_VARIABLE, e.getMessage());
return;
}
context.getState().setOk();
}
// send values from cache.
// return true if the meta fields has been sent, otherwise, return false.
// the meta fields must be sent right before the first batch of data(or eos flag).

View File

@ -20,6 +20,8 @@ package org.apache.doris.qe;
import org.apache.doris.analysis.LiteralExpr;
import org.apache.doris.analysis.SetType;
import org.apache.doris.analysis.SetVar;
import org.apache.doris.analysis.SetVar.SetVarType;
import org.apache.doris.analysis.StringLiteral;
import org.apache.doris.analysis.VariableExpr;
import org.apache.doris.catalog.Env;
import org.apache.doris.catalog.Type;
@ -648,6 +650,33 @@ public class VariableMgr {
return ImmutableMap.copyOf(result);
}
public static void setAllVarsToDefaultValue(SessionVariable sessionVariable, SetType setType)
throws DdlException {
for (Map.Entry<String, VarContext> entry : ctxByDisplayVarName.entrySet()) {
VarContext varCtx = entry.getValue();
SetType newSetType = null;
// some variables are GLOBAL only or SESSION only
if ((varCtx.getFlag() & GLOBAL) != 0) {
newSetType = SetType.GLOBAL;
} else if ((varCtx.getFlag() & SESSION_ONLY) != 0) {
newSetType = SetType.SESSION;
}
SetVar setVar = new SetVar(newSetType != null ? newSetType : setType, entry.getKey(),
new StringLiteral(varCtx.defaultValue), SetVarType.SET_SESSION_VAR);
//skip read only variables
if ((varCtx.getFlag() & READ_ONLY) == 0) {
setVar(sessionVariable, setVar);
}
}
}
public static String getDefaultValue(String key) {
VarContext varContext = ctxByDisplayVarName.get(key);
return varContext == null ? null : varContext.defaultValue;
}
// Dump all fields. Used for `show variables`
public static List<List<String>> dump(SetType type, SessionVariable sessionVar, PatternMatcher matcher) {
List<List<String>> rows = Lists.newArrayList();