enhance: add GetLoadState, RefreshLoad, GrantPrivilege and RevokePriv…

…ilege Signed-off-by: Ruichen Bao <ruichen.bao@zju.edu.cn>
milvus-io · Jan 7, 2025 · f581ca0 · f581ca0
1 parent 9c93747
commit f581ca0
Show file tree

Hide file tree

Showing 7 changed files with 225 additions and 2 deletions.
diff --git a/src/impl/MilvusClientImplV2.cpp b/src/impl/MilvusClientImplV2.cpp
@@ -473,6 +473,40 @@ MilvusClientImplV2::ShowPartitions(const std::string& collection_name, const std
         pre, &MilvusConnection::ShowPartitions, post);
 }
 
+Status
+MilvusClientImplV2::GetLoadState(const std::string& collection_name, LoadState& state,
+                                 const std::string& partition_name, int timeout) {
+    auto pre = [&collection_name, &partition_name]() {
+        proto::milvus::GetLoadStateRequest rpc_request;
+        rpc_request.set_collection_name(collection_name);
+        if (!partition_name.empty()) {
+            rpc_request.add_partition_names(partition_name);
+        }
+        return rpc_request;
+    };
+
+    auto post = [&state](const proto::milvus::GetLoadStateResponse& response) {
+        state.SetCode(static_cast<LoadStateCode>(response.state()));
+    };
+
+    return apiHandler<proto::milvus::GetLoadStateRequest, proto::milvus::GetLoadStateResponse>(
+        pre, &MilvusConnection::GetLoadState, post, GrpcOpts{timeout});
+}
+
+Status
+MilvusClientImplV2::RefreshLoad(const std::string& collection_name, int timeout) {
+    auto pre = [&collection_name]() {
+        proto::milvus::LoadCollectionRequest rpc_request;
+        rpc_request.set_collection_name(collection_name);
+        rpc_request.set_refresh(true);
+        rpc_request.set_replica_number(1);
+        return rpc_request;
+    };
+
+    return apiHandler<proto::milvus::LoadCollectionRequest, proto::common::Status>(
+        pre, &MilvusConnection::LoadCollection, GrpcOpts{timeout});
+}
+
 Status
 MilvusClientImplV2::CreateAlias(const std::string& collection_name, const std::string& alias) {
     auto pre = [&collection_name, &alias]() {
@@ -1145,6 +1179,64 @@ MilvusClientImplV2::ListRoles(std::vector<std::string>& roles, int timeout) {
         pre, &MilvusConnection::SelectRole, post, GrpcOpts{timeout});
 }
 
+Status
+MilvusClientImplV2::GrantPrivilege(const std::string& role_name, const std::string& object_type,
+                                   const std::string& privilege, const std::string& object_name,
+                                   const std::string& db_name, int timeout) {
+    auto pre = [&role_name, &object_type, &privilege, &object_name, &db_name]() {
+        proto::milvus::OperatePrivilegeRequest rpc_request;
+
+        auto* role_entity = rpc_request.mutable_entity()->mutable_role();
+        role_entity->set_name(role_name);
+        auto* object_entity = rpc_request.mutable_entity()->mutable_object();
+        object_entity->set_name(object_type);
+        rpc_request.mutable_entity()->set_object_name(object_name);
+        auto* grantor = rpc_request.mutable_entity()->mutable_grantor();
+        auto* privilege_entity = grantor->mutable_privilege();
+        privilege_entity->set_name(privilege);
+
+        if (!db_name.empty()) {
+            rpc_request.mutable_entity()->set_db_name(db_name);
+        }
+
+        rpc_request.set_type(proto::milvus::OperatePrivilegeType::Grant);
+
+        return rpc_request;
+    };
+
+    return apiHandler<proto::milvus::OperatePrivilegeRequest, proto::common::Status>(
+        pre, &MilvusConnection::OperatePrivilege, GrpcOpts{timeout});
+}
+
+Status
+MilvusClientImplV2::RevokePrivilege(const std::string& role_name, const std::string& object_type,
+                                    const std::string& privilege, const std::string& object_name,
+                                    const std::string& db_name, int timeout) {
+    auto pre = [&role_name, &object_type, &privilege, &object_name, &db_name]() {
+        proto::milvus::OperatePrivilegeRequest rpc_request;
+
+        auto* role_entity = rpc_request.mutable_entity()->mutable_role();
+        role_entity->set_name(role_name);
+        auto* object_entity = rpc_request.mutable_entity()->mutable_object();
+        object_entity->set_name(object_type);
+        rpc_request.mutable_entity()->set_object_name(object_name);
+        auto* grantor = rpc_request.mutable_entity()->mutable_grantor();
+        auto* privilege_entity = grantor->mutable_privilege();
+        privilege_entity->set_name(privilege);
+
+        if (!db_name.empty()) {
+            rpc_request.mutable_entity()->set_db_name(db_name);
+        }
+
+        rpc_request.set_type(proto::milvus::OperatePrivilegeType::Revoke);
+
+        return rpc_request;
+    };
+
+    return apiHandler<proto::milvus::OperatePrivilegeRequest, proto::common::Status>(
+        pre, &MilvusConnection::OperatePrivilege, 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
@@ -99,6 +99,13 @@ class MilvusClientImplV2 : public MilvusClientV2 {
     ShowPartitions(const std::string& collection_name, const std::vector<std::string>& partition_names,
                    PartitionsInfo& partitions_info) final;
 
+    Status
+    GetLoadState(const std::string& collection_name, LoadState& state, 
+                 const std::string& partition_name, int timeout) final;
+
+    Status
+    RefreshLoad(const std::string& collection_name, int timeout) final;
+
     Status
     CreateAlias(const std::string& collection_name, const std::string& alias) final;
 
@@ -182,6 +189,16 @@ class MilvusClientImplV2 : public MilvusClientV2 {
     Status 
     ListRoles(std::vector<std::string>& roles, int timeout) final;
 
+    Status
+    GrantPrivilege(const std::string& role_name, const std::string& object_type,
+                   const std::string& privilege, const std::string& object_name,
+                   const std::string& db_name, int timeout) final;
+
+    Status
+    RevokePrivilege(const std::string& role_name, const std::string& object_type,
+                    const std::string& privilege, const std::string& object_name,
+                    const std::string& db_name, int timeout) final;
+
     Status
     CalcDistance(const CalcDistanceArguments& arguments, DistanceArray& results) final;
 

diff --git a/src/impl/MilvusConnection.cpp b/src/impl/MilvusConnection.cpp
@@ -217,6 +217,13 @@ MilvusConnection::GetPartitionStatistics(const proto::milvus::GetPartitionStatis
     return grpcCall("GetPartitionStatistics", &Stub::GetPartitionStatistics, request, response, options);
 }
 
+Status
+MilvusConnection::GetLoadState(const proto::milvus::GetLoadStateRequest& request,
+                               proto::milvus::GetLoadStateResponse& response,
+                               const GrpcContextOptions& options) {
+    return grpcCall("GetLoadState", &Stub::GetLoadState, request, response, options);
+}
+
 Status
 MilvusConnection::CreateAlias(const proto::milvus::CreateAliasRequest& request, proto::common::Status& response,
                               const GrpcContextOptions& options) {
@@ -411,16 +418,22 @@ MilvusConnection::OperateUserRole(const proto::milvus::OperateUserRoleRequest& r
 
 Status
 MilvusConnection::SelectGrant(const proto::milvus::SelectGrantRequest& request,
-                            proto::milvus::SelectGrantResponse& response, const GrpcContextOptions& options) {
+                              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) {
+                             proto::milvus::SelectRoleResponse& response, const GrpcContextOptions& options) {
     return grpcCall("SelectRole", &Stub::SelectRole, request, response, options);
 }
 
+Status
+MilvusConnection::OperatePrivilege(const proto::milvus::OperatePrivilegeRequest& request,
+                                   proto::common::Status& response, const GrpcContextOptions& options) {
+    return grpcCall("OperatePrivilege", &Stub::OperatePrivilege, 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
@@ -128,6 +128,10 @@ class MilvusConnection {
     GetPartitionStatistics(const proto::milvus::GetPartitionStatisticsRequest& request,
                            proto::milvus::GetPartitionStatisticsResponse& response, const GrpcContextOptions& options);
 
+    Status
+    GetLoadState(const proto::milvus::GetLoadStateRequest& request,
+                 proto::milvus::GetLoadStateResponse& response, const GrpcContextOptions& options);
+
     Status
     CreateAlias(const proto::milvus::CreateAliasRequest& request, proto::common::Status& response,
                 const GrpcContextOptions& options);
@@ -261,6 +265,10 @@ class MilvusConnection {
     SelectRole(const proto::milvus::SelectRoleRequest& request, proto::milvus::SelectRoleResponse& response,
                const GrpcContextOptions& options);
 
+    Status
+    OperatePrivilege(const proto::milvus::OperatePrivilegeRequest& request, proto::common::Status& response,
+                     const GrpcContextOptions& options);
+
     void
     SetHeader(const std::string& key, const std::string& value);
 

diff --git a/src/impl/types/LoadState.cpp b/src/impl/types/LoadState.cpp
@@ -0,0 +1,45 @@
+#include "milvus/types/LoadState.h"
+
+namespace milvus {
+
+LoadState::LoadState() : code_(LoadStateCode::NotExist) {
+    UpdateStateDesc();
+}
+
+LoadState::LoadState(LoadStateCode code) : code_(code) {
+    UpdateStateDesc();
+}
+
+LoadStateCode LoadState::GetCode() const {
+    return code_;
+}
+
+const std::string& LoadState::GetDesc() const {
+    return state_desc_;
+}
+
+void LoadState::SetCode(LoadStateCode code) {
+    code_ = code;
+    UpdateStateDesc();
+}
+
+void LoadState::UpdateStateDesc() {
+    switch (code_) {
+        case LoadStateCode::NotExist:
+            state_desc_ = "NotExist";
+            break;
+        case LoadStateCode::NotLoad:
+            state_desc_ = "NotLoad";
+            break;
+        case LoadStateCode::Loading:
+            state_desc_ = "Loading";
+            break;
+        case LoadStateCode::Loaded:
+            state_desc_ = "Loaded";
+            break;
+        default:
+            state_desc_ = "Unknown";
+    }
+}
+
+}  // namespace milvus
diff --git a/src/include/milvus/MilvusClientV2.h b/src/include/milvus/MilvusClientV2.h
@@ -34,6 +34,7 @@
 #include "types/HybridTimestamp.h"
 #include "types/IndexDesc.h"
 #include "types/IndexState.h"
+#include "types/LoadState.h"
 #include "types/PartitionInfo.h"
 #include "types/PartitionStat.h"
 #include "types/ProgressMonitor.h"
@@ -267,6 +268,14 @@ class MilvusClientV2 {
                            PartitionStat& partition_stat,
                            const ProgressMonitor& progress_monitor = ProgressMonitor::Forever()) = 0;
 
+
+    virtual Status
+    GetLoadState(const std::string& collection_name, LoadState& state, 
+                 const std::string& partition_name = "", int timeout = 0) = 0;
+
+    virtual Status
+    RefreshLoad(const std::string& collection_name, int timeout = 0) = 0;
+
     /**
      * If the partition_names is empty, list all partitions brief information's. \n
      * If the partition_names is specified, return the specified partition's loading process state.
@@ -459,6 +468,16 @@ class MilvusClientV2 {
     virtual Status 
     ListRoles(std::vector<std::string>& roles, int timeout = 0) = 0;
 
+    virtual Status
+    GrantPrivilege(const std::string& role_name, const std::string& object_type,
+                   const std::string& privilege, const std::string& object_name,
+                   const std::string& db_name = "", int timeout = 0) = 0;
+
+    virtual Status
+    RevokePrivilege(const std::string& role_name, const std::string& object_type,
+                    const std::string& privilege, const std::string& object_name,
+                    const std::string& db_name = "", int timeout = 0) = 0;
+
     /**
      * Calculate distance between two vector arrays.
      *

diff --git a/src/include/milvus/types/LoadState.h b/src/include/milvus/types/LoadState.h
@@ -0,0 +1,29 @@
+#pragma once
+
+#include <string>
+
+namespace milvus {
+
+enum class LoadStateCode {
+    NotExist = 0,
+    NotLoad = 1,
+    Loading = 2,
+    Loaded = 3
+};
+
+class LoadState {
+public:
+    LoadState();
+    explicit LoadState(LoadStateCode code);
+
+    LoadStateCode GetCode() const;
+    const std::string& GetDesc() const;
+    void SetCode(LoadStateCode code);
+
+private:
+    LoadStateCode code_;
+    std::string state_desc_;
+    void UpdateStateDesc();
+};
+
+}  // namespace milvus