[Enhancement](function) support last_query_id function (#40739)

Vallishp · dataroaring · commit fc338ade74e7 · 2024-10-05T19:46:41.000+08:00
support display previous query id.

mysql&gt; select LAST_QUERY_ID();
No connection. Trying to reconnect...
Connection id:    0
Current database: *** NONE ***

+-----------------+
| last_query_id() |
+-----------------+
| Not Available   |
+-----------------+
1 row in set (0.08 sec)

mysql&gt; select LAST_QUERY_ID();
+-----------------------------------+
| last_query_id()                   |
+-----------------------------------+
| e133f2641e0a411c-b73f1f6c67bedcf9 |
+-----------------------------------+
1 row in set (0.01 sec)
diff --git a/fe/fe-core/src/main/java/org/apache/doris/catalog/BuiltinScalarFunctions.java b/fe/fe-core/src/main/java/org/apache/doris/catalog/BuiltinScalarFunctions.java
@@ -261,6 +261,7 @@
 import org.apache.doris.nereids.trees.expressions.functions.scalar.L1Distance;
 import org.apache.doris.nereids.trees.expressions.functions.scalar.L2Distance;
 import org.apache.doris.nereids.trees.expressions.functions.scalar.LastDay;
+import org.apache.doris.nereids.trees.expressions.functions.scalar.LastQueryId;
 import org.apache.doris.nereids.trees.expressions.functions.scalar.Least;
 import org.apache.doris.nereids.trees.expressions.functions.scalar.Left;
 import org.apache.doris.nereids.trees.expressions.functions.scalar.Length;
@@ -928,7 +929,8 @@ public class BuiltinScalarFunctions implements FunctionHelper {
             scalar(YearsDiff.class, "years_diff"),
             scalar(YearsSub.class, "years_sub"),
             scalar(MultiMatch.class, "multi_match"),
-            scalar(SessionUser.class, "session_user"));
+            scalar(SessionUser.class, "session_user"),
+            scalar(LastQueryId.class, "last_query_id"));
 
     public static final BuiltinScalarFunctions INSTANCE = new BuiltinScalarFunctions();
 
diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/expression/rules/FoldConstantRuleOnFE.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/expression/rules/FoldConstantRuleOnFE.java
@@ -20,6 +20,7 @@
 import org.apache.doris.catalog.EncryptKey;
 import org.apache.doris.catalog.Env;
 import org.apache.doris.cluster.ClusterNamespace;
+import org.apache.doris.common.util.DebugUtil;
 import org.apache.doris.nereids.exceptions.AnalysisException;
 import org.apache.doris.nereids.rules.expression.AbstractExpressionRewriteRule;
 import org.apache.doris.nereids.rules.expression.ExpressionListenerMatcher;
@@ -62,6 +63,7 @@
 import org.apache.doris.nereids.trees.expressions.functions.scalar.Date;
 import org.apache.doris.nereids.trees.expressions.functions.scalar.EncryptKeyRef;
 import org.apache.doris.nereids.trees.expressions.functions.scalar.If;
+import org.apache.doris.nereids.trees.expressions.functions.scalar.LastQueryId;
 import org.apache.doris.nereids.trees.expressions.functions.scalar.Password;
 import org.apache.doris.nereids.trees.expressions.functions.scalar.SessionUser;
 import org.apache.doris.nereids.trees.expressions.functions.scalar.User;
@@ -84,6 +86,7 @@
 import org.apache.doris.nereids.util.ExpressionUtils;
 import org.apache.doris.qe.ConnectContext;
 import org.apache.doris.qe.GlobalVariable;
+import org.apache.doris.thrift.TUniqueId;
 
 import com.google.common.base.Preconditions;
 import com.google.common.base.Strings;
@@ -166,7 +169,8 @@ public List<ExpressionPatternMatcher<? extends Expression>> buildRules() {
                 matches(Array.class, this::visitArray),
                 matches(Date.class, this::visitDate),
                 matches(Version.class, this::visitVersion),
-                matches(SessionUser.class, this::visitSessionUser)
+                matches(SessionUser.class, this::visitSessionUser),
+                matches(LastQueryId.class, this::visitLastQueryId)
         );
     }
 
@@ -334,6 +338,16 @@ public Expression visitSessionUser(SessionUser user, ExpressionRewriteContext co
         return new VarcharLiteral(res);
     }
 
+    @Override
+    public Expression visitLastQueryId(LastQueryId queryId, ExpressionRewriteContext context) {
+        String res = "Not Available";
+        TUniqueId id = context.cascadesContext.getConnectContext().getLastQueryId();
+        if (id != null) {
+            res = DebugUtil.printId(id);
+        }
+        return new VarcharLiteral(res);
+    }
+
     @Override
     public Expression visitConnectionId(ConnectionId connectionId, ExpressionRewriteContext context) {
         return new BigIntLiteral(context.cascadesContext.getConnectContext().getConnectionId());
diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/LastQueryId.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/LastQueryId.java
@@ -0,0 +1,59 @@
+// 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.nereids.trees.expressions.functions.scalar;
+
+import org.apache.doris.catalog.FunctionSignature;
+import org.apache.doris.nereids.trees.expressions.functions.AlwaysNotNullable;
+import org.apache.doris.nereids.trees.expressions.functions.ExplicitlyCastableSignature;
+import org.apache.doris.nereids.trees.expressions.shape.LeafExpression;
+import org.apache.doris.nereids.trees.expressions.visitor.ExpressionVisitor;
+import org.apache.doris.nereids.types.VarcharType;
+
+import com.google.common.collect.ImmutableList;
+
+import java.util.List;
+
+/**
+ * ScalarFunction 'last_query_id'.
+ */
+public class LastQueryId extends ScalarFunction
+        implements LeafExpression, ExplicitlyCastableSignature, AlwaysNotNullable {
+
+    public static final List<FunctionSignature> SIGNATURES = ImmutableList.of(
+            FunctionSignature.ret(VarcharType.SYSTEM_DEFAULT).args()
+    );
+
+    public LastQueryId() {
+        super("last_query_id", ImmutableList.of());
+    }
+
+    @Override
+    public List<FunctionSignature> getSignatures() {
+        return SIGNATURES;
+    }
+
+    @Override
+    public <R, C> R accept(ExpressionVisitor<R, C> visitor, C context) {
+        return visitor.visitLastQueryId(this, context);
+    }
+
+    @Override
+    public boolean isDeterministic() {
+        return false;
+    }
+}
diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/visitor/ScalarFunctionVisitor.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/visitor/ScalarFunctionVisitor.java
@@ -264,6 +264,7 @@
 import org.apache.doris.nereids.trees.expressions.functions.scalar.L1Distance;
 import org.apache.doris.nereids.trees.expressions.functions.scalar.L2Distance;
 import org.apache.doris.nereids.trees.expressions.functions.scalar.LastDay;
+import org.apache.doris.nereids.trees.expressions.functions.scalar.LastQueryId;
 import org.apache.doris.nereids.trees.expressions.functions.scalar.Least;
 import org.apache.doris.nereids.trees.expressions.functions.scalar.Left;
 import org.apache.doris.nereids.trees.expressions.functions.scalar.Length;
@@ -2212,4 +2213,8 @@ default R visitStructElement(StructElement structElement, C context) {
     default R visitMultiMatch(MultiMatch multiMatch, C context) {
         return visitScalarFunction(multiMatch, context);
     }
+
+    default R visitLastQueryId(LastQueryId queryId, C context) {
+        return visitScalarFunction(queryId, context);
+    }
 }
diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectContext.java b/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectContext.java
@@ -122,6 +122,7 @@ public enum ConnectType {
     protected volatile TUniqueId queryId = null;
     protected volatile AtomicInteger instanceIdGenerator = new AtomicInteger();
     protected volatile String traceId;
+    protected volatile TUniqueId lastQueryId = null;
     // id for this connection
     protected volatile int connectionId;
     // Timestamp when the connection is make
@@ -866,6 +867,9 @@ public void setKilled() {
     }
 
     public void setQueryId(TUniqueId queryId) {
+        if (this.queryId != null) {
+            this.lastQueryId = this.queryId.deepCopy();
+        }
         this.queryId = queryId;
         if (connectScheduler != null && !Strings.isNullOrEmpty(traceId)) {
             connectScheduler.putTraceId2QueryId(traceId, queryId);
@@ -884,6 +888,10 @@ public TUniqueId queryId() {
         return queryId;
     }
 
+    public TUniqueId getLastQueryId() {
+        return lastQueryId;
+    }
+
     public TUniqueId nextInstanceId() {
         return new TUniqueId(queryId.hi, queryId.lo + instanceIdGenerator.incrementAndGet());
     }
diff --git a/regression-test/suites/query_p0/system/test_query_sys.groovy b/regression-test/suites/query_p0/system/test_query_sys.groovy
@@ -43,6 +43,8 @@ suite("test_query_sys", "query,p0") {
     sql "select pi();"
     sql "select e();"
     sql "select sleep(2);"
+    sql "select last_query_id();"
+    sql "select LAST_QUERY_ID();"
 
     // INFORMATION_SCHEMA
     sql "SELECT table_name FROM INFORMATION_SCHEMA.TABLES where table_schema=\"test_query_db\" and TABLE_TYPE = \"BASE TABLE\" order by table_name"
