enhance: grant/revoke v2 optional db and collection params (milvus-io…

…#2386) issue: milvus-io/milvus#37031 Signed-off-by: shaoting-huang <[email protected]>
weiliu1031 · Nov 28, 2024 · 053c568 · 053c568
1 parent 3449513
commit 053c568
Show file tree

Hide file tree

Showing 4 changed files with 38 additions and 20 deletions.
diff --git a/pymilvus/client/grpc_handler.py b/pymilvus/client/grpc_handler.py
@@ -1876,8 +1876,8 @@ def grant_privilege_v2(
         self,
         role_name: str,
         privilege: str,
-        db_name: str,
-        collection_name: str,
+        db_name: Optional[str] = None,
+        collection_name: Optional[str] = None,
         timeout: Optional[float] = None,
         **kwargs,
     ):
@@ -1896,8 +1896,8 @@ def revoke_privilege_v2(
         self,
         role_name: str,
         privilege: str,
-        db_name: 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
@@ -1484,9 +1484,11 @@ def operate_privilege_v2_request(
             role_name=role_name,
             privilege=privilege,
             operate_privilege_type=operate_privilege_type,
-            db_name=db_name,
-            collection_name=collection_name,
         )
+        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
@@ -1008,28 +1008,38 @@ def grant_privilege_v2(
         self,
         role_name: str,
         privilege: str,
-        db_name: str,
-        collection_name: str,
+        db_name: Optional[str] = None,
+        collection_name: Optional[str] = None,
         timeout: Optional[float] = None,
         **kwargs,
     ):
         conn = self._get_connection()
         conn.grant_privilege_v2(
-            role_name, privilege, db_name, collection_name, timeout=timeout, **kwargs
+            role_name,
+            privilege,
+            db_name=db_name,
+            collection_name=collection_name,
+            timeout=timeout,
+            **kwargs,
         )
 
     def revoke_privilege_v2(
         self,
         role_name: str,
         privilege: str,
-        db_name: str,
-        collection_name: str,
+        db_name: Optional[str] = None,
+        collection_name: Optional[str] = None,
         timeout: Optional[float] = None,
         **kwargs,
     ):
         conn = self._get_connection()
         conn.revoke_privilege_v2(
-            role_name, privilege, db_name, collection_name, timeout=timeout, **kwargs
+            role_name,
+            privilege,
+            db_name=db_name,
+            collection_name=collection_name,
+            timeout=timeout,
+            **kwargs,
         )
 
     def create_alias(

diff --git a/pymilvus/orm/role.py b/pymilvus/orm/role.py
@@ -10,6 +10,8 @@
 # or implied. See the License for the specific language governing permissions and limitations under
 # the License.
 
+from typing import Optional
+
 from .connections import connections
 
 INCLUDE_USER_INFO, NOT_INCLUDE_USER_INFO = True, False
@@ -176,13 +178,15 @@ 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: str, collection_name: str):
+    def grant_v2(
+        self, privilege: str, db_name: Optional[str] = None, collection_name: Optional[str] = None
+    ):
         """Grant a privilege for the role
             :param privilege: privilege name.
             :type  privilege: str
-            :param db_name: db name.
+            :param db_name: db name. Optional
             :type  db_name: str
-            :param collection_name: collection name.
+            :param collection_name: collection name. Optional
             :type  collection_name: str
 
         :example:
@@ -193,16 +197,18 @@ def grant_v2(self, privilege: str, db_name: str, collection_name: str):
             >>> role.grant_v2("Insert", db_name, collection_name)
         """
         return self._get_connection().grant_privilege_v2(
-            self._name, privilege, db_name, collection_name
+            self._name, privilege, db_name=db_name, collection_name=collection_name
         )
 
-    def revoke_v2(self, privilege: str, db_name: str, collection_name: str):
+    def revoke_v2(
+        self, privilege: str, db_name: Optional[str] = None, collection_name: Optional[str] = None
+    ):
         """Revoke a privilege for the role
             :param privilege: privilege name.
             :type  privilege: str
-            :param db_name: db name.
+            :param db_name: db name. Optional
             :type  db_name: str
-            :param collection_name: collection name.
+            :param collection_name: collection name. Optional
             :type  collection_name: str
 
         :example:
@@ -213,7 +219,7 @@ def revoke_v2(self, privilege: str, db_name: str, collection_name: str):
             >>> role.revoke_v2("Insert", db_name, collection_name)
         """
         return self._get_connection().revoke_privilege_v2(
-            self._name, privilege, db_name, collection_name
+            self._name, privilege, db_name=db_name, collection_name=collection_name
         )
 
     def list_grant(self, object: str, object_name: str, db_name: str = ""):