[Improment]Add schema table workload_group_privileges

be/src/exec/schema_scanner.cpp

@@ -48,6 +48,7 @@
 #include "exec/schema_scanner/schema_user_scanner.h"
 #include "exec/schema_scanner/schema_variables_scanner.h"
 #include "exec/schema_scanner/schema_views_scanner.h"
+#include "exec/schema_scanner/schema_workload_group_privileges.h"
 #include "exec/schema_scanner/schema_workload_groups_scanner.h"
 #include "exec/schema_scanner/schema_workload_sched_policy_scanner.h"
 #include "olap/hll.h"
@@ -170,6 +171,8 @@ std::unique_ptr<SchemaScanner> SchemaScanner::create(TSchemaTableType::type type
         return SchemaWorkloadSchedulePolicyScanner::create_unique();
     case TSchemaTableType::SCH_TABLE_OPTIONS:
         return SchemaTableOptionsScanner::create_unique();
+    case TSchemaTableType::SCH_WORKLOAD_GROUP_PRIVILEGES:
+        return SchemaWorkloadGroupPrivilegesScanner::create_unique();
     default:
         return SchemaDummyScanner::create_unique();
         break;

be/src/exec/schema_scanner/schema_workload_group_privileges.cpp

@@ -0,0 +1,136 @@
+// 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.
+
+#include "exec/schema_scanner/schema_workload_group_privileges.h"
+
+#include "runtime/client_cache.h"
+#include "runtime/exec_env.h"
+#include "runtime/runtime_state.h"
+#include "util/thrift_rpc_helper.h"
+#include "vec/common/string_ref.h"
+#include "vec/core/block.h"
+#include "vec/data_types/data_type_factory.hpp"
+
+namespace doris {
+std::vector<SchemaScanner::ColumnDesc> SchemaWorkloadGroupPrivilegesScanner::_s_tbls_columns = {
+        {"GRANTEE", TYPE_VARCHAR, sizeof(StringRef), true},
+        {"WORKLOAD_GROUP_NAME", TYPE_VARCHAR, sizeof(StringRef), true},
+        {"PRIVILEGE_TYPE", TYPE_VARCHAR, sizeof(StringRef), true},
+        {"IS_GRANTABLE", TYPE_VARCHAR, sizeof(StringRef), true},
+};
+
+SchemaWorkloadGroupPrivilegesScanner::SchemaWorkloadGroupPrivilegesScanner()
+        : SchemaScanner(_s_tbls_columns, TSchemaTableType::SCH_WORKLOAD_GROUPS) {}
+
+SchemaWorkloadGroupPrivilegesScanner::~SchemaWorkloadGroupPrivilegesScanner() {}
+
+Status SchemaWorkloadGroupPrivilegesScanner::start(RuntimeState* state) {
+    _block_rows_limit = state->batch_size();
+    _rpc_timeout = state->execution_timeout() * 1000;
+    return Status::OK();
+}
+
+Status SchemaWorkloadGroupPrivilegesScanner::_get_workload_group_privs_block_from_fe() {
+    TNetworkAddress master_addr = ExecEnv::GetInstance()->master_info()->network_address;
+
+    TSchemaTableRequestParams schema_table_request_params;
+    for (int i = 0; i < _s_tbls_columns.size(); i++) {
+        schema_table_request_params.__isset.columns_name = true;
+        schema_table_request_params.columns_name.emplace_back(_s_tbls_columns[i].name);
+    }
+    schema_table_request_params.__set_current_user_ident(*_param->common_param->current_user_ident);
+
+    TFetchSchemaTableDataRequest request;
+    request.__set_schema_table_name(TSchemaTableName::WORKLOAD_GROUP_PRIVILEGES);
+    request.__set_schema_table_params(schema_table_request_params);
+
+    TFetchSchemaTableDataResult result;
+
+    RETURN_IF_ERROR(ThriftRpcHelper::rpc<FrontendServiceClient>(
+            master_addr.hostname, master_addr.port,
+            [&request, &result](FrontendServiceConnection& client) {
+                client->fetchSchemaTableData(result, request);
+            },
+            _rpc_timeout));
+
+    Status status(Status::create(result.status));
+    if (!status.ok()) {
+        LOG(WARNING) << "fetch workload group privileges from FE failed, errmsg=" << status;
+        return status;
+    }
+    std::vector<TRow> result_data = result.data_batch;
+
+    _workload_groups_privs_block = vectorized::Block::create_unique();
+    for (int i = 0; i < _s_tbls_columns.size(); ++i) {
+        TypeDescriptor descriptor(_s_tbls_columns[i].type);
+        auto data_type = vectorized::DataTypeFactory::instance().create_data_type(descriptor, true);
+        _workload_groups_privs_block->insert(vectorized::ColumnWithTypeAndName(
+                data_type->create_column(), data_type, _s_tbls_columns[i].name));
+    }
+
+    if (result_data.size() > 0) {
+        int col_size = result_data[0].column_value.size();
+        if (col_size != _s_tbls_columns.size()) {
+            return Status::InternalError<false>(
+                    "workload group privileges schema is not match for FE and BE");
+        }
+    }
+
+    _workload_groups_privs_block->reserve(result_data.size());
+
+    for (int i = 0; i < result_data.size(); i++) {
+        TRow row = result_data[i];
+
+        for (int j = 0; j < _s_tbls_columns.size(); j++) {
+            RETURN_IF_ERROR(insert_block_column(row.column_value[j], j,
+                                                _workload_groups_privs_block.get(),
+                                                _s_tbls_columns[j].type));
+        }
+    }
+    return Status::OK();
+}
+
+Status SchemaWorkloadGroupPrivilegesScanner::get_next_block(vectorized::Block* block, bool* eos) {
+    if (!_is_init) {
+        return Status::InternalError("Used before initialized.");
+    }
+
+    if (nullptr == block || nullptr == eos) {
+        return Status::InternalError("input pointer is nullptr.");
+    }
+
+    if (_workload_groups_privs_block == nullptr) {
+        RETURN_IF_ERROR(_get_workload_group_privs_block_from_fe());
+        _total_rows = _workload_groups_privs_block->rows();
+    }
+
+    if (_row_idx == _total_rows) {
+        *eos = true;
+        return Status::OK();
+    }
+
+    int current_batch_rows = std::min(_block_rows_limit, _total_rows - _row_idx);
+    vectorized::MutableBlock mblock = vectorized::MutableBlock::build_mutable_block(block);
+    RETURN_IF_ERROR(
+            mblock.add_rows(_workload_groups_privs_block.get(), _row_idx, current_batch_rows));
+    _row_idx += current_batch_rows;
+
+    *eos = _row_idx == _total_rows;
+    return Status::OK();
+}
+
+} // namespace doris

be/src/exec/schema_scanner/schema_workload_group_privileges.h

@@ -0,0 +1,52 @@
+// 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.
+
+#pragma once
+
+#include <vector>
+
+#include "common/status.h"
+#include "exec/schema_scanner.h"
+
+namespace doris {
+class RuntimeState;
+namespace vectorized {
+class Block;
+} // namespace vectorized
+
+class SchemaWorkloadGroupPrivilegesScanner : public SchemaScanner {
+    ENABLE_FACTORY_CREATOR(SchemaWorkloadGroupPrivilegesScanner);
+
+public:
+    SchemaWorkloadGroupPrivilegesScanner();
+    ~SchemaWorkloadGroupPrivilegesScanner() override;
+
+    Status start(RuntimeState* state) override;
+    Status get_next_block(vectorized::Block* block, bool* eos) override;
+
+    static std::vector<SchemaScanner::ColumnDesc> _s_tbls_columns;
+
+private:
+    Status _get_workload_group_privs_block_from_fe();
+
+    int _block_rows_limit = 4096;
+    int _row_idx = 0;
+    int _total_rows = 0;
+    std::unique_ptr<vectorized::Block> _workload_groups_privs_block = nullptr;
+    int _rpc_timeout = 3000;
+};
+}; // namespace doris

fe/fe-core/src/main/java/org/apache/doris/analysis/SchemaTableType.java

@@ -78,7 +78,9 @@ public enum SchemaTableType {
     SCH_WORKLOAD_POLICY("WORKLOAD_POLICY", "WORKLOAD_POLICY",
             TSchemaTableType.SCH_WORKLOAD_POLICY),
     SCH_TABLE_OPTIONS("TABLE_OPTIONS", "TABLE_OPTIONS",
-            TSchemaTableType.SCH_TABLE_OPTIONS);
+            TSchemaTableType.SCH_TABLE_OPTIONS),
+    SCH_WORKLOAD_GROUP_PRIVILEGES("WORKLOAD_GROUP_PRIVILEGES",
+            "WORKLOAD_GROUP_PRIVILEGES", TSchemaTableType.SCH_WORKLOAD_GROUP_PRIVILEGES);
 
     private static final String dbName = "INFORMATION_SCHEMA";
     private static SelectList fullSelectLists;

fe/fe-core/src/main/java/org/apache/doris/catalog/SchemaTable.java

@@ -532,6 +532,14 @@ public class SchemaTable extends Table {
                                     .column("PARTITION_NUM", ScalarType.createType(PrimitiveType.INT))
                                     .column("PROPERTIES", ScalarType.createStringType())
                                     .build()))
+            .put("workload_group_privileges",
+                    new SchemaTable(SystemIdGenerator.getNextId(), "workload_group_privileges", TableType.SCHEMA,
+                            builder().column("GRANTEE", ScalarType.createVarchar(NAME_CHAR_LEN))
+                                    .column("WORKLOAD_GROUP_NAME", ScalarType.createVarchar(256))
+                                    .column("PRIVILEGE_TYPE", ScalarType.createVarchar(PRIVILEGE_TYPE_LEN))
+                                    .column("IS_GRANTABLE", ScalarType.createVarchar(IS_GRANTABLE_LEN))
+                                    .build())
+            )
             .build();
 
     protected SchemaTable(long id, String name, TableType type, List<Column> baseSchema) {

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

@@ -82,6 +82,7 @@
 import java.io.DataOutput;
 import java.io.IOException;
 import java.util.Collections;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
@@ -1407,6 +1408,47 @@ private void getUserAuthInfo(List<List<String>> userAuthInfos, UserIdentity user
         userAuthInfos.add(userAuthInfo);
     }
 
+    public void getUserRoleWorkloadGroupPrivs(List<List<String>> result, UserIdentity currentUserIdentity) {
+        readLock();
+        try {
+            boolean isCurrentUserAdmin = checkGlobalPriv(currentUserIdentity, PrivPredicate.ADMIN);
+            Map<String, List<User>> nameToUsers = userManager.getNameToUsers();
+            for (List<User> users : nameToUsers.values()) {
+                for (User user : users) {
+                    if (!user.isSetByDomainResolver()) {
+                        if (!isCurrentUserAdmin && !currentUserIdentity.equals(user.getUserIdentity())) {
+                            continue;
+                        }
+                        String isGrantable = checkGlobalPriv(user.getUserIdentity(), PrivPredicate.ADMIN) ? "YES"
+                                : "NO";
+
+                        // workload group
+                        for (PrivEntry entry : getUserWorkloadGroupPrivTable(user.getUserIdentity()).entries) {
+                            WorkloadGroupPrivEntry workloadGroupPrivEntry = (WorkloadGroupPrivEntry) entry;
+                            PrivBitSet savedPrivs = workloadGroupPrivEntry.getPrivSet().copy();
+
+                            List<String> row = Lists.newArrayList();
+                            row.add(user.getUserIdentity().toString());
+                            row.add(workloadGroupPrivEntry.getOrigWorkloadGroupName());
+                            row.add(savedPrivs.toString());
+                            row.add(isGrantable);
+                            result.add(row);
+                        }
+                    }
+                }
+            }
+
+            Set<String> currentUserRole = null;
+            if (!isCurrentUserAdmin) {
+                currentUserRole = userRoleManager.getRolesByUser(currentUserIdentity, false);
+                currentUserRole = currentUserRole == null ? new HashSet<>() : currentUserRole;
+            }
+            roleManager.getRoleWorkloadGroupPrivs(result, currentUserRole);
+        } finally {
+            readUnlock();
+        }
+    }
+
     private ResourcePrivTable getUserCloudClusterPrivTable(UserIdentity userIdentity) {
         ResourcePrivTable table = new ResourcePrivTable();
         Set<String> roles = userRoleManager.getRolesByUser(userIdentity);

fe/fe-core/src/main/java/org/apache/doris/mysql/privilege/RoleManager.java

@@ -37,6 +37,7 @@
 import org.apache.doris.qe.ConnectContext;
 import org.apache.doris.resource.workloadgroup.WorkloadGroupMgr;
 
+import com.aliyuncs.utils.StringUtils;
 import com.google.common.base.Joiner;
 import com.google.common.base.Strings;
 import com.google.common.collect.Lists;
@@ -211,6 +212,31 @@ public void getRoleInfo(List<List<String>> results) {
         }
     }
 
+    public void getRoleWorkloadGroupPrivs(List<List<String>> result, Set<String> limitedRole) {
+        for (Role role : roles.values()) {
+            if (ClusterNamespace.getNameFromFullName(role.getRoleName()).startsWith(DEFAULT_ROLE_PREFIX)) {
+                continue;
+            }
+
+            if (limitedRole != null && !limitedRole.contains(role.getRoleName())) {
+                continue;
+            }
+            String isGrantable = role.checkGlobalPriv(PrivPredicate.ADMIN) ? "YES" : "NO";
+
+            for (Map.Entry<WorkloadGroupPattern, PrivBitSet> entry : role.getWorkloadGroupPatternToPrivs().entrySet()) {
+                List<String> row = Lists.newArrayList();
+                row.add(role.getRoleName());
+                row.add(entry.getKey().getworkloadGroupName());
+                if (StringUtils.isEmpty(entry.getValue().toString())) {
+                    continue;
+                }
+                row.add(entry.getValue().toString());
+                row.add(isGrantable);
+                result.add(row);
+            }
+        }
+    }
+
     public Role createDefaultRole(UserIdentity userIdent) throws DdlException {
         String userDefaultRoleName = getUserDefaultRoleName(userIdent);
         if (roles.containsKey(userDefaultRoleName)) {

fe/fe-core/src/main/java/org/apache/doris/tablefunction/MetadataGenerator.java

@@ -113,6 +113,8 @@ public class MetadataGenerator {
 
     private static final ImmutableMap<String, Integer> TABLE_OPTIONS_COLUMN_TO_INDEX;
 
+    private static final ImmutableMap<String, Integer> WORKLOAD_GROUP_PRIVILEGES_COLUMN_TO_INDEX;
+
     static {
         ImmutableMap.Builder<String, Integer> activeQueriesbuilder = new ImmutableMap.Builder();
         List<Column> activeQueriesColList = SchemaTable.TABLE_MAP.get("active_queries").getFullSchema();
@@ -146,6 +148,13 @@ public class MetadataGenerator {
             optionBuilder.put(optionColList.get(i).getName().toLowerCase(), i);
         }
         TABLE_OPTIONS_COLUMN_TO_INDEX = optionBuilder.build();
+
+        ImmutableMap.Builder<String, Integer> wgPrivsBuilder = new ImmutableMap.Builder();
+        List<Column> wgPrivsColList = SchemaTable.TABLE_MAP.get("workload_group_privileges").getFullSchema();
+        for (int i = 0; i < wgPrivsColList.size(); i++) {
+            wgPrivsBuilder.put(wgPrivsColList.get(i).getName().toLowerCase(), i);
+        }
+        WORKLOAD_GROUP_PRIVILEGES_COLUMN_TO_INDEX = wgPrivsBuilder.build();
     }
 
     public static TFetchSchemaTableDataResult getMetadataTable(TFetchSchemaTableDataRequest request) throws TException {
@@ -229,6 +238,10 @@ public static TFetchSchemaTableDataResult getSchemaTableData(TFetchSchemaTableDa
                 result = tableOptionsMetadataResult(schemaTableParams);
                 columnIndex = TABLE_OPTIONS_COLUMN_TO_INDEX;
                 break;
+            case WORKLOAD_GROUP_PRIVILEGES:
+                result = workloadGroupPrivsMetadataResult(schemaTableParams);
+                columnIndex = WORKLOAD_GROUP_PRIVILEGES_COLUMN_TO_INDEX;
+                break;
             default:
                 return errorResult("invalid schema table name.");
         }
@@ -537,6 +550,30 @@ private static TFetchSchemaTableDataResult workloadSchedPolicyMetadataResult(TSc
         return result;
     }
 
+    private static TFetchSchemaTableDataResult workloadGroupPrivsMetadataResult(TSchemaTableRequestParams params) {
+        if (!params.isSetCurrentUserIdent()) {
+            return errorResult("current user ident is not set.");
+        }
+        UserIdentity currentUserIdentity = UserIdentity.fromThrift(params.getCurrentUserIdent());
+
+        List<List<String>> rows = new ArrayList<>();
+        Env.getCurrentEnv().getAuth().getUserRoleWorkloadGroupPrivs(rows, currentUserIdentity);
+        List<TRow> dataBatch = Lists.newArrayList();
+        for (List<String> privRow : rows) {
+            TRow trow = new TRow();
+            String workloadGroupName = privRow.get(1);
+            trow.addToColumnValue(new TCell().setStringVal(privRow.get(0))); // GRANTEE
+            trow.addToColumnValue(new TCell().setStringVal(workloadGroupName)); // WORKLOAD_GROUP_NAME
+            trow.addToColumnValue(new TCell().setStringVal(privRow.get(2))); // PRIVILEGE_TYPE
+            trow.addToColumnValue(new TCell().setStringVal(privRow.get(3))); // IS_GRANTABLE
+            dataBatch.add(trow);
+        }
+        TFetchSchemaTableDataResult result = new TFetchSchemaTableDataResult();
+        result.setDataBatch(dataBatch);
+        result.setStatus(new TStatus(TStatusCode.OK));
+        return result;
+    }
+
     private static TFetchSchemaTableDataResult queriesMetadataResult(TSchemaTableRequestParams tSchemaTableParams,
             TFetchSchemaTableDataRequest parentRequest) {
         TFetchSchemaTableDataResult result = new TFetchSchemaTableDataResult();