|
| 1 | +# Licensed to the Apache Software Foundation (ASF) under one |
| 2 | +# or more contributor license agreements. See the NOTICE file |
| 3 | +# distributed with this work for additional information |
| 4 | +# regarding copyright ownership. The ASF licenses this file |
| 5 | +# to you under the Apache License, Version 2.0 (the |
| 6 | +# "License"); you may not use this file except in compliance |
| 7 | +# with the License. You may obtain a copy of the License at |
| 8 | +# |
| 9 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 10 | +# |
| 11 | +# Unless required by applicable law or agreed to in writing, |
| 12 | +# software distributed under the License is distributed on an |
| 13 | +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| 14 | +# KIND, either express or implied. See the License for the |
| 15 | +# specific language governing permissions and limitations |
| 16 | +# under the License. |
| 17 | +# |
| 18 | +import logging |
| 19 | +import nuvolaris.config as cfg |
| 20 | +from pymilvus import MilvusClient |
| 21 | + |
| 22 | +class MilvusAdminClient: |
| 23 | + """ |
| 24 | + Simple Milvus Client used to perform Mivlus administration Tasks |
| 25 | + """ |
| 26 | + def __init__(self, db_name="default"): |
| 27 | + self.admin_username = cfg.get("milvus.admin.user", "MILVUS_ROOT_USER", "root") |
| 28 | + self.milvus_api_host = cfg.get("milvus.host", "MILVUS_API_HOST", "nuvolaris-milvus") |
| 29 | + self.milvus_api_port = cfg.get("milvus.host", "MILVUS_API_PORT", "19530") |
| 30 | + self.admin_password = cfg.get("milvus.password.root", "MILVUS_ROOT_PASSWORD", "An0therPa55") |
| 31 | + |
| 32 | + self.milvus_api_url = f"http://{self.milvus_api_host}:{self.milvus_api_port}" |
| 33 | + self.client = MilvusClient( |
| 34 | + uri=self.milvus_api_url, |
| 35 | + token=f"{self.admin_username}:{self.admin_password}", |
| 36 | + db_name=db_name |
| 37 | + ) |
| 38 | + |
| 39 | + print(f"{self.admin_username}:{self.admin_password}") |
| 40 | + |
| 41 | + def close_connection(self): |
| 42 | + try: |
| 43 | + self.client.close() |
| 44 | + logging.info("MILVUS client connection closed") |
| 45 | + except Exception as ex: |
| 46 | + logging.warning("cannot close MILVUS client connection", ex) |
| 47 | + |
| 48 | + def add_user(self, username, password): |
| 49 | + """ |
| 50 | + adds a new MILVUS user to the predefined |
| 51 | + param: username |
| 52 | + param: password |
| 53 | + return: True if user has been successfully created |
| 54 | + """ |
| 55 | + try: |
| 56 | + self.client.create_user(username, password) |
| 57 | + created_user = self.client.describe_user(username) |
| 58 | + return 'user_name' in created_user |
| 59 | + except Exception as ex: |
| 60 | + logging.error(f"Could not create milvus user {username}", ex) |
| 61 | + return False |
| 62 | + |
| 63 | + def add_role(self, role): |
| 64 | + """ |
| 65 | + adds a new MILVUS role |
| 66 | + param: role |
| 67 | + return: True if role has been successfully created |
| 68 | + """ |
| 69 | + try: |
| 70 | + self.client.create_role(role) |
| 71 | + created_role = self.client.describe_role(role) |
| 72 | + return 'role_name' in created_role |
| 73 | + except Exception as ex: |
| 74 | + logging.error(f"Could not create milvus role {role}", ex) |
| 75 | + return False |
| 76 | + |
| 77 | + def add_default_privileges_to_role(self, role): |
| 78 | + """ |
| 79 | + adds default privileges to a role |
| 80 | + param: role |
| 81 | + return: True if role has been successfully created |
| 82 | + """ |
| 83 | + try: |
| 84 | + self.client.grant_privilege( |
| 85 | + role_name=role, |
| 86 | + object_type='Global', # value here can be Global, Collection or User, object type also depends on the API defined in privilegeName |
| 87 | + object_name='*', # value here can be * or a specific user name if object type is 'User' |
| 88 | + privilege='CreateCollection' |
| 89 | + ) |
| 90 | + |
| 91 | + self.client.grant_privilege( |
| 92 | + role_name=role, |
| 93 | + object_type='Collection', # value here can be Global, Collection or User, object type also depends on the API defined in privilegeName |
| 94 | + object_name='*', # value here can be * or a specific user name if object type is 'User' |
| 95 | + privilege='*' |
| 96 | + ) |
| 97 | + |
| 98 | + return True |
| 99 | + except Exception as ex: |
| 100 | + logging.error(f"Could not create milvus role {role}", ex) |
| 101 | + return False |
| 102 | + |
| 103 | + def assign_role(self, username, role): |
| 104 | + """ |
| 105 | + assign a tole to a user |
| 106 | + param: username |
| 107 | + param: role |
| 108 | + return: True if role has been successfully created |
| 109 | + """ |
| 110 | + try: |
| 111 | + self.client.grant_role(user_name=username,role_name=role) |
| 112 | + user_detail = self.client.describe_user(username) |
| 113 | + return role in user_detail['roles'] |
| 114 | + except Exception as ex: |
| 115 | + logging.error(f"Could not assign MILVUS role {role}", ex) |
| 116 | + return False |
| 117 | + |
| 118 | + def setup_user(self, username, password): |
| 119 | + """ |
| 120 | + Creates a user into MILVUS and assign permission on collection |
| 121 | + param: username |
| 122 | + param: role |
| 123 | + return: True if role has been successfully created |
| 124 | + """ |
| 125 | + role = f"{username}_role" |
| 126 | + self.add_user(username,password) |
| 127 | + self.add_role(role) |
| 128 | + self.add_default_privileges_to_role(role) |
| 129 | + self.assign_role(username, role) |
0 commit comments