Skip to content

Commit

Permalink
[Feature] Add function get_query_dump (backport #48105) (#48212)
Browse files Browse the repository at this point in the history
Co-authored-by: zihe.liu <[email protected]>
  • Loading branch information
mergify[bot] and ZiheLiu authored Jul 12, 2024
1 parent 900727b commit 27978d7
Show file tree
Hide file tree
Showing 12 changed files with 685 additions and 61 deletions.
12 changes: 12 additions & 0 deletions fe/fe-core/src/main/java/com/starrocks/catalog/Function.java
Original file line number Diff line number Diff line change
Expand Up @@ -150,6 +150,9 @@ public enum CompareMode {
private boolean isNullable = true;

private Vector<Pair<String, Expr>> defaultArgExprs;

private boolean isMetaFunction = false;

// Only used for serialization
protected Function() {
}
Expand Down Expand Up @@ -225,6 +228,7 @@ public Function(Function other) {
isPolymorphic = other.isPolymorphic;
couldApplyDictOptimize = other.couldApplyDictOptimize;
isNullable = other.isNullable;
isMetaFunction = other.isMetaFunction;
}

public FunctionName getFunctionName() {
Expand Down Expand Up @@ -341,6 +345,14 @@ public boolean isPolymorphic() {
return isPolymorphic;
}

public boolean isMetaFunction() {
return isMetaFunction;
}

public void setMetaFunction(boolean metaFunction) {
isMetaFunction = metaFunction;
}

public long getFunctionId() {
return functionId;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -15,25 +15,15 @@
package com.starrocks.http.rest;

import com.google.common.base.Strings;
import com.google.gson.Gson;
import com.google.gson.GsonBuilder;
import com.starrocks.catalog.Database;
import com.starrocks.catalog.InternalCatalog;
import com.starrocks.common.DdlException;
import com.starrocks.common.Pair;
import com.starrocks.http.ActionController;
import com.starrocks.http.BaseRequest;
import com.starrocks.http.BaseResponse;
import com.starrocks.http.IllegalArgException;
import com.starrocks.persist.gson.GsonUtils;
import com.starrocks.qe.ConnectContext;
import com.starrocks.qe.StmtExecutor;
import com.starrocks.server.GlobalStateMgr;
import com.starrocks.sql.ast.StatementBase;
import com.starrocks.sql.optimizer.dump.DumpInfo;
import com.starrocks.sql.optimizer.dump.QueryDumpDeserializer;
import com.starrocks.sql.optimizer.dump.QueryDumpInfo;
import com.starrocks.sql.optimizer.dump.QueryDumpSerializer;
import com.starrocks.sql.parser.SqlParser;
import com.starrocks.sql.optimizer.dump.QueryDumper;
import io.netty.handler.codec.http.HttpMethod;
import io.netty.handler.codec.http.HttpResponseStatus;
import org.apache.commons.lang3.StringUtils;
Expand All @@ -50,25 +40,17 @@

public class QueryDumpAction extends RestBaseAction {
private static final Logger LOG = LogManager.getLogger(QueryDumpAction.class);
private static final String DB = "db";

public static final String URL = "/api/query_dump";
private static final String DB = "db";
private static final String MOCK = "mock";

private static final Gson GSON = new GsonBuilder()
.addSerializationExclusionStrategy(new GsonUtils.HiddenAnnotationExclusionStrategy())
.addDeserializationExclusionStrategy(new GsonUtils.HiddenAnnotationExclusionStrategy())
.enableComplexMapKeySerialization()
.disableHtmlEscaping()
.registerTypeAdapter(QueryDumpInfo.class, new QueryDumpSerializer())
.registerTypeAdapter(QueryDumpInfo.class, new QueryDumpDeserializer())
.create();

public QueryDumpAction(ActionController controller) {
super(controller);
}

public static void registerAction(ActionController controller) throws IllegalArgException {
controller.registerHandler(HttpMethod.POST, "/api/query_dump", new QueryDumpAction(controller));
controller.registerHandler(HttpMethod.POST, URL, new QueryDumpAction(controller));
}

@Override
Expand All @@ -78,52 +60,24 @@ public void executeWithoutPassword(BaseRequest request, BaseResponse response) t
boolean enableMock = request.getSingleParameter(MOCK) == null ||
"true".equalsIgnoreCase(StringUtils.trim(request.getSingleParameter(MOCK)));

String catalogName = "";
String dbName = "";
if (!Strings.isNullOrEmpty(catalogDbName)) {
String[] catalogDbNames = catalogDbName.split("\\.");

String catalogName = InternalCatalog.DEFAULT_INTERNAL_CATALOG_NAME;
catalogName = InternalCatalog.DEFAULT_INTERNAL_CATALOG_NAME;
if (catalogDbNames.length == 2) {
catalogName = catalogDbNames[0];
}
String dbName = catalogDbNames[catalogDbNames.length - 1];
context.setCurrentCatalog(catalogName);
Database db = GlobalStateMgr.getCurrentState().getMetadataMgr().getDb(catalogName, dbName);
if (db == null) {
response.getContent().append("Database [" + dbName + "] does not exists");
sendResult(request, response, HttpResponseStatus.NOT_FOUND);
return;
}
context.setDatabase(db.getFullName());
dbName = catalogDbNames[catalogDbNames.length - 1];
}
context.setIsHTTPQueryDump(true);

String query = request.getContent();
if (Strings.isNullOrEmpty(query)) {
response.getContent().append("not valid parameter");
sendResult(request, response, HttpResponseStatus.BAD_REQUEST);
return;
}

StatementBase parsedStmt;
try {
parsedStmt = SqlParser.parse(query, context.getSessionVariable()).get(0);
StmtExecutor executor = new StmtExecutor(context, parsedStmt);
executor.execute();
} catch (Exception e) {
LOG.warn("execute query failed. ", e);
response.getContent().append("execute query failed. " + e.getMessage());
sendResult(request, response, HttpResponseStatus.BAD_REQUEST);
return;
}
Pair<HttpResponseStatus, String> statusAndRes = QueryDumper.dumpQuery(catalogName, dbName, query, enableMock);

DumpInfo dumpInfo = context.getDumpInfo();
if (dumpInfo != null) {
dumpInfo.setDesensitizedInfo(enableMock);
response.getContent().append(GSON.toJson(dumpInfo, QueryDumpInfo.class));
sendResult(request, response);
} else {
response.getContent().append("not use cbo planner, try again.");
sendResult(request, response, HttpResponseStatus.BAD_REQUEST);
}
response.getContent().append(statusAndRes.second);
sendResult(request, response, statusAndRes.first);
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,100 @@
// Copyright 2021-present StarRocks, Inc. All rights reserved.
//
// Licensed 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
//
// https://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 com.starrocks.sql.optimizer.dump;

import com.google.gson.Gson;
import com.google.gson.GsonBuilder;
import com.starrocks.catalog.Database;
import com.starrocks.common.Pair;
import com.starrocks.persist.gson.GsonUtils;
import com.starrocks.qe.ConnectContext;
import com.starrocks.qe.StmtExecutor;
import com.starrocks.server.GlobalStateMgr;
import com.starrocks.sql.ast.StatementBase;
import com.starrocks.sql.parser.SqlParser;
import io.netty.handler.codec.http.HttpResponseStatus;
import org.apache.commons.lang3.StringUtils;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;

public class QueryDumper {
private static final Logger LOG = LogManager.getLogger(QueryDumper.class);

private static final Gson GSON = new GsonBuilder()
.addSerializationExclusionStrategy(new GsonUtils.HiddenAnnotationExclusionStrategy())
.addDeserializationExclusionStrategy(new GsonUtils.HiddenAnnotationExclusionStrategy())
.enableComplexMapKeySerialization()
.disableHtmlEscaping()
.registerTypeAdapter(QueryDumpInfo.class, new QueryDumpSerializer())
.registerTypeAdapter(QueryDumpInfo.class, new QueryDumpDeserializer())
.create();

public static Pair<HttpResponseStatus, String> dumpQuery(String catalogName, String dbName, String query,
boolean enableMock) {
ConnectContext context = ConnectContext.get();
if (context == null) {
return Pair.create(HttpResponseStatus.BAD_REQUEST,
"There is no ConnectContext for this thread: " + Thread.currentThread().getName());
}

final String prevCatalog = context.getCurrentCatalog();
final String prevDb = context.getDatabase();
final boolean prevIsHTTPQueryDump = context.isHTTPQueryDump();

try {
if (StringUtils.isEmpty(query)) {
return Pair.create(HttpResponseStatus.BAD_REQUEST, "query is empty");
}

if (!StringUtils.isEmpty(catalogName)) {
context.setCurrentCatalog(catalogName);
}

if (!StringUtils.isEmpty(dbName)) {
Database db = GlobalStateMgr.getCurrentState().getMetadataMgr().getDb(catalogName, dbName);
if (db == null) {
return Pair.create(HttpResponseStatus.NOT_FOUND,
String.format("Database [%s.%s] does not exists", catalogName, dbName));
}
context.setDatabase(db.getFullName());
}

context.setIsHTTPQueryDump(true);

StatementBase parsedStmt;
try {
parsedStmt = SqlParser.parse(query, context.getSessionVariable()).get(0);
StmtExecutor executor = new StmtExecutor(context, parsedStmt);
executor.execute();
} catch (Exception e) {
LOG.warn("execute query failed. ", e);
return Pair.create(HttpResponseStatus.BAD_REQUEST, "execute query failed. " + e.getMessage());
}

DumpInfo dumpInfo = context.getDumpInfo();
if (dumpInfo != null) {
dumpInfo.setDesensitizedInfo(enableMock);
String dumpStr = GSON.toJson(dumpInfo, QueryDumpInfo.class);
return Pair.create(HttpResponseStatus.OK, dumpStr);
} else {
return Pair.create(HttpResponseStatus.BAD_REQUEST, "not use cbo planner, try again.");
}
} finally {
context.setCurrentCatalog(prevCatalog);
context.setDatabase(prevDb);
context.setIsHTTPQueryDump(prevIsHTTPQueryDump);
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -44,8 +44,10 @@
import com.starrocks.sql.analyzer.Authorizer;
import com.starrocks.sql.optimizer.CachingMvPlanContextBuilder;
import com.starrocks.sql.optimizer.OptExpression;
import com.starrocks.sql.optimizer.dump.QueryDumper;
import com.starrocks.sql.optimizer.operator.scalar.ConstantOperator;
import com.starrocks.sql.optimizer.rewrite.ConstantFunction;
import io.netty.handler.codec.http.HttpResponseStatus;
import org.apache.commons.collections4.MapUtils;
import org.apache.commons.collections4.SetUtils;
import org.apache.commons.lang3.tuple.Pair;
Expand Down Expand Up @@ -346,4 +348,20 @@ public static ConstantOperator inspectMvPlan(ConstantOperator mvName, ConstantOp
return ConstantOperator.createVarchar("failed");
}
}

@ConstantFunction(name = "get_query_dump", argTypes = {VARCHAR, BOOLEAN}, returnType = VARCHAR, isMetaFunction = true)
public static ConstantOperator getQueryDump(ConstantOperator query, ConstantOperator enableMock) {
com.starrocks.common.Pair<HttpResponseStatus, String> statusAndRes =
QueryDumper.dumpQuery("", "", query.getVarchar(), enableMock.getBoolean());
if (statusAndRes.first != HttpResponseStatus.OK) {
ErrorReport.reportSemanticException(ErrorCode.ERR_INVALID_PARAMETER, "get_query_dump: " + statusAndRes.second);
}
return ConstantOperator.createVarchar(statusAndRes.second);
}

@ConstantFunction(name = "get_query_dump", argTypes = {VARCHAR}, returnType = VARCHAR, isMetaFunction = true)
public static ConstantOperator getQueryDump(ConstantOperator query) {
return getQueryDump(query, ConstantOperator.createBoolean(false));
}

}
Original file line number Diff line number Diff line change
Expand Up @@ -117,7 +117,10 @@ public Function getMetaFunction(FunctionName name, Type[] args) {
if (invoker == null || !invoker.isMetaFunction) {
return null;
}
return new Function(name, Lists.newArrayList(args), Type.VARCHAR, false);

Function function = new Function(name, Lists.newArrayList(args), Type.VARCHAR, false);
function.setMetaFunction(true);
return function;
}

public ScalarOperator evaluation(CallOperator root) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -12,11 +12,11 @@
// See the License for the specific language governing permissions and
// limitations under the License.


package com.starrocks.sql.optimizer.rewrite.scalar;

import com.starrocks.analysis.BinaryType;
import com.starrocks.catalog.Type;
import com.starrocks.sql.analyzer.SemanticException;
import com.starrocks.sql.optimizer.operator.scalar.BinaryPredicateOperator;
import com.starrocks.sql.optimizer.operator.scalar.CallOperator;
import com.starrocks.sql.optimizer.operator.scalar.CastOperator;
Expand Down Expand Up @@ -48,9 +48,19 @@ public FoldConstantsRule(boolean needMonotonicFunc) {

@Override
public ScalarOperator visitCall(CallOperator call, ScalarOperatorRewriteContext context) {
if (call.isAggregate() || notAllConstant(call.getChildren())) {
if (call.isAggregate()) {
return call;
}

if (notAllConstant(call.getChildren())) {
if (call.getFunction() != null && call.getFunction().isMetaFunction()) {
String errMsg = String.format("Meta function %s does not support non-constant arguments",
call.getFunction().getFunctionName());
throw new SemanticException(errMsg);
}
return call;
}

return ScalarOperatorEvaluator.INSTANCE.evaluation(call, needMonotonicFunc);
}

Expand Down
Loading

0 comments on commit 27978d7

Please sign in to comment.