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

issue: milvus-io/milvus#30040

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

@@ -1233,9 +1233,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/grpc_gen/milvus_pb2.pyi

@@ -1863,12 +1863,14 @@ class AllocTimestampResponse(_message.Message):
     def __init__(self, status: _Optional[_Union[_common_pb2.Status, _Mapping]] = ..., timestamp: _Optional[int] = ...) -> None: ...
 
 class CreateDatabaseRequest(_message.Message):
-    __slots__ = ("base", "db_name")
+    __slots__ = ("base", "db_name", "properties")
     BASE_FIELD_NUMBER: _ClassVar[int]
     DB_NAME_FIELD_NUMBER: _ClassVar[int]
+    PROPERTIES_FIELD_NUMBER: _ClassVar[int]
     base: _common_pb2.MsgBase
     db_name: str
-    def __init__(self, base: _Optional[_Union[_common_pb2.MsgBase, _Mapping]] = ..., db_name: _Optional[str] = ...) -> None: ...
+    properties: _containers.RepeatedCompositeFieldContainer[_common_pb2.KeyValuePair]
+    def __init__(self, base: _Optional[_Union[_common_pb2.MsgBase, _Mapping]] = ..., db_name: _Optional[str] = ..., properties: _Optional[_Iterable[_Union[_common_pb2.KeyValuePair, _Mapping]]] = ...) -> None: ...
 
 class DropDatabaseRequest(_message.Message):
     __slots__ = ("base", "db_name")

pymilvus/grpc_gen/schema_pb2.pyi

@@ -91,20 +91,22 @@ class FieldSchema(_message.Message):
     def __init__(self, fieldID: _Optional[int] = ..., name: _Optional[str] = ..., is_primary_key: bool = ..., description: _Optional[str] = ..., data_type: _Optional[_Union[DataType, str]] = ..., type_params: _Optional[_Iterable[_Union[_common_pb2.KeyValuePair, _Mapping]]] = ..., index_params: _Optional[_Iterable[_Union[_common_pb2.KeyValuePair, _Mapping]]] = ..., autoID: bool = ..., state: _Optional[_Union[FieldState, str]] = ..., element_type: _Optional[_Union[DataType, str]] = ..., default_value: _Optional[_Union[ValueField, _Mapping]] = ..., is_dynamic: bool = ..., is_partition_key: bool = ..., is_clustering_key: bool = ..., nullable: bool = ...) -> None: ...
 
 class CollectionSchema(_message.Message):
-    __slots__ = ("name", "description", "autoID", "fields", "enable_dynamic_field", "properties")
+    __slots__ = ("name", "description", "autoID", "fields", "enable_dynamic_field", "properties", "partition_key_isolation")
     NAME_FIELD_NUMBER: _ClassVar[int]
     DESCRIPTION_FIELD_NUMBER: _ClassVar[int]
     AUTOID_FIELD_NUMBER: _ClassVar[int]
     FIELDS_FIELD_NUMBER: _ClassVar[int]
     ENABLE_DYNAMIC_FIELD_FIELD_NUMBER: _ClassVar[int]
     PROPERTIES_FIELD_NUMBER: _ClassVar[int]
+    PARTITION_KEY_ISOLATION_FIELD_NUMBER: _ClassVar[int]
     name: str
     description: str
     autoID: bool
     fields: _containers.RepeatedCompositeFieldContainer[FieldSchema]
     enable_dynamic_field: bool
     properties: _containers.RepeatedCompositeFieldContainer[_common_pb2.KeyValuePair]
-    def __init__(self, name: _Optional[str] = ..., description: _Optional[str] = ..., autoID: bool = ..., fields: _Optional[_Iterable[_Union[FieldSchema, _Mapping]]] = ..., enable_dynamic_field: bool = ..., properties: _Optional[_Iterable[_Union[_common_pb2.KeyValuePair, _Mapping]]] = ...) -> None: ...
+    partition_key_isolation: bool
+    def __init__(self, name: _Optional[str] = ..., description: _Optional[str] = ..., autoID: bool = ..., fields: _Optional[_Iterable[_Union[FieldSchema, _Mapping]]] = ..., enable_dynamic_field: bool = ..., properties: _Optional[_Iterable[_Union[_common_pb2.KeyValuePair, _Mapping]]] = ..., partition_key_isolation: bool = ...) -> None: ...
 
 class BoolArray(_message.Message):
     __slots__ = ("data",)

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)