Merge pull request #42 from d4rkstar/main

fix(milvus): milvus 2.5.2 permissions

Author: sciabarracom

nuvolaris/milvus_admin_client.py

@@ -18,8 +18,9 @@
 import logging
 
 import nuvolaris.config as cfg
+# from pymilvus import MilvusClient, connections,  db
 from nuvolaris.milvus_simple_client import MilvusSimpleClient as MilvusClient
-#from pymilvus import MilvusClient, connections,  db
+
 
 class MilvusAdminClient:
     """
@@ -33,6 +34,15 @@ def __init__(self, db_name="default"):
         self.milvus_url   = f"http://{self.milvus_api_host}:{self.milvus_api_port}"
         self.milvus_admin_token = f"root:{self.admin_password}"
 
+        #self.global_privileges_v1 = ['CreateCollection', 'DropCollection', 'DescribeCollection', 'ShowCollections',
+        #                 'RenameCollection']
+        self.global_privileges_v1 = []
+
+        # references:
+        # https://milvus.io/docs/privilege_group.md
+        # https://milvus.io/docs/grant_privileges.md#Grant-a-privilege-or-a-privilege-group-to-a-role
+        self.global_privileges_v2 = ['CollectionAdmin','DatabaseAdmin']
+
     def setup_user(self, username, password,database):
         """
         Creates a user into MILVUS, creates a corresponding database
@@ -53,12 +63,11 @@ def setup_user(self, username, password,database):
             # rest of action are performed specifying the database
             client = MilvusClient(uri=self.milvus_url,token=self.milvus_admin_token, db_name=database)
             client.create_role(role_name=role,db_name=database)
-            client.grant_privilege(role_name=role, object_type='Global',  object_name='*', privilege='CreateCollection', db_name=database)
-            client.grant_privilege(role_name=role, object_type='Global',  object_name='*', privilege='DropCollection', db_name=database)
-            client.grant_privilege(role_name=role, object_type='Global',  object_name='*', privilege='DescribeCollection', db_name=database)
-            client.grant_privilege(role_name=role, object_type='Global',  object_name='*', privilege='ShowCollections', db_name=database)
-            client.grant_privilege(role_name=role, object_type='Global',  object_name='*', privilege='RenameCollection', db_name=database)
-            client.grant_privilege(role_name=role, object_type='Collection',  object_name='*', privilege='*', db_name=database)
+            for priv in self.global_privileges_v1:
+                client.grant_privilege(role_name=role, object_type='Global',  object_name='*', privilege=priv, db_name=database)
+            for priv in self.global_privileges_v2:
+                client.grant_privilege_v2(role_name=role, object_type='Global',  object_name='*', collection_name='*', privilege=priv, db_name=database)
+
             client.grant_role(user_name=username,role_name=role,db_name=database)
             client.close()
             return True
@@ -82,7 +91,16 @@ def remove_user(self, username, database):
                 client.drop_collection(collection_name=collection)
             client.close()
 
-            client = MilvusClient(uri=self.milvus_url,token=self.milvus_admin_token)                    
+            client = MilvusClient(uri=self.milvus_url,token=self.milvus_admin_token)
+
+            for privilege in self.global_privileges_v1:
+                client.revoke_privilege(role_name=role, object_type='Global', object_name='*', privilege=privilege,
+                                       db_name=database)
+            for privilege in self.global_privileges_v2:
+                client.revoke_privilege_v2(role_name=role, object_type='Global', object_name='*', collection_name='*',
+                                          privilege=privilege, db_name=database)
+
+
             client.drop_role(role_name=role,db_name=database)
             client.drop_user(user_name=username)                
             client.drop_database(db_name=database)

nuvolaris/milvus_simple_client.py

@@ -17,6 +17,7 @@
 #
 from types import NoneType
 from typing import Optional
+from requests.exceptions import HTTPError
 
 import requests
 
@@ -29,6 +30,10 @@ def __init__(self, code: int, message: str):
     def __str__(self):
         return f"{self.message} ({self.code})"
 
+class MilvusUnauthorizedException(MilvusSimpleException):
+    def __init__(self):
+        super().__init__(0, "Unauthorized")
+
 
 class MilvusSimpleClient:
 
@@ -46,12 +51,17 @@ def _request(self, endpoint: str, json=None, method: str = "POST", api_level="v2
         if json is None:
             json = {}
         response = requests.request(method, url, headers=headers, json=json)
-        response.raise_for_status()
-        res = response.json()
-        if type(res.get('data')) is not NoneType:
-            return res['data']
-        else:
-            raise MilvusSimpleException(code=res.get('code'), message=res.get('message'))
+        try:
+            response.raise_for_status()
+            res = response.json()
+            if type(res.get('data')) is not NoneType:
+                return res['data']
+            else:
+                raise MilvusSimpleException(code=res.get('code'), message=res.get('message'))
+        except HTTPError as e:
+            if e.response.status_code == 403:
+                raise MilvusUnauthorizedException()
+
 
     def close(self):
         pass
@@ -162,6 +172,31 @@ def revoke_privilege(self, role_name: str, object_type: str, object_name: str, p
             payload["dbName"] = db_name
         return self._request(f"roles/revoke_privilege", json=payload)
 
+
+    def grant_privilege_v2(self, role_name: str, object_type: str, object_name: str, collection_name: str,  privilege: str, db_name: Optional[str] = None):
+        payload = {
+            "roleName": role_name,
+            "objectType": object_type,
+            "objectName": object_name,
+            "privilege": privilege,
+            "collectionName": collection_name,
+        }
+        if db_name is not None:
+            payload["dbName"] = db_name
+        return self._request(f"roles/grant_privilege_v2", json=payload)
+
+    def revoke_privilege_v2(self, role_name: str, object_type: str, object_name: str, collection_name: str,  privilege: str, db_name: Optional[str] = None):
+        payload = {
+            "roleName": role_name,
+            "objectType": object_type,
+            "objectName": object_name,
+            "privilege": privilege,
+            "collectionName": collection_name,
+        }
+        if db_name is not None:
+            payload["dbName"] = db_name
+        return self._request(f"roles/revoke_privilege_v2", json=payload)
+
     # Collection operations
 
     def create_collection(self,

tests/kind/whisk.yaml

@@ -173,6 +173,6 @@ spec:
       ledgers: 25
     replicas: 1
     password:
-      root: x£VqD7G6712o        
+      root: An0therPa55
       etcd: 97Vk2{qe8o>S
-      s3: 8_d$8zCrl7£m                                            
+      s3: 8_d$8zCrl7£m