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 c160d7e77b916f..6a645be18e6135 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 @@ -76,6 +76,7 @@ import com.google.common.collect.Maps; import com.google.common.collect.Sets; import io.netty.util.concurrent.FastThreadLocal; +import org.apache.commons.lang3.StringUtils; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; import org.json.JSONObject; @@ -1028,15 +1029,75 @@ public String getCurrentConnectedFEIp() { public int getExecTimeout() { if (executor != null && executor.isSyncLoadKindStmt()) { // particular for insert stmt, we can expand other type of timeout in the same way - return Math.max(sessionVariable.getInsertTimeoutS(), sessionVariable.getQueryTimeoutS()); + return Math.max(getInsertTimeoutS(), getQueryTimeoutS()); } else if (executor != null && executor.isAnalyzeStmt()) { return sessionVariable.getAnalyzeTimeoutS(); } else { // normal query stmt - return sessionVariable.getQueryTimeoutS(); + return getQueryTimeoutS(); } } + /** + * First, retrieve from the user's attributes. If not, retrieve from the session variable + * + * @return insertTimeoutS + */ + public int getInsertTimeoutS() { + int userInsertTimeout = getInsertTimeoutSFromProperty(); + if (userInsertTimeout > 0) { + return userInsertTimeout; + } + return sessionVariable.getInsertTimeoutS(); + } + + private int getInsertTimeoutSFromProperty() { + if (env == null || env.getAuth() == null || StringUtils.isEmpty(getQualifiedUser())) { + return 0; + } + return env.getAuth().getInsertTimeout(getQualifiedUser()); + } + + /** + * First, retrieve from the user's attributes. If not, retrieve from the session variable + * + * @return queryTimeoutS + */ + public int getQueryTimeoutS() { + int userQueryTimeout = getQueryTimeoutSFromProperty(); + if (userQueryTimeout > 0) { + return userQueryTimeout; + } + return sessionVariable.getQueryTimeoutS(); + } + + private int getQueryTimeoutSFromProperty() { + if (env == null || env.getAuth() == null || StringUtils.isEmpty(getQualifiedUser())) { + return 0; + } + return env.getAuth().getQueryTimeout(getQualifiedUser()); + } + + /** + * First, retrieve from the user's attributes. If not, retrieve from the session variable + * + * @return maxExecMemByte + */ + public long getMaxExecMemByte() { + long userLimit = getMaxExecMemByteFromProperty(); + if (userLimit > 0) { + return userLimit; + } + return sessionVariable.getMaxExecMemByte(); + } + + private long getMaxExecMemByteFromProperty() { + if (env == null || env.getAuth() == null || StringUtils.isEmpty(getQualifiedUser())) { + return 0L; + } + return env.getAuth().getExecMemLimit(getQualifiedUser()); + } + public void setResultAttachedInfo(Map resultAttachedInfo) { this.resultAttachedInfo = resultAttachedInfo; } diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/Coordinator.java b/fe/fe-core/src/main/java/org/apache/doris/qe/Coordinator.java index 58631524024721..05594f9c235f92 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/Coordinator.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/Coordinator.java @@ -403,9 +403,7 @@ private void setFromUserProperty(ConnectContext connectContext) { this.queryOptions.setResourceLimit(resourceLimit); } // set exec mem limit - long maxExecMemByte = connectContext.getSessionVariable().getMaxExecMemByte(); - long memLimit = maxExecMemByte > 0 ? maxExecMemByte : - Env.getCurrentEnv().getAuth().getExecMemLimit(qualifiedUser); + long memLimit = connectContext.getMaxExecMemByte(); if (memLimit > 0) { // overwrite the exec_mem_limit from session variable; this.queryOptions.setMemLimit(memLimit); 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 4bb30b14a632be..6ad9532a29784b 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 @@ -23,6 +23,7 @@ import org.apache.doris.mysql.privilege.Auth; import org.apache.doris.thrift.TUniqueId; +import mockit.Expectations; import mockit.Mocked; import org.junit.Assert; import org.junit.Before; @@ -202,4 +203,73 @@ public void testThreadLocal() { Assert.assertNotNull(ConnectContext.get()); Assert.assertEquals(ctx, ConnectContext.get()); } + + @Test + public void testGetMaxExecMemByte() { + ConnectContext context = new ConnectContext(); + context.setQualifiedUser("a"); + context.setEnv(env); + long sessionValue = 2097153L; + long propertyValue = 2097154L; + // only session + context.getSessionVariable().setMaxExecMemByte(sessionValue); + long result = context.getMaxExecMemByte(); + Assert.assertEquals(sessionValue, result); + // has property + new Expectations() { + { + auth.getExecMemLimit(anyString); + minTimes = 0; + result = propertyValue; + } + }; + result = context.getMaxExecMemByte(); + Assert.assertEquals(propertyValue, result); + } + + @Test + public void testGetQueryTimeoutS() { + ConnectContext context = new ConnectContext(); + context.setQualifiedUser("a"); + context.setEnv(env); + int sessionValue = 1; + int propertyValue = 2; + // only session + context.getSessionVariable().setQueryTimeoutS(sessionValue); + long result = context.getQueryTimeoutS(); + Assert.assertEquals(sessionValue, result); + // has property + new Expectations() { + { + auth.getQueryTimeout(anyString); + minTimes = 0; + result = propertyValue; + } + }; + result = context.getQueryTimeoutS(); + Assert.assertEquals(propertyValue, result); + } + + @Test + public void testInsertQueryTimeoutS() { + ConnectContext context = new ConnectContext(); + context.setQualifiedUser("a"); + context.setEnv(env); + int sessionValue = 1; + int propertyValue = 2; + // only session + context.getSessionVariable().setInsertTimeoutS(sessionValue); + long result = context.getInsertTimeoutS(); + Assert.assertEquals(sessionValue, result); + // has property + new Expectations() { + { + auth.getInsertTimeout(anyString); + minTimes = 0; + result = propertyValue; + } + }; + result = context.getInsertTimeoutS(); + Assert.assertEquals(propertyValue, result); + } } diff --git a/regression-test/suites/account_p0/test_property_session.groovy b/regression-test/suites/account_p0/test_property_session.groovy new file mode 100644 index 00000000000000..57b2dad747b123 --- /dev/null +++ b/regression-test/suites/account_p0/test_property_session.groovy @@ -0,0 +1,56 @@ +// 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. + +import org.junit.Assert; + +suite("test_property_session") { + String suiteName = "test_property_session" + String userName = "${suiteName}_user" + String pwd = 'C123_567p' + sql """drop user if exists ${userName}""" + sql """CREATE USER '${userName}' IDENTIFIED BY '${pwd}'""" + + //cloud-mode + if (isCloudMode()) { + def clusters = sql " SHOW CLUSTERS; " + assertTrue(!clusters.isEmpty()) + def validCluster = clusters[0][0] + sql """GRANT USAGE_PRIV ON CLUSTER `${validCluster}` TO ${userName}"""; + } + sql """GRANT select_PRIV ON *.*.* TO ${userName}"""; + connect(user=userName, password="${pwd}", url=context.config.jdbcUrl) { + sql """ + set query_timeout=1; + """ + test { + sql """ + select sleep(3); + """ + exception "timeout" + } + } + + // the priority of property should be higher than session + sql """set property for '${userName}' 'query_timeout' = '10';""" + connect(user=userName, password="${pwd}", url=context.config.jdbcUrl) { + sql """ + select sleep(3); + """ + } + + sql """drop user if exists ${userName}""" +}