diff --git a/fe/fe-core/src/main/antlr4/org/apache/doris/nereids/DorisParser.g4 b/fe/fe-core/src/main/antlr4/org/apache/doris/nereids/DorisParser.g4 index b51b342e24b895a..c88ad1baa22a60e 100644 --- a/fe/fe-core/src/main/antlr4/org/apache/doris/nereids/DorisParser.g4 +++ b/fe/fe-core/src/main/antlr4/org/apache/doris/nereids/DorisParser.g4 @@ -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) @@ -763,8 +765,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)? diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/parser/LogicalPlanBuilder.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/parser/LogicalPlanBuilder.java index 7804959aa79f9ee..1d1e7c6a6bc536d 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/parser/LogicalPlanBuilder.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/parser/LogicalPlanBuilder.java @@ -101,6 +101,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; @@ -492,6 +493,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; @@ -4810,6 +4812,15 @@ public LogicalPlan visitDropEncryptkey(DropEncryptkeyContext ctx) { return new DropEncryptkeyCommand(new EncryptKeyName(nameParts), ctx.EXISTS() != null); } + @Override + public LogicalPlan visitCreateWorkloadGroup(CreateWorkloadGroupContext ctx) { + String workloadGroupName = ctx.name.getText(); + boolean ifNotExists = ctx.EXISTS() != null; + Map 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; diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/PlanType.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/PlanType.java index 48e64acaab63771..21967e17d1b99a5 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/PlanType.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/PlanType.java @@ -240,6 +240,7 @@ public enum PlanType { RECOVER_PARTITION_COMMAND, REPLAY_COMMAND, CREATE_ENCRYPTKEY_COMMAND, + CREATE_WORKLOAD_GROUP_COMMAND, CREATE_FILE_COMMAND, CREATE_ROUTINE_LOAD_COMMAND } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/CreateWorkloadGroupCommand.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/CreateWorkloadGroupCommand.java new file mode 100644 index 000000000000000..ad4db8cbfdfae82 --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/CreateWorkloadGroupCommand.java @@ -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 properties; + + /** + * Constructor for CreateWorkloadGroupCommand + */ + public CreateWorkloadGroupCommand(String workloadGroupName, boolean ifNotExists, Map 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 accept(PlanVisitor visitor, C context) { + return visitor.visitCreateWorkloadGroupCommand(this, context); + } + + @Override + public StmtType stmtType() { + return StmtType.CREATE; + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/visitor/CommandVisitor.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/visitor/CommandVisitor.java index 8561c934ec976fb..17be050d8ba1232 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/visitor/CommandVisitor.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/visitor/CommandVisitor.java @@ -44,6 +44,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; @@ -534,6 +535,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); } diff --git a/fe/fe-core/src/main/java/org/apache/doris/resource/workloadgroup/WorkloadGroupMgr.java b/fe/fe-core/src/main/java/org/apache/doris/resource/workloadgroup/WorkloadGroupMgr.java index 22cbcf4a0a15023..94695cc8d5b8ac1 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/resource/workloadgroup/WorkloadGroupMgr.java +++ b/fe/fe-core/src/main/java/org/apache/doris/resource/workloadgroup/WorkloadGroupMgr.java @@ -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"); @@ -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 properties = Maps.newHashMap(); // 100 is cgroup v2 default cpu_share value diff --git a/regression-test/suites/workload_manager_p0/test_nereids_workload_test.groovy b/regression-test/suites/workload_manager_p0/test_nereids_workload_test.groovy index 3f7c55f5c99cf91..680f5a28ca982ef 100644 --- a/regression-test/suites/workload_manager_p0/test_nereids_workload_test.groovy +++ b/regression-test/suites/workload_manager_p0/test_nereids_workload_test.groovy @@ -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;") -} \ No newline at end of file +}