From a225cbfb66f06e48fce6f350b2948fff6b502e4c Mon Sep 17 00:00:00 2001
From: sthuang <167743503+shaoting-huang@users.noreply.github.com>
Date: Wed, 4 Dec 2024 11:50:32 +0800
Subject: [PATCH] fix: remove grant/revoke v2 optional collection name (#2403)

Since built-in privilege group level adjustment is not supported, remove
the optional collection name from the grant/revoke v2 API. The empty
db_name uses default db which follows the same rule as grant/revoke v1.
issue: https://github.com/milvus-io/milvus/issues/37031

Signed-off-by: shaoting-huang <shaoting.huang@zilliz.com>
---
 pymilvus/client/grpc_handler.py         |  4 +--
 pymilvus/client/prepare.py              |  3 +--
 pymilvus/milvus_client/milvus_client.py |  8 +++---
 pymilvus/orm/role.py                    | 34 +++++++++++++------------
 4 files changed, 25 insertions(+), 24 deletions(-)

diff --git a/pymilvus/client/grpc_handler.py b/pymilvus/client/grpc_handler.py
index be55eac90..9bb508752 100644
--- a/pymilvus/client/grpc_handler.py
+++ b/pymilvus/client/grpc_handler.py
@@ -1877,8 +1877,8 @@ def grant_privilege_v2(
         self,
         role_name: str,
         privilege: str,
+        collection_name: str,
         db_name: Optional[str] = None,
-        collection_name: Optional[str] = None,
         timeout: Optional[float] = None,
         **kwargs,
     ):
@@ -1897,8 +1897,8 @@ def revoke_privilege_v2(
         self,
         role_name: str,
         privilege: str,
+        collection_name: str,
         db_name: Optional[str] = None,
-        collection_name: Optional[str] = None,
         timeout: Optional[float] = None,
         **kwargs,
     ):
diff --git a/pymilvus/client/prepare.py b/pymilvus/client/prepare.py
index 44ce44266..e0e591af8 100644
--- a/pymilvus/client/prepare.py
+++ b/pymilvus/client/prepare.py
@@ -1484,12 +1484,11 @@ def operate_privilege_v2_request(
         check_pass_param(
             role_name=role_name,
             privilege=privilege,
+            collection_name=collection_name,
             operate_privilege_type=operate_privilege_type,
         )
         if db_name:
             check_pass_param(db_name=db_name)
-        if collection_name:
-            check_pass_param(collection_name=collection_name)
         return milvus_types.OperatePrivilegeV2Request(
             role=milvus_types.RoleEntity(name=role_name),
             grantor=milvus_types.GrantorEntity(
diff --git a/pymilvus/milvus_client/milvus_client.py b/pymilvus/milvus_client/milvus_client.py
index b97147e9b..4aaccebf4 100644
--- a/pymilvus/milvus_client/milvus_client.py
+++ b/pymilvus/milvus_client/milvus_client.py
@@ -1007,8 +1007,8 @@ def grant_privilege_v2(
         self,
         role_name: str,
         privilege: str,
+        collection_name: str,
         db_name: Optional[str] = None,
-        collection_name: Optional[str] = None,
         timeout: Optional[float] = None,
         **kwargs,
     ):
@@ -1016,8 +1016,8 @@ def grant_privilege_v2(
         conn.grant_privilege_v2(
             role_name,
             privilege,
+            collection_name,
             db_name=db_name,
-            collection_name=collection_name,
             timeout=timeout,
             **kwargs,
         )
@@ -1026,8 +1026,8 @@ def revoke_privilege_v2(
         self,
         role_name: str,
         privilege: str,
+        collection_name: str,
         db_name: Optional[str] = None,
-        collection_name: Optional[str] = None,
         timeout: Optional[float] = None,
         **kwargs,
     ):
@@ -1035,8 +1035,8 @@ def revoke_privilege_v2(
         conn.revoke_privilege_v2(
             role_name,
             privilege,
+            collection_name,
             db_name=db_name,
-            collection_name=collection_name,
             timeout=timeout,
             **kwargs,
         )
diff --git a/pymilvus/orm/role.py b/pymilvus/orm/role.py
index 313231ad4..21e1b935f 100644
--- a/pymilvus/orm/role.py
+++ b/pymilvus/orm/role.py
@@ -178,48 +178,50 @@ def revoke(self, object: str, object_name: str, privilege: str, db_name: str = "
             self._name, object, object_name, privilege, db_name
         )
 
-    def grant_v2(
-        self, privilege: str, db_name: Optional[str] = None, collection_name: Optional[str] = None
-    ):
+    def grant_v2(self, privilege: str, collection_name: str, db_name: Optional[str] = None):
         """Grant a privilege for the role
             :param privilege: privilege name.
             :type  privilege: str
-            :param db_name: db name. Optional
-            :type  db_name: str
-            :param collection_name: collection name. Optional
+            :param collection_name: collection name.
             :type  collection_name: str
+            :param db_name: db name. Optional. If None, use the default db.
+            :type  db_name: str
 
         :example:
             >>> from pymilvus import connections
             >>> from pymilvus.orm.role import Role
             >>> connections.connect()
             >>> role = Role(role_name)
-            >>> role.grant_v2("Insert", db_name, collection_name)
+            >>> role.grant_v2("Insert", collection_name, db_name=db_name)
         """
         return self._get_connection().grant_privilege_v2(
-            self._name, privilege, db_name=db_name, collection_name=collection_name
+            self._name,
+            privilege,
+            collection_name,
+            db_name=db_name,
         )
 
-    def revoke_v2(
-        self, privilege: str, db_name: Optional[str] = None, collection_name: Optional[str] = None
-    ):
+    def revoke_v2(self, privilege: str, collection_name: str, db_name: Optional[str] = None):
         """Revoke a privilege for the role
             :param privilege: privilege name.
             :type  privilege: str
-            :param db_name: db name. Optional
-            :type  db_name: str
-            :param collection_name: collection name. Optional
+            :param collection_name: collection name.
             :type  collection_name: str
+            :param db_name: db name. Optional. If None, use the default db.
+            :type  db_name: str
 
         :example:
             >>> from pymilvus import connections
             >>> from pymilvus.orm.role import Role
             >>> connections.connect()
             >>> role = Role(role_name)
-            >>> role.revoke_v2("Insert", db_name, collection_name)
+            >>> role.revoke_v2("Insert", collection_name, db_name=db_name)
         """
         return self._get_connection().revoke_privilege_v2(
-            self._name, privilege, db_name=db_name, collection_name=collection_name
+            self._name,
+            privilege,
+            collection_name,
+            db_name=db_name,
         )
 
     def list_grant(self, object: str, object_name: str, db_name: str = ""):