[Feature] Add function get_query_dump

fe/fe-core/src/main/java/com/starrocks/catalog/Function.java

@@ -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() {
     }
@@ -225,6 +228,7 @@ public Function(Function other) {
         isPolymorphic = other.isPolymorphic;
         couldApplyDictOptimize = other.couldApplyDictOptimize;
         isNullable = other.isNullable;
+        isMetaFunction = other.isMetaFunction;
     }
 
     public FunctionName getFunctionName() {
@@ -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;
     }

fe/fe-core/src/main/java/com/starrocks/http/rest/QueryDumpAction.java

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

fe/fe-core/src/main/java/com/starrocks/sql/optimizer/dump/QueryDumper.java

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

fe/fe-core/src/main/java/com/starrocks/sql/optimizer/function/MetaFunctions.java

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

fe/fe-core/src/main/java/com/starrocks/sql/optimizer/rewrite/ScalarOperatorEvaluator.java

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

fe/fe-core/src/main/java/com/starrocks/sql/optimizer/rewrite/scalar/FoldConstantsRule.java

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