From 2697f72d778df16ab3da6e538f62fa6de25ab3cb Mon Sep 17 00:00:00 2001
From: DongLiang-0 <46414265+DongLiang-0@users.noreply.github.com>
Date: Thu, 27 Oct 2022 10:03:39 +0800
Subject: [PATCH] [Improvement][SET-PROPERTY] Support for set query_timeout
 property  (#13444)

---
 .../SET-PROPERTY.md                           |  8 ++++++
 .../SET-PROPERTY.md                           |  8 ++++++
 .../doris/mysql/nio/AcceptListener.java       |  2 ++
 .../mysql/privilege/CommonUserProperties.java | 11 ++++++++
 .../doris/mysql/privilege/PaloAuth.java       |  9 +++++++
 .../doris/mysql/privilege/UserProperty.java   | 20 ++++++++++++++
 .../mysql/privilege/UserPropertyMgr.java      |  9 +++++++
 .../org/apache/doris/qe/ConnectContext.java   | 27 +++++++++++++++----
 .../org/apache/doris/qe/ConnectScheduler.java |  1 +
 .../doris/catalog/UserPropertyTest.java       |  4 +++
 .../doris/planner/ResourceTagQueryTest.java   |  2 +-
 .../apache/doris/qe/ConnectContextTest.java   | 13 +++++++++
 12 files changed, 108 insertions(+), 6 deletions(-)

diff --git a/docs/en/docs/sql-manual/sql-reference/Account-Management-Statements/SET-PROPERTY.md b/docs/en/docs/sql-manual/sql-reference/Account-Management-Statements/SET-PROPERTY.md
index 100064fb4a..aa16dbbed9 100644
--- a/docs/en/docs/sql-manual/sql-reference/Account-Management-Statements/SET-PROPERTY.md
+++ b/docs/en/docs/sql-manual/sql-reference/Account-Management-Statements/SET-PROPERTY.md
@@ -60,6 +60,8 @@ Super user privileges:
 
  resource_tags: Specifies the user's resource tag permissions.
 
+ query_timeout: Specifies the user's query timeout permissions.
+
     Note: If the attributes `cpu_resource_limit`, `exec_mem_limit` are not set, the value in the session variable will be used by default.
 
 Ordinary user rights:
@@ -156,6 +158,12 @@ Data, etl program automatically retains the next use.
     SET PROPERTY FOR 'jack' 'exec_mem_limit' = '2147483648';
     ````
 
+13. Modify the user's query timeout limit, in second
+
+    ```sql
+    SET PROPERTY FOR 'jack' 'query_timeout' = '500';
+    ````
+
 ### Keywords
 
     SET, PROPERTY
diff --git a/docs/zh-CN/docs/sql-manual/sql-reference/Account-Management-Statements/SET-PROPERTY.md b/docs/zh-CN/docs/sql-manual/sql-reference/Account-Management-Statements/SET-PROPERTY.md
index b8d911ad62..7929edfd32 100644
--- a/docs/zh-CN/docs/sql-manual/sql-reference/Account-Management-Statements/SET-PROPERTY.md
+++ b/docs/zh-CN/docs/sql-manual/sql-reference/Account-Management-Statements/SET-PROPERTY.md
@@ -60,6 +60,8 @@ key:
 
 ​        resource_tags：指定用户的资源标签权限。
 
+​        query_timeout：指定用户的查询超时权限。
+
     注：`cpu_resource_limit`, `exec_mem_limit` 两个属性如果未设置，则默认使用会话变量中值。
 
 普通用户权限：
@@ -155,6 +157,12 @@ key:
     ```sql
     SET PROPERTY FOR 'jack' 'exec_mem_limit' = '2147483648';
     ```
+
+13. 修改用户的查询超时限制，单位秒
+
+    ```sql
+    SET PROPERTY FOR 'jack' 'query_timeout' = '500';
+    ```
     
 ### Keywords
 
diff --git a/fe/fe-core/src/main/java/org/apache/doris/mysql/nio/AcceptListener.java b/fe/fe-core/src/main/java/org/apache/doris/mysql/nio/AcceptListener.java
index b4f28be2c5..8744a272cd 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/mysql/nio/AcceptListener.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/mysql/nio/AcceptListener.java
@@ -78,6 +78,8 @@ public class AcceptListener implements ChannelListener<AcceptingChannel<StreamCo
                         throw new AfterConnectedException("Reach limit of connections");
                     }
                     context.setStartTime();
+                    context.setUserQueryTimeout(
+                            context.getEnv().getAuth().getQueryTimeout(context.getQualifiedUser()));
                     ConnectProcessor processor = new ConnectProcessor(context);
                     context.startAcceptQuery(processor);
                 } catch (AfterConnectedException e) {
diff --git a/fe/fe-core/src/main/java/org/apache/doris/mysql/privilege/CommonUserProperties.java b/fe/fe-core/src/main/java/org/apache/doris/mysql/privilege/CommonUserProperties.java
index 0b83d12898..255388abf8 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/mysql/privilege/CommonUserProperties.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/mysql/privilege/CommonUserProperties.java
@@ -51,6 +51,9 @@ public class CommonUserProperties implements Writable {
     @SerializedName("execMemLimit")
     private long execMemLimit = -1;
 
+    @SerializedName("queryTimeout")
+    private long queryTimeout = -1;
+
     private String[] sqlBlockRulesSplit = {};
 
     long getMaxConn() {
@@ -111,6 +114,14 @@ public class CommonUserProperties implements Writable {
         this.execMemLimit = execMemLimit;
     }
 
+    public long getQueryTimeout() {
+        return queryTimeout;
+    }
+
+    public void setQueryTimeout(long timeout) {
+        this.queryTimeout = timeout;
+    }
+
     public static CommonUserProperties read(DataInput in) throws IOException {
         String json = Text.readString(in);
         CommonUserProperties commonUserProperties = GsonUtils.GSON.fromJson(json, CommonUserProperties.class);
diff --git a/fe/fe-core/src/main/java/org/apache/doris/mysql/privilege/PaloAuth.java b/fe/fe-core/src/main/java/org/apache/doris/mysql/privilege/PaloAuth.java
index 92ea8384ad..4729c4552a 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/mysql/privilege/PaloAuth.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/mysql/privilege/PaloAuth.java
@@ -1363,6 +1363,15 @@ public class PaloAuth implements Writable {
         }
     }
 
+    public long getQueryTimeout(String qualifiedUser) {
+        readLock();
+        try {
+            return propertyMgr.getQueryTimeout(qualifiedUser);
+        } finally {
+            readUnlock();
+        }
+    }
+
     public long getMaxQueryInstances(String qualifiedUser) {
         readLock();
         try {
diff --git a/fe/fe-core/src/main/java/org/apache/doris/mysql/privilege/UserProperty.java b/fe/fe-core/src/main/java/org/apache/doris/mysql/privilege/UserProperty.java
index a30f2d8d5c..741493a4c4 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/mysql/privilege/UserProperty.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/mysql/privilege/UserProperty.java
@@ -63,6 +63,7 @@ public class UserProperty implements Writable {
     private static final String PROP_SQL_BLOCK_RULES = "sql_block_rules";
     private static final String PROP_CPU_RESOURCE_LIMIT = "cpu_resource_limit";
     private static final String PROP_EXEC_MEM_LIMIT = "exec_mem_limit";
+    private static final String PROP_USER_QUERY_TIMEOUT = "query_timeout";
     // advanced properties end
 
     private static final String PROP_LOAD_CLUSTER = "load_cluster";
@@ -108,6 +109,7 @@ public class UserProperty implements Writable {
         ADVANCED_PROPERTIES.add(Pattern.compile("^" + PROP_CPU_RESOURCE_LIMIT + "$", Pattern.CASE_INSENSITIVE));
         ADVANCED_PROPERTIES.add(Pattern.compile("^" + PROP_RESOURCE_TAGS + "$", Pattern.CASE_INSENSITIVE));
         ADVANCED_PROPERTIES.add(Pattern.compile("^" + PROP_EXEC_MEM_LIMIT + "$", Pattern.CASE_INSENSITIVE));
+        ADVANCED_PROPERTIES.add(Pattern.compile("^" + PROP_USER_QUERY_TIMEOUT + "$", Pattern.CASE_INSENSITIVE));
 
         COMMON_PROPERTIES.add(Pattern.compile("^" + PROP_QUOTA + ".", Pattern.CASE_INSENSITIVE));
         COMMON_PROPERTIES.add(Pattern.compile("^" + PROP_DEFAULT_LOAD_CLUSTER + "$", Pattern.CASE_INSENSITIVE));
@@ -130,6 +132,10 @@ public class UserProperty implements Writable {
         return this.commonProperties.getMaxConn();
     }
 
+    public long getQueryTimeout() {
+        return this.commonProperties.getQueryTimeout();
+    }
+
     public long getMaxQueryInstances() {
         return commonProperties.getMaxQueryInstances(); // maxQueryInstances;
     }
@@ -176,6 +182,7 @@ public class UserProperty implements Writable {
         int cpuResourceLimit = this.commonProperties.getCpuResourceLimit();
         Set<Tag> resourceTags = this.commonProperties.getResourceTags();
         long execMemLimit = this.commonProperties.getExecMemLimit();
+        long queryTimeout = this.commonProperties.getQueryTimeout();
 
         UserResource newResource = resource.getCopiedUserResource();
         String newDefaultLoadCluster = defaultLoadCluster;
@@ -314,6 +321,15 @@ public class UserProperty implements Writable {
             } else if (keyArr[0].equalsIgnoreCase(PROP_EXEC_MEM_LIMIT)) {
                 // set property "exec_mem_limit" = "2147483648";
                 execMemLimit = getLongProperty(key, value, keyArr, PROP_EXEC_MEM_LIMIT);
+            } else if (keyArr[0].equalsIgnoreCase(PROP_USER_QUERY_TIMEOUT)) {
+                if (keyArr.length != 1) {
+                    throw new DdlException(PROP_MAX_USER_CONNECTIONS + " format error");
+                }
+                try {
+                    queryTimeout = Long.parseLong(value);
+                } catch (NumberFormatException e) {
+                    throw new DdlException(PROP_USER_QUERY_TIMEOUT + " is not number");
+                }
             } else {
                 throw new DdlException("Unknown user property(" + key + ")");
             }
@@ -326,6 +342,7 @@ public class UserProperty implements Writable {
         this.commonProperties.setCpuResourceLimit(cpuResourceLimit);
         this.commonProperties.setResourceTags(resourceTags);
         this.commonProperties.setExecMemLimit(execMemLimit);
+        this.commonProperties.setQueryTimeout(queryTimeout);
         resource = newResource;
         if (newDppConfigs.containsKey(newDefaultLoadCluster)) {
             defaultLoadCluster = newDefaultLoadCluster;
@@ -452,6 +469,9 @@ public class UserProperty implements Writable {
         // exec mem limit
         result.add(Lists.newArrayList(PROP_EXEC_MEM_LIMIT, String.valueOf(commonProperties.getExecMemLimit())));
 
+        // query timeout
+        result.add(Lists.newArrayList(PROP_USER_QUERY_TIMEOUT, String.valueOf(commonProperties.getQueryTimeout())));
+
         // resource tag
         result.add(Lists.newArrayList(PROP_RESOURCE_TAGS, Joiner.on(", ").join(commonProperties.getResourceTags())));
 
diff --git a/fe/fe-core/src/main/java/org/apache/doris/mysql/privilege/UserPropertyMgr.java b/fe/fe-core/src/main/java/org/apache/doris/mysql/privilege/UserPropertyMgr.java
index 7d824971f7..df666e7ea5 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/mysql/privilege/UserPropertyMgr.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/mysql/privilege/UserPropertyMgr.java
@@ -130,6 +130,15 @@ public class UserPropertyMgr implements Writable {
         property.update(properties);
     }
 
+    public long getQueryTimeout(String qualifiedUser) {
+        UserProperty existProperty = propertyMap.get(qualifiedUser);
+        existProperty = getLdapPropertyIfNull(qualifiedUser, existProperty);
+        if (existProperty == null) {
+            return 0;
+        }
+        return existProperty.getQueryTimeout();
+    }
+
     public long getMaxConn(String qualifiedUser) {
         UserProperty existProperty = propertyMap.get(qualifiedUser);
         existProperty = getLdapPropertyIfNull(qualifiedUser, existProperty);
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
index b6e7ac9d84..53b3c6cde7 100644
--- 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
@@ -147,6 +147,12 @@ public class ConnectContext {
 
     private SessionContext sessionContext;
 
+    private long userQueryTimeout;
+
+    public void setUserQueryTimeout(long queryTimeout) {
+        this.userQueryTimeout = queryTimeout;
+    }
+
     private StatementContext statementContext;
 
     public SessionContext getSessionContext() {
@@ -562,12 +568,23 @@ public class ConnectContext {
                 killConnection = true;
             }
         } else {
-            if (delta > sessionVariable.getQueryTimeoutS() * 1000) {
-                LOG.warn("kill query timeout, remote: {}, query timeout: {}",
-                        getMysqlChannel().getRemoteHostPortString(), sessionVariable.getQueryTimeoutS());
+            if (userQueryTimeout > 0) {
+                // user set query_timeout property
+                if (delta > userQueryTimeout * 1000) {
+                    LOG.warn("kill query timeout, remote: {}, query timeout: {}",
+                            getMysqlChannel().getRemoteHostPortString(), userQueryTimeout);
 
-                // Only kill
-                killFlag = true;
+                    killFlag = true;
+                }
+            } else {
+                // default use session query_timeout
+                if (delta > sessionVariable.getQueryTimeoutS() * 1000) {
+                    LOG.warn("kill query timeout, remote: {}, query timeout: {}",
+                            getMysqlChannel().getRemoteHostPortString(), sessionVariable.getQueryTimeoutS());
+
+                    // Only kill
+                    killFlag = true;
+                }
             }
         }
         if (killFlag) {
diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectScheduler.java b/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectScheduler.java
index 66702d438b..e31dfea298 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectScheduler.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectScheduler.java
@@ -193,6 +193,7 @@ public class ConnectScheduler {
                     return;
                 }
 
+                context.setUserQueryTimeout(context.getEnv().getAuth().getQueryTimeout(context.getQualifiedUser()));
                 context.setStartTime();
                 ConnectProcessor processor = new ConnectProcessor(context);
                 processor.loop();
diff --git a/fe/fe-core/src/test/java/org/apache/doris/catalog/UserPropertyTest.java b/fe/fe-core/src/test/java/org/apache/doris/catalog/UserPropertyTest.java
index e7e8435871..c8c3613935 100644
--- a/fe/fe-core/src/test/java/org/apache/doris/catalog/UserPropertyTest.java
+++ b/fe/fe-core/src/test/java/org/apache/doris/catalog/UserPropertyTest.java
@@ -107,6 +107,7 @@ public class UserPropertyTest {
         properties.add(Pair.of("max_qUERY_instances", "3000"));
         properties.add(Pair.of("sql_block_rules", "rule1,rule2"));
         properties.add(Pair.of("cpu_resource_limit", "2"));
+        properties.add(Pair.of("query_timeout", "500"));
 
         UserProperty userProperty = new UserProperty();
         userProperty.update(properties);
@@ -118,6 +119,7 @@ public class UserPropertyTest {
         Assert.assertEquals(3000, userProperty.getMaxQueryInstances());
         Assert.assertEquals(new String[]{"rule1", "rule2"}, userProperty.getSqlBlockRules());
         Assert.assertEquals(2, userProperty.getCpuResourceLimit());
+        Assert.assertEquals(500, userProperty.getQueryTimeout());
 
         // fetch property
         List<List<String>> rows = userProperty.fetchProperty();
@@ -141,6 +143,8 @@ public class UserPropertyTest {
                 Assert.assertEquals("rule1,rule2", value);
             } else if (key.equalsIgnoreCase("cpu_resource_limit")) {
                 Assert.assertEquals("2", value);
+            } else if (key.equalsIgnoreCase("query_timeout")) {
+                Assert.assertEquals("500", value);
             }
         }
 
diff --git a/fe/fe-core/src/test/java/org/apache/doris/planner/ResourceTagQueryTest.java b/fe/fe-core/src/test/java/org/apache/doris/planner/ResourceTagQueryTest.java
index 391fb6c0d7..3d531a2024 100644
--- a/fe/fe-core/src/test/java/org/apache/doris/planner/ResourceTagQueryTest.java
+++ b/fe/fe-core/src/test/java/org/apache/doris/planner/ResourceTagQueryTest.java
@@ -279,7 +279,7 @@ public class ResourceTagQueryTest {
         Assert.assertEquals(1000000, execMemLimit);
 
         List<List<String>> userProps = Env.getCurrentEnv().getAuth().getUserProperties(PaloAuth.ROOT_USER);
-        Assert.assertEquals(16, userProps.size());
+        Assert.assertEquals(17, userProps.size());
     }
 
     private void checkTableReplicaAllocation(OlapTable tbl) throws InterruptedException {
diff --git a/fe/fe-core/src/test/java/org/apache/doris/qe/ConnectContextTest.java b/fe/fe-core/src/test/java/org/apache/doris/qe/ConnectContextTest.java
index 2a984d9564..8fdcdadb1c 100644
--- a/fe/fe-core/src/test/java/org/apache/doris/qe/ConnectContextTest.java
+++ b/fe/fe-core/src/test/java/org/apache/doris/qe/ConnectContextTest.java
@@ -21,6 +21,7 @@ import org.apache.doris.catalog.Env;
 import org.apache.doris.mysql.MysqlCapability;
 import org.apache.doris.mysql.MysqlChannel;
 import org.apache.doris.mysql.MysqlCommand;
+import org.apache.doris.mysql.privilege.PaloAuth;
 import org.apache.doris.thrift.TUniqueId;
 
 import mockit.Expectations;
@@ -43,6 +44,10 @@ public class ConnectContextTest {
     private Env env;
     @Mocked
     private ConnectScheduler connectScheduler;
+    @Mocked
+    private PaloAuth paloAuth;
+    @Mocked
+    private String qualifiedUser;
 
     @Before
     public void setUp() throws Exception {
@@ -166,6 +171,14 @@ public class ConnectContextTest {
         ctx.checkTimeout(now);
         Assert.assertTrue(ctx.isKilled());
 
+        // user query timeout
+        ctx.setStartTime();
+        ctx.setUserQueryTimeout(1);
+        now = ctx.getStartTime() + paloAuth.getQueryTimeout(qualifiedUser) * 1000 + 1;
+        ctx.setExecutor(executor);
+        ctx.checkTimeout(now);
+        Assert.assertTrue(ctx.isKilled());
+
         // Kill
         ctx.kill(true);
         Assert.assertTrue(ctx.isKilled());