database/doris
2
0
Fork 0
Code Issues Pull Requests Actions 3 Packages Projects Releases Wiki Activity

[feature-wip](workload-group) Support setting user default workload group (#20180)

Browse Source 
Issue Number: close #xxx

SET PROPERTY 'default_workload_group' = 'group_name';
This commit is contained in:
luozenglin
2023-05-31 09:18:25 +08:00
committed by GitHub
parent aae04d9680
commit 8a54be3318
9 changed files with 70 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
fe/fe-core/src/main/java/org/apache/doris/mysql/privilege/Auth.java
View File

@ -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<String> allDomains) {
readLock();
try {

12
fe/fe-core/src/main/java/org/apache/doris/mysql/privilege/CommonUserProperties.java
View File

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

16
fe/fe-core/src/main/java/org/apache/doris/mysql/privilege/UserProperty.java
View File

@ -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<Pattern> 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<String, DppConfig> 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));

9
fe/fe-core/src/main/java/org/apache/doris/mysql/privilege/UserPropertyMgr.java
View File

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

2
fe/fe-core/src/main/java/org/apache/doris/qe/SessionVariable.java
View File

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

2
fe/fe-core/src/main/java/org/apache/doris/qe/StmtExecutor.java
View File

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

8
fe/fe-core/src/main/java/org/apache/doris/resource/resourcegroup/ResourceGroupMgr.java
View File

@ -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<TPipelineResourceGroup> getResourceGroup(String groupName) throws UserException {
public List<TPipelineResourceGroup> getResourceGroup(ConnectContext context) throws UserException {
String groupName = context.getSessionVariable().getResourceGroup();
if (Strings.isNullOrEmpty(groupName)) {
groupName = Env.getCurrentEnv().getAuth().getWorkloadGroup(context.getQualifiedUser());
}
List<TPipelineResourceGroup> resourceGroups = Lists.newArrayList();
readLock();
try {

2
fe/fe-core/src/test/java/org/apache/doris/planner/ResourceTagQueryTest.java
View File

@ -278,7 +278,7 @@ public class ResourceTagQueryTest {
Assert.assertEquals(1000000, execMemLimit);
List<List<String>> 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 {

19
fe/fe-core/src/test/java/org/apache/doris/resource/resourcegroup/ResourceGroupMgrTest.java
View File

@ -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<String, String> 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<TPipelineResourceGroup> tResourceGroups1 = resourceGroupMgr.getResourceGroup(name1);
context.getSessionVariable().setResourceGroup(name1);
List<TPipelineResourceGroup> 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<String, String> 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<String, String> properties = Maps.newHashMap();
String name = "g1";
@ -202,7 +210,8 @@ public class ResourceGroupMgrTest {
AlterResourceGroupStmt stmt2 = new AlterResourceGroupStmt(name, newProperties);
resourceGroupMgr.alterResourceGroup(stmt2);
List<TPipelineResourceGroup> tResourceGroups = resourceGroupMgr.getResourceGroup(name);
context.getSessionVariable().setResourceGroup(name);
List<TPipelineResourceGroup> tResourceGroups = resourceGroupMgr.getResourceGroup(context);
Assert.assertEquals(1, tResourceGroups.size());
TPipelineResourceGroup tResourceGroup1 = tResourceGroups.get(0);
Assert.assertEquals(tResourceGroup1.getProperties().get(ResourceGroup.CPU_SHARE), "5");