[Enhancement] (nereids)implement CreateWorkloadGroupCommand in nereids (#44970)

Issue Number: close #42589

Author: msridhar78

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

fe/fe-core/src/main/java/org/apache/doris/nereids/parser/LogicalPlanBuilder.java

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

fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/PlanType.java

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

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<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;
+    }
+}

fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/visitor/CommandVisitor.java

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

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<String, String> properties = Maps.newHashMap();
         // 100 is cgroup v2 default cpu_share value

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;")
-}
+}