enhance: enable setting properties during create database (#2168) (#2174

)

issue: milvus-io/milvus#30040
pr: #2168

Signed-off-by: Wei Liu <[email protected]>

examples/database.py

@@ -117,35 +117,38 @@ def collection_read_write(collection, db_name):
     # create collection within default
     col1_db1 = create_collection("col1_db1", "default")
 
+    db1Name = "db1"
     # create db1
-    if "db1" not in db.list_database():
+    if db1Name not in db.list_database():
         print("\ncreate database: db1")
-        db.create_database(db_name="db1")
+        db.create_database(db_name=db1Name, properties={"key1":"value1"})
+        db_info = db.describe_database(db_name=db1Name)
+        print(db_info)
 
     # use database db1
-    db.using_database(db_name="db1")
+    db.using_database(db_name=db1Name)
     # create collection within default
-    col2_db1 = create_collection("col1_db1", "db1")
+    col2_db1 = create_collection("col1_db1", db1Name)
 
     # verify read and write
-    collection_read_write(col2_db1, "db1")
+    collection_read_write(col2_db1, db1Name)
 
     # list collections within db1
     print("\nlist collections of database db1:")
     print(utility.list_collections())
 
     # set properties of db1
-    print("\nset properties of db1:")
-    db_info = db.describe_database(db_name="db1")
+    db_info = db.describe_database(db_name=db1Name)
     print(db_info)
-    db.set_properties(db_name="db1", properties={"key": "value"})
-    db_info = db.describe_database(db_name="db1")
+    print("\nset properties of db1:")
+    db.set_properties(db_name=db1Name, properties={"key": "value"})
+    db_info = db.describe_database(db_name=db1Name)
     print(db_info)
 
     print("\ndrop collection: col1_db2 from db1")
     col2_db1.drop()
     print("\ndrop database: db1")
-    db.drop_database(db_name="db1")
+    db.drop_database(db_name=db1Name)
 
     # list database
     print("\nlist databases:")

pymilvus/client/grpc_handler.py

@@ -1277,8 +1277,8 @@ def get_loading_progress(
         return response.progress
 
     @retry_on_rpc_failure()
-    def create_database(self, db_name: str, timeout: Optional[float] = None):
-        request = Prepare.create_database_req(db_name)
+    def create_database(self, db_name: str, timeout: Optional[float] = None, **kwargs):
+        request = Prepare.create_database_req(db_name, **kwargs)
         status = self._stub.CreateDatabase(request, timeout=timeout)
         check_status(status)
 

pymilvus/client/prepare.py

@@ -1227,9 +1227,17 @@ def register_request(cls, user: str, host: str, **kwargs):
         )
 
     @classmethod
-    def create_database_req(cls, db_name: str):
+    def create_database_req(cls, db_name: str, **kwargs):
         check_pass_param(db_name=db_name)
-        return milvus_types.CreateDatabaseRequest(db_name=db_name)
+
+        req = milvus_types.CreateDatabaseRequest(db_name=db_name)
+        properties = kwargs.get("properties")
+        if is_legal_collection_properties(properties):
+            properties = [
+                common_types.KeyValuePair(key=str(k), value=str(v)) for k, v in properties.items()
+            ]
+            req.properties.extend(properties)
+        return req
 
     @classmethod
     def drop_database_req(cls, db_name: str):

pymilvus/orm/collection.py

@@ -347,7 +347,9 @@ def set_properties(self, properties: dict, timeout: Optional[float] = None, **kw
 
         Args:
             properties (``dict``): collection properties.
-                 only support collection TTL with key `collection.ttl.seconds`
+                 support collection TTL with key `collection.ttl.seconds`
+                 support collection replica number with key `collection.replica.number`
+                 support collection resource groups with key `collection.resource_groups`.
             timeout (float, optional): an optional duration of time in seconds to allow
                 for the RPCs. If timeout is not set, the client keeps waiting until the
                 server responds or an error occurs.

pymilvus/orm/db.py

@@ -17,14 +17,17 @@ def using_database(db_name: str, using: str = "default"):
     _get_connection(using).reset_db_name(db_name)
 
 
-def create_database(db_name: str, using: str = "default", timeout: Optional[float] = None):
+def create_database(
+    db_name: str, using: str = "default", timeout: Optional[float] = None, **kwargs
+):
     """Create a database using provided database name
-
-    :param db_name: Database name
-    :type  db_name: str
-
+    Args:
+        db_name (``str``): Database name
+        properties (``dict``): database properties.
+            support database replica number with key `database.replica.number`
+            support database resource groups with key `database.resource_groups`
     """
-    _get_connection(using).create_database(db_name, timeout=timeout)
+    _get_connection(using).create_database(db_name, timeout=timeout, **kwargs)
 
 
 def drop_database(db_name: str, using: str = "default", timeout: Optional[float] = None):
@@ -52,14 +55,12 @@ def set_properties(
     using: str = "default",
     timeout: Optional[float] = None,
 ):
-    """Alter a database using provided database name
-
-    :param db_name: Database name
-    :type  db_name: str
-
-    :param properties: Properties to be set
-    :type  properties: dict
-
+    """Set properties for a database using provided database name
+    Args:
+        db_name (``str``): Database name
+        properties (``dict``): database properties.
+            support database replica number with key `database.replica.number`
+            support database resource groups with key `database.resource_groups`
     """
     _get_connection(using).alter_database(db_name, properties=properties, timeout=timeout)