enhance: add CreateRole, DropRole, GrantRole, RevokeRole, DescribeRol…

…e and ListRoles Signed-off-by: Ruichen Bao <ruichen.bao@zju.edu.cn>
milvus-io · Jan 6, 2025 · 9c93747 · 9c93747
1 parent 2b2231a
commit 9c93747
Show file tree

Hide file tree

Showing 7 changed files with 235 additions and 0 deletions.
diff --git a/src/impl/MilvusClientImplV2.cpp b/src/impl/MilvusClientImplV2.cpp
@@ -979,6 +979,7 @@ MilvusClientImplV2::DescribeUser(const std::string& username, UserResult& result
         proto::milvus::SelectUserRequest rpc_request;
         auto* user_entity = rpc_request.mutable_user();
         user_entity->set_name(username);
+        rpc_request.set_include_role_info(true);
         return rpc_request;
     };
 
@@ -1045,6 +1046,105 @@ MilvusClientImplV2::DropUser(const std::string& username, int timeout) {
     return apiHandler<proto::milvus::DeleteCredentialRequest, proto::common::Status>(pre, &MilvusConnection::DeleteCredential, GrpcOpts{timeout});
 }
 
+Status
+MilvusClientImplV2::CreateRole(const std::string& role_name, int timeout) {
+    auto pre = [&role_name]() {
+        proto::milvus::CreateRoleRequest rpc_request;
+        proto::milvus::RoleEntity* role_entity = rpc_request.mutable_entity();
+        role_entity->set_name(role_name);
+        return rpc_request;
+    };
+
+    return apiHandler<proto::milvus::CreateRoleRequest, proto::common::Status>(pre, &MilvusConnection::CreateRole, GrpcOpts{timeout});
+}
+
+Status
+MilvusClientImplV2::DropRole(const std::string& role_name, int timeout) {
+    auto pre = [&role_name]() {
+        proto::milvus::DropRoleRequest rpc_request;
+        rpc_request.set_role_name(role_name);
+        return rpc_request;
+    };
+
+    return apiHandler<proto::milvus::DropRoleRequest, proto::common::Status>(pre, &MilvusConnection::DropRole, GrpcOpts{timeout});
+}
+
+Status
+MilvusClientImplV2::GrantRole(const std::string& username, const std::string& role_name, int timeout) {
+    auto pre = [&username, &role_name]() {
+        proto::milvus::OperateUserRoleRequest rpc_request;
+        rpc_request.set_username(username);
+        rpc_request.set_role_name(role_name);
+        rpc_request.set_type(proto::milvus::OperateUserRoleType::AddUserToRole);
+        return rpc_request;
+    };
+
+    return apiHandler<proto::milvus::OperateUserRoleRequest, proto::common::Status>(pre, &MilvusConnection::OperateUserRole, GrpcOpts{timeout});
+}
+
+Status
+MilvusClientImplV2::RevokeRole(const std::string& username, const std::string& role_name, int timeout) {
+    auto pre = [&username, &role_name]() {
+        proto::milvus::OperateUserRoleRequest rpc_request;
+        rpc_request.set_username(username);
+        rpc_request.set_role_name(role_name);
+        rpc_request.set_type(proto::milvus::OperateUserRoleType::RemoveUserFromRole);
+        return rpc_request;
+    };
+
+    return apiHandler<proto::milvus::OperateUserRoleRequest, proto::common::Status>(
+        pre, &MilvusConnection::OperateUserRole, GrpcOpts{timeout});
+}
+
+Status
+MilvusClientImplV2::DescribeRole(const std::string& role_name, RoleDesc& role_desc, int timeout) {
+    auto pre = [&role_name]() {
+        proto::milvus::SelectGrantRequest rpc_request;
+        auto* entity = rpc_request.mutable_entity();
+        auto* role = entity->mutable_role();
+        role->set_name(role_name);
+        return rpc_request;
+    };
+
+    auto post = [&role_desc, &role_name](const proto::milvus::SelectGrantResponse& response) {
+        std::vector<Privilege> privileges;
+        for (const auto& entity : response.entities()) {
+            if (entity.role().name() == role_name) {
+                Privilege p;
+                p.object_type = entity.object().name();
+                p.object_name = entity.object_name();
+                p.db_name = entity.db_name();
+                p.role_name = entity.role().name();
+                p.privilege = entity.grantor().privilege().name();
+                p.grantor_name = entity.grantor().user().name();
+                privileges.push_back(p);
+            }
+        }
+        role_desc = RoleDesc(role_name, privileges);
+    };
+
+    return apiHandler<proto::milvus::SelectGrantRequest, proto::milvus::SelectGrantResponse>(
+        pre, &MilvusConnection::SelectGrant, post, GrpcOpts{timeout});
+}
+
+Status
+MilvusClientImplV2::ListRoles(std::vector<std::string>& roles, int timeout) {
+    auto pre = []() {
+        proto::milvus::SelectRoleRequest rpc_request;
+        return rpc_request;
+    };
+
+    auto post = [&roles](const proto::milvus::SelectRoleResponse& response) {
+        roles.clear();
+        for (const auto& result : response.results()) {
+            roles.emplace_back(result.role().name());
+        }
+    };
+
+    return apiHandler<proto::milvus::SelectRoleRequest, proto::milvus::SelectRoleResponse>(
+        pre, &MilvusConnection::SelectRole, post, GrpcOpts{timeout});
+}
+
 Status
 MilvusClientImplV2::CalcDistance(const CalcDistanceArguments& arguments, DistanceArray& results) {
     auto validate = [&arguments]() { return arguments.Validate(); };

diff --git a/src/impl/MilvusClientImplV2.h b/src/impl/MilvusClientImplV2.h
@@ -164,6 +164,24 @@ class MilvusClientImplV2 : public MilvusClientV2 {
     Status
     DropUser(const std::string& username, int timeout) final;
 
+    Status
+    CreateRole(const std::string& role_name, int timeout) final;
+
+    Status
+    DropRole(const std::string& role_name, int timeout) final;
+
+    Status
+    GrantRole(const std::string& username, const std::string& role_name, int timeout) final;
+
+    Status
+    RevokeRole(const std::string& username, const std::string& role_name, int timeout) final;
+
+    Status
+    DescribeRole(const std::string& role_name, RoleDesc& role_desc, int timeout) final;
+
+    Status 
+    ListRoles(std::vector<std::string>& roles, int timeout) final;
+
     Status
     CalcDistance(const CalcDistanceArguments& arguments, DistanceArray& results) final;
 

diff --git a/src/impl/MilvusConnection.cpp b/src/impl/MilvusConnection.cpp
@@ -391,6 +391,36 @@ MilvusConnection::SelectUser(const proto::milvus::SelectUserRequest& request,
     return grpcCall("SelectUser", &Stub::SelectUser, request, response, options);
 }
 
+Status
+MilvusConnection::CreateRole(const proto::milvus::CreateRoleRequest& request,
+                             proto::common::Status& response, const GrpcContextOptions& options) {
+    return grpcCall("CreateRole", &Stub::CreateRole, request, response, options);
+}
+
+Status
+MilvusConnection::DropRole(const proto::milvus::DropRoleRequest& request,
+                           proto::common::Status& response, const GrpcContextOptions& options) {
+    return grpcCall("DropRole", &Stub::DropRole, request, response, options);
+}
+
+Status
+MilvusConnection::OperateUserRole(const proto::milvus::OperateUserRoleRequest& request,
+                                  proto::common::Status& response, const GrpcContextOptions& options) {
+    return grpcCall("OperateUserRole", &Stub::OperateUserRole, request, response, options);
+}
+
+Status
+MilvusConnection::SelectGrant(const proto::milvus::SelectGrantRequest& request,
+                            proto::milvus::SelectGrantResponse& response, const GrpcContextOptions& options) {
+    return grpcCall("SelectGrant", &Stub::SelectGrant, request, response, options);
+}
+
+Status
+MilvusConnection::SelectRole(const proto::milvus::SelectRoleRequest& request,
+                           proto::milvus::SelectRoleResponse& response, const GrpcContextOptions& options) {
+    return grpcCall("SelectRole", &Stub::SelectRole, request, response, options);
+}
+
 void
 MilvusConnection::SetHeader(const std::string& key, const std::string& value) {
     headers_[key] = value;

diff --git a/src/impl/MilvusConnection.h b/src/impl/MilvusConnection.h
@@ -241,6 +241,26 @@ class MilvusConnection {
     SelectUser(const proto::milvus::SelectUserRequest& request, proto::milvus::SelectUserResponse& response,
                const GrpcContextOptions& options);
 
+    Status
+    CreateRole(const proto::milvus::CreateRoleRequest& request, proto::common::Status& response,
+               const GrpcContextOptions& options);
+
+    Status
+    DropRole(const proto::milvus::DropRoleRequest& request, proto::common::Status& response,
+             const GrpcContextOptions& options);
+
+    Status
+    OperateUserRole(const proto::milvus::OperateUserRoleRequest& request, proto::common::Status& response,
+                    const GrpcContextOptions& options);
+
+    Status
+    SelectGrant(const proto::milvus::SelectGrantRequest& request, proto::milvus::SelectGrantResponse& response,
+                const GrpcContextOptions& options);
+
+    Status
+    SelectRole(const proto::milvus::SelectRoleRequest& request, proto::milvus::SelectRoleResponse& response,
+               const GrpcContextOptions& options);
+
     void
     SetHeader(const std::string& key, const std::string& value);
 

diff --git a/src/impl/types/RoleDesc.cpp b/src/impl/types/RoleDesc.cpp
@@ -0,0 +1,18 @@
+#include "milvus/types/RoleDesc.h"
+
+namespace milvus {
+
+RoleDesc::RoleDesc() {}
+
+RoleDesc::RoleDesc(const std::string& role, const std::vector<Privilege>& privileges)
+    : role_(role), privileges_(privileges) {}
+
+const std::string& RoleDesc::GetRole() const {
+    return role_;
+}
+
+const std::vector<Privilege>& RoleDesc::GetPrivileges() const {
+    return privileges_;
+}
+
+}  // namespace milvus
diff --git a/src/include/milvus/MilvusClientV2.h b/src/include/milvus/MilvusClientV2.h
@@ -39,6 +39,7 @@
 #include "types/ProgressMonitor.h"
 #include "types/QueryArguments.h"
 #include "types/QueryResults.h"
+#include "types/RoleDesc.h"
 #include "types/SearchArguments.h"
 #include "types/SearchResults.h"
 #include "types/SegmentInfo.h"
@@ -440,6 +441,24 @@ class MilvusClientV2 {
     virtual Status
     DropUser(const std::string& username, int timeout = 0) = 0;
 
+    virtual Status
+    CreateRole(const std::string& role_name, int timeout = 0) = 0;
+
+    virtual Status
+    DropRole(const std::string& role_name, int timeout = 0) = 0;
+
+    virtual Status
+    GrantRole(const std::string& username, const std::string& role_name, int timeout = 0) = 0;
+
+    virtual Status
+    RevokeRole(const std::string& username, const std::string& role_name, int timeout = 0) = 0;
+
+    virtual Status 
+    DescribeRole(const std::string& role_name, RoleDesc& role_desc, int timeout = 0) = 0;
+
+    virtual Status 
+    ListRoles(std::vector<std::string>& roles, int timeout = 0) = 0;
+
     /**
      * Calculate distance between two vector arrays.
      *

diff --git a/src/include/milvus/types/RoleDesc.h b/src/include/milvus/types/RoleDesc.h
@@ -0,0 +1,30 @@
+#pragma once
+
+#include <string>
+#include <vector>
+
+namespace milvus {
+
+struct Privilege {
+    std::string object_type;
+    std::string object_name;
+    std::string db_name;
+    std::string role_name;
+    std::string privilege;
+    std::string grantor_name;
+};
+
+class RoleDesc {
+public:
+    RoleDesc();
+    RoleDesc(const std::string& role, const std::vector<Privilege>& privileges);
+
+    const std::string& GetRole() const;
+    const std::vector<Privilege>& GetPrivileges() const;
+
+private:
+    std::string role_;
+    std::vector<Privilege> privileges_;
+};
+
+}  // namespace milvus