diff --git a/fe/fe-core/src/main/java/org/apache/doris/mysql/privilege/Auth.java b/fe/fe-core/src/main/java/org/apache/doris/mysql/privilege/Auth.java index 9c3afdfd71..64e4c0b055 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/mysql/privilege/Auth.java +++ b/fe/fe-core/src/main/java/org/apache/doris/mysql/privilege/Auth.java @@ -943,6 +943,15 @@ public class Auth implements Writable { } } + public String getWorkloadGroup(String qualifiedUser) { + readLock(); + try { + return propertyMgr.getWorkloadGroup(qualifiedUser); + } finally { + readUnlock(); + } + } + public void getAllDomains(Set allDomains) { readLock(); try { 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 27e06be196..b1e0dca03b 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 @@ -21,6 +21,7 @@ import org.apache.doris.common.io.Text; import org.apache.doris.common.io.Writable; import org.apache.doris.persist.gson.GsonUtils; import org.apache.doris.resource.Tag; +import org.apache.doris.resource.resourcegroup.ResourceGroupMgr; import com.google.common.collect.Sets; import com.google.gson.annotations.SerializedName; @@ -57,6 +58,9 @@ public class CommonUserProperties implements Writable { @SerializedName("insertTimeout") private int insertTimeout = -1; + @SerializedName("workloadGroup") + private String workloadGroup = ResourceGroupMgr.DEFAULT_GROUP_NAME; + private String[] sqlBlockRulesSplit = {}; long getMaxConn() { @@ -133,6 +137,14 @@ public class CommonUserProperties implements Writable { this.insertTimeout = insertTimeout; } + public String getWorkloadGroup() { + return workloadGroup; + } + + public void setWorkloadGroup(String workloadGroup) { + this.workloadGroup = workloadGroup; + } + 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/UserProperty.java b/fe/fe-core/src/main/java/org/apache/doris/mysql/privilege/UserProperty.java index 10ee95e406..e312b21520 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 @@ -69,6 +69,8 @@ public class UserProperty implements Writable { private static final String PROP_QUOTA = "quota"; private static final String PROP_DEFAULT_LOAD_CLUSTER = "default_load_cluster"; + private static final String PROP_WORKLOAD_GROUP = "default_workload_group"; + // for system user public static final Set ADVANCED_PROPERTIES = Sets.newHashSet(); // for normal user @@ -114,6 +116,7 @@ public class UserProperty implements Writable { COMMON_PROPERTIES.add(Pattern.compile("^" + PROP_DEFAULT_LOAD_CLUSTER + "$", Pattern.CASE_INSENSITIVE)); COMMON_PROPERTIES.add(Pattern.compile("^" + PROP_LOAD_CLUSTER + "." + DppConfig.CLUSTER_NAME_REGEX + ".", Pattern.CASE_INSENSITIVE)); + COMMON_PROPERTIES.add(Pattern.compile("^" + PROP_WORKLOAD_GROUP + "$", Pattern.CASE_INSENSITIVE)); } public UserProperty() { @@ -151,6 +154,10 @@ public class UserProperty implements Writable { return commonProperties.getCpuResourceLimit(); } + public String getWorkloadGroup() { + return commonProperties.getWorkloadGroup(); + } + @Deprecated public WhiteList getWhiteList() { return whiteList; @@ -174,6 +181,7 @@ public class UserProperty implements Writable { long execMemLimit = this.commonProperties.getExecMemLimit(); int queryTimeout = this.commonProperties.getQueryTimeout(); int insertTimeout = this.commonProperties.getInsertTimeout(); + String workloadGroup = this.commonProperties.getWorkloadGroup(); String newDefaultLoadCluster = defaultLoadCluster; Map newDppConfigs = Maps.newHashMap(clusterToDppConfig); @@ -293,6 +301,11 @@ public class UserProperty implements Writable { } catch (NumberFormatException e) { throw new DdlException(PROP_USER_INSERT_TIMEOUT + " is not number"); } + } else if (keyArr[0].equalsIgnoreCase(PROP_WORKLOAD_GROUP)) { + if (keyArr.length != 1) { + throw new DdlException(PROP_WORKLOAD_GROUP + " format error"); + } + workloadGroup = value; } else { throw new DdlException("Unknown user property(" + key + ")"); } @@ -307,6 +320,7 @@ public class UserProperty implements Writable { this.commonProperties.setExecMemLimit(execMemLimit); this.commonProperties.setQueryTimeout(queryTimeout); this.commonProperties.setInsertTimeout(insertTimeout); + this.commonProperties.setWorkloadGroup(workloadGroup); if (newDppConfigs.containsKey(newDefaultLoadCluster)) { defaultLoadCluster = newDefaultLoadCluster; } else { @@ -437,6 +451,8 @@ public class UserProperty implements Writable { // resource tag result.add(Lists.newArrayList(PROP_RESOURCE_TAGS, Joiner.on(", ").join(commonProperties.getResourceTags()))); + result.add(Lists.newArrayList(PROP_WORKLOAD_GROUP, String.valueOf(commonProperties.getWorkloadGroup()))); + // load cluster if (defaultLoadCluster != null) { result.add(Lists.newArrayList(PROP_DEFAULT_LOAD_CLUSTER, defaultLoadCluster)); 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 ff68cd028e..4458a36a79 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 @@ -173,6 +173,15 @@ public class UserPropertyMgr implements Writable { return existProperty.getExecMemLimit(); } + public String getWorkloadGroup(String qualifiedUser) { + UserProperty existProperty = propertyMap.get(qualifiedUser); + existProperty = getLdapPropertyIfNull(qualifiedUser, existProperty); + if (existProperty == null) { + return null; + } + return existProperty.getWorkloadGroup(); + } + private UserProperty getLdapPropertyIfNull(String qualifiedUser, UserProperty existProperty) { if (existProperty == null && Env.getCurrentEnv().getAuth().getLdapManager().doesUserExist(qualifiedUser)) { return LDAP_PROPERTY; diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/SessionVariable.java b/fe/fe-core/src/main/java/org/apache/doris/qe/SessionVariable.java index c0e31bf30c..aadb76b158 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/SessionVariable.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/SessionVariable.java @@ -405,7 +405,7 @@ public class SessionVariable implements Serializable, Writable { public long sqlMode = SqlModeHelper.MODE_DEFAULT; @VariableMgr.VarAttr(name = RESOURCE_VARIABLE) - public String resourceGroup = "normal"; + public String resourceGroup = ""; // this is used to make mysql client happy @VariableMgr.VarAttr(name = AUTO_COMMIT) diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/StmtExecutor.java b/fe/fe-core/src/main/java/org/apache/doris/qe/StmtExecutor.java index 458102dc6d..a890b43d4e 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/StmtExecutor.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/StmtExecutor.java @@ -1097,7 +1097,7 @@ public class StmtExecutor { if (parsedStmt instanceof QueryStmt && Config.enable_resource_group && context.sessionVariable.enablePipelineEngine()) { analyzer.setResourceGroups(analyzer.getEnv().getResourceGroupMgr() - .getResourceGroup(context.sessionVariable.resourceGroup)); + .getResourceGroup(context)); } } profile.getSummaryProfile().setQueryAnalysisFinishTime(); diff --git a/fe/fe-core/src/main/java/org/apache/doris/resource/resourcegroup/ResourceGroupMgr.java b/fe/fe-core/src/main/java/org/apache/doris/resource/resourcegroup/ResourceGroupMgr.java index a11907f34c..d764c64ce8 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/resource/resourcegroup/ResourceGroupMgr.java +++ b/fe/fe-core/src/main/java/org/apache/doris/resource/resourcegroup/ResourceGroupMgr.java @@ -32,8 +32,10 @@ import org.apache.doris.common.proc.ProcResult; import org.apache.doris.persist.DropResourceGroupOperationLog; import org.apache.doris.persist.gson.GsonPostProcessable; import org.apache.doris.persist.gson.GsonUtils; +import org.apache.doris.qe.ConnectContext; import org.apache.doris.thrift.TPipelineResourceGroup; +import com.google.common.base.Strings; import com.google.common.collect.ImmutableList; import com.google.common.collect.Lists; import com.google.common.collect.Maps; @@ -99,7 +101,11 @@ public class ResourceGroupMgr implements Writable, GsonPostProcessable { } } - public List getResourceGroup(String groupName) throws UserException { + public List getResourceGroup(ConnectContext context) throws UserException { + String groupName = context.getSessionVariable().getResourceGroup(); + if (Strings.isNullOrEmpty(groupName)) { + groupName = Env.getCurrentEnv().getAuth().getWorkloadGroup(context.getQualifiedUser()); + } List resourceGroups = Lists.newArrayList(); readLock(); try { 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 0f15ac486b..ef4a58e534 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 @@ -278,7 +278,7 @@ public class ResourceTagQueryTest { Assert.assertEquals(1000000, execMemLimit); List> userProps = Env.getCurrentEnv().getAuth().getUserProperties(Auth.ROOT_USER); - Assert.assertEquals(9, userProps.size()); + Assert.assertEquals(10, userProps.size()); } private void checkTableReplicaAllocation(OlapTable tbl) throws InterruptedException { diff --git a/fe/fe-core/src/test/java/org/apache/doris/resource/resourcegroup/ResourceGroupMgrTest.java b/fe/fe-core/src/test/java/org/apache/doris/resource/resourcegroup/ResourceGroupMgrTest.java index f6562f584a..562e71017d 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/resource/resourcegroup/ResourceGroupMgrTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/resource/resourcegroup/ResourceGroupMgrTest.java @@ -25,6 +25,7 @@ import org.apache.doris.common.Config; import org.apache.doris.common.DdlException; import org.apache.doris.common.UserException; import org.apache.doris.persist.EditLog; +import org.apache.doris.qe.ConnectContext; import org.apache.doris.thrift.TPipelineResourceGroup; import com.google.common.collect.Maps; @@ -128,6 +129,7 @@ public class ResourceGroupMgrTest { @Test public void testGetResourceGroup() throws UserException { Config.enable_resource_group = true; + ConnectContext context = new ConnectContext(); ResourceGroupMgr resourceGroupMgr = new ResourceGroupMgr(); Map properties1 = Maps.newHashMap(); properties1.put(ResourceGroup.CPU_SHARE, "10"); @@ -135,14 +137,16 @@ public class ResourceGroupMgrTest { String name1 = "g1"; CreateResourceGroupStmt stmt1 = new CreateResourceGroupStmt(false, name1, properties1); resourceGroupMgr.createResourceGroup(stmt1); - List tResourceGroups1 = resourceGroupMgr.getResourceGroup(name1); + context.getSessionVariable().setResourceGroup(name1); + List tResourceGroups1 = resourceGroupMgr.getResourceGroup(context); Assert.assertEquals(1, tResourceGroups1.size()); TPipelineResourceGroup tResourceGroup1 = tResourceGroups1.get(0); Assert.assertEquals(name1, tResourceGroup1.getName()); Assert.assertTrue(tResourceGroup1.getProperties().containsKey(ResourceGroup.CPU_SHARE)); try { - resourceGroupMgr.getResourceGroup("g2"); + context.getSessionVariable().setResourceGroup("g2"); + resourceGroupMgr.getResourceGroup(context); Assert.fail(); } catch (UserException e) { Assert.assertTrue(e.getMessage().contains("does not exist")); @@ -152,6 +156,7 @@ public class ResourceGroupMgrTest { @Test public void testDropResourceGroup() throws UserException { Config.enable_resource_group = true; + ConnectContext context = new ConnectContext(); ResourceGroupMgr resourceGroupMgr = new ResourceGroupMgr(); Map properties = Maps.newHashMap(); properties.put(ResourceGroup.CPU_SHARE, "10"); @@ -159,12 +164,14 @@ public class ResourceGroupMgrTest { String name = "g1"; CreateResourceGroupStmt createStmt = new CreateResourceGroupStmt(false, name, properties); resourceGroupMgr.createResourceGroup(createStmt); - Assert.assertEquals(1, resourceGroupMgr.getResourceGroup(name).size()); + context.getSessionVariable().setResourceGroup(name); + Assert.assertEquals(1, resourceGroupMgr.getResourceGroup(context).size()); DropResourceGroupStmt dropStmt = new DropResourceGroupStmt(false, name); resourceGroupMgr.dropResourceGroup(dropStmt); try { - resourceGroupMgr.getResourceGroup(name); + context.getSessionVariable().setResourceGroup(name); + resourceGroupMgr.getResourceGroup(context); Assert.fail(); } catch (UserException e) { Assert.assertTrue(e.getMessage().contains("does not exist")); @@ -181,6 +188,7 @@ public class ResourceGroupMgrTest { @Test public void testAlterResourceGroup() throws UserException { Config.enable_resource_group = true; + ConnectContext context = new ConnectContext(); ResourceGroupMgr resourceGroupMgr = new ResourceGroupMgr(); Map properties = Maps.newHashMap(); String name = "g1"; @@ -202,7 +210,8 @@ public class ResourceGroupMgrTest { AlterResourceGroupStmt stmt2 = new AlterResourceGroupStmt(name, newProperties); resourceGroupMgr.alterResourceGroup(stmt2); - List tResourceGroups = resourceGroupMgr.getResourceGroup(name); + context.getSessionVariable().setResourceGroup(name); + List tResourceGroups = resourceGroupMgr.getResourceGroup(context); Assert.assertEquals(1, tResourceGroups.size()); TPipelineResourceGroup tResourceGroup1 = tResourceGroups.get(0); Assert.assertEquals(tResourceGroup1.getProperties().get(ResourceGroup.CPU_SHARE), "5");