Skip to content

Commit

Permalink
[Enhancement] (nereids)implement CreateWorkloadGroupCommand in nereids (
Browse files Browse the repository at this point in the history
#44970)

Issue Number: close #42589
  • Loading branch information
msridhar78 authored Dec 9, 2024
1 parent 90cbfc5 commit 9c4c0b3
Show file tree
Hide file tree
Showing 7 changed files with 130 additions and 9 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -183,7 +183,9 @@ supportedCreateStatement
| CREATE (EXTERNAL)? TABLE (IF NOT EXISTS)? name=multipartIdentifier
LIKE existedTable=multipartIdentifier
(WITH ROLLUP (rollupNames=identifierList)?)? #createTableLike
| CREATE ROLE (IF NOT EXISTS)? name=identifier (COMMENT STRING_LITERAL)? #createRole
| CREATE ROLE (IF NOT EXISTS)? name=identifier (COMMENT STRING_LITERAL)? #createRole
| CREATE WORKLOAD GROUP (IF NOT EXISTS)?
name=identifierOrText properties=propertyClause? #createWorkloadGroup
| CREATE ROW POLICY (IF NOT EXISTS)? name=identifier
ON table=multipartIdentifier
AS type=(RESTRICTIVE | PERMISSIVE)
Expand Down Expand Up @@ -764,8 +766,6 @@ unsupportedCreateStatement
name=identifierOrText properties=propertyClause? #createResource
| CREATE STORAGE VAULT (IF NOT EXISTS)?
name=identifierOrText properties=propertyClause? #createStorageVault
| CREATE WORKLOAD GROUP (IF NOT EXISTS)?
name=identifierOrText properties=propertyClause? #createWorkloadGroup
| CREATE WORKLOAD POLICY (IF NOT EXISTS)? name=identifierOrText
(CONDITIONS LEFT_PAREN workloadPolicyConditions RIGHT_PAREN)?
(ACTIONS LEFT_PAREN workloadPolicyActions RIGHT_PAREN)?
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -104,6 +104,7 @@
import org.apache.doris.nereids.DorisParser.CreateTableContext;
import org.apache.doris.nereids.DorisParser.CreateTableLikeContext;
import org.apache.doris.nereids.DorisParser.CreateViewContext;
import org.apache.doris.nereids.DorisParser.CreateWorkloadGroupContext;
import org.apache.doris.nereids.DorisParser.CteContext;
import org.apache.doris.nereids.DorisParser.DataTypeWithNullableContext;
import org.apache.doris.nereids.DorisParser.DateCeilContext;
Expand Down Expand Up @@ -502,6 +503,7 @@
import org.apache.doris.nereids.trees.plans.commands.CreateTableCommand;
import org.apache.doris.nereids.trees.plans.commands.CreateTableLikeCommand;
import org.apache.doris.nereids.trees.plans.commands.CreateViewCommand;
import org.apache.doris.nereids.trees.plans.commands.CreateWorkloadGroupCommand;
import org.apache.doris.nereids.trees.plans.commands.DeleteFromCommand;
import org.apache.doris.nereids.trees.plans.commands.DeleteFromUsingCommand;
import org.apache.doris.nereids.trees.plans.commands.DropCatalogRecycleBinCommand;
Expand Down Expand Up @@ -4862,6 +4864,15 @@ public LogicalPlan visitDropEncryptkey(DropEncryptkeyContext ctx) {
return new DropEncryptkeyCommand(new EncryptKeyName(nameParts), ctx.EXISTS() != null);
}

@Override
public LogicalPlan visitCreateWorkloadGroup(CreateWorkloadGroupContext ctx) {
String workloadGroupName = stripQuotes(ctx.name.getText());
boolean ifNotExists = ctx.EXISTS() != null;
Map<String, String> properties = ctx.propertyClause() != null
? Maps.newHashMap(visitPropertyClause(ctx.propertyClause())) : Maps.newHashMap();
return new CreateWorkloadGroupCommand(workloadGroupName, ifNotExists, properties);
}

@Override
public LogicalPlan visitDropFile(DropFileContext ctx) {
String dbName = null;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -247,6 +247,7 @@ public enum PlanType {
RECOVER_PARTITION_COMMAND,
REPLAY_COMMAND,
CREATE_ENCRYPTKEY_COMMAND,
CREATE_WORKLOAD_GROUP_COMMAND,
CREATE_FILE_COMMAND,
CREATE_ROUTINE_LOAD_COMMAND
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,100 @@
// 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.plans.commands;

import org.apache.doris.analysis.StmtType;
import org.apache.doris.catalog.Env;
import org.apache.doris.common.AnalysisException;
import org.apache.doris.common.ErrorCode;
import org.apache.doris.common.ErrorReport;
import org.apache.doris.common.FeNameFormat;
import org.apache.doris.mysql.privilege.PrivPredicate;
import org.apache.doris.nereids.trees.plans.PlanType;
import org.apache.doris.nereids.trees.plans.visitor.PlanVisitor;
import org.apache.doris.qe.ConnectContext;
import org.apache.doris.qe.StmtExecutor;
import org.apache.doris.resource.workloadgroup.WorkloadGroup;
import org.apache.doris.resource.workloadgroup.WorkloadGroupMgr;

import org.apache.commons.lang3.StringUtils;

import java.util.Map;

/**
* Create workload group command
*/
public class CreateWorkloadGroupCommand extends Command implements ForwardWithSync {
private final boolean ifNotExists;
private final String workloadGroupName;
private final Map<String, String> properties;

/**
* Constructor for CreateWorkloadGroupCommand
*/
public CreateWorkloadGroupCommand(String workloadGroupName, boolean ifNotExists, Map<String, String> properties) {
super(PlanType.CREATE_WORKLOAD_GROUP_COMMAND);
this.workloadGroupName = workloadGroupName;
this.ifNotExists = ifNotExists;
this.properties = properties;
}

private void validate(ConnectContext ctx) throws AnalysisException {
// check auth
if (!Env.getCurrentEnv().getAccessManager().checkGlobalPriv(ConnectContext.get(), PrivPredicate.ADMIN)) {
ErrorReport.reportAnalysisException(ErrorCode.ERR_SPECIFIC_ACCESS_DENIED_ERROR, "ADMIN");
}

// check name
FeNameFormat.checkWorkloadGroupName(workloadGroupName);

if (properties == null || properties.isEmpty()) {
throw new AnalysisException("Workload Group properties can't be empty");
}

if (properties.containsKey(WorkloadGroup.INTERNAL_TYPE)) {
throw new AnalysisException(WorkloadGroup.INTERNAL_TYPE + " can not be create or modified ");
}

String tagStr = properties.get(WorkloadGroup.TAG);
if (!StringUtils.isEmpty(tagStr) && (WorkloadGroupMgr.DEFAULT_GROUP_NAME.equals(workloadGroupName)
|| WorkloadGroupMgr.INTERNAL_GROUP_NAME.equals(workloadGroupName))) {
throw new AnalysisException(
WorkloadGroupMgr.INTERNAL_GROUP_NAME + " and " + WorkloadGroupMgr.DEFAULT_GROUP_NAME
+ " group can not set tag");
}
}

@Override
public void run(ConnectContext ctx, StmtExecutor executor) throws Exception {
validate(ctx);
// Create workload group
WorkloadGroup workloadGroup = WorkloadGroup.create(workloadGroupName, properties);
WorkloadGroupMgr workloadGroupMgr = Env.getCurrentEnv().getWorkloadGroupMgr();
workloadGroupMgr.createWorkloadGroup(workloadGroup, ifNotExists);
}

@Override
public <R, C> R accept(PlanVisitor<R, C> visitor, C context) {
return visitor.visitCreateWorkloadGroupCommand(this, context);
}

@Override
public StmtType stmtType() {
return StmtType.CREATE;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,7 @@
import org.apache.doris.nereids.trees.plans.commands.CreateTableCommand;
import org.apache.doris.nereids.trees.plans.commands.CreateTableLikeCommand;
import org.apache.doris.nereids.trees.plans.commands.CreateViewCommand;
import org.apache.doris.nereids.trees.plans.commands.CreateWorkloadGroupCommand;
import org.apache.doris.nereids.trees.plans.commands.DeleteFromCommand;
import org.apache.doris.nereids.trees.plans.commands.DeleteFromUsingCommand;
import org.apache.doris.nereids.trees.plans.commands.DropCatalogRecycleBinCommand;
Expand Down Expand Up @@ -562,6 +563,10 @@ default R visitShowTableIdCommand(ShowTableIdCommand showTableIdCommand, C conte
return visitCommand(showTableIdCommand, context);
}

default R visitCreateWorkloadGroupCommand(CreateWorkloadGroupCommand createWorkloadGroupCommand, C context) {
return visitCommand(createWorkloadGroupCommand, context);
}

default R visitShowEncryptKeysCommand(ShowEncryptKeysCommand showEncryptKeysCommand, C context) {
return visitCommand(showEncryptKeysCommand, context);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -357,13 +357,12 @@ private String getWorkloadGroupNameAndCheckPriv(ConnectContext context) throws A
return groupName;
}

public void createWorkloadGroup(CreateWorkloadGroupStmt stmt) throws DdlException {
WorkloadGroup workloadGroup = WorkloadGroup.create(stmt.getWorkloadGroupName(), stmt.getProperties());
public void createWorkloadGroup(WorkloadGroup workloadGroup, boolean isIfNotExists) throws DdlException {
String workloadGroupName = workloadGroup.getName();
writeLock();
try {
if (nameToWorkloadGroup.containsKey(workloadGroupName)) {
if (stmt.isIfNotExists()) {
if (isIfNotExists) {
return;
}
throw new DdlException("workload group " + workloadGroupName + " already exist");
Expand All @@ -382,6 +381,11 @@ public void createWorkloadGroup(CreateWorkloadGroupStmt stmt) throws DdlExceptio
LOG.info("Create workload group success: {}", workloadGroup);
}

public void createWorkloadGroup(CreateWorkloadGroupStmt stmt) throws DdlException {
WorkloadGroup workloadGroup = WorkloadGroup.create(stmt.getWorkloadGroupName(), stmt.getProperties());
createWorkloadGroup(workloadGroup, stmt.isIfNotExists());
}

public void createInternalWorkloadGroup() {
Map<String, String> properties = Maps.newHashMap();
// 100 is cgroup v2 default cpu_share value
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -18,10 +18,10 @@
suite("test_nereids_workload_test") {
sql "drop workload group if exists test_nereids_wg1;"
sql "drop workload group if exists test_nereids_wg2;"
sql "create workload group test_nereids_wg1 properties('cpu_share'='1024');"
sql "create workload group test_nereids_wg2 properties('cpu_share'='1024');"
checkNereidsExecute("create workload group test_nereids_wg1 properties('cpu_share'='1024');")
checkNereidsExecute("create workload group test_nereids_wg2 properties('cpu_share'='1024');")
qt_check_workload_check1("select NAME from information_schema.workload_groups where NAME='test_nereids_wg1';")
checkNereidsExecute("drop workload group test_nereids_wg1;")
qt_check_workload_check2("select NAME from information_schema.workload_groups where NAME='test_nereids_wg1';")
checkNereidsExecute("drop workload group if exists test_nereids_wg2;")
}
}

0 comments on commit 9c4c0b3

Please sign in to comment.