From 46bfda5fcb42129c02fece0a60be92429ac3e09c Mon Sep 17 00:00:00 2001
From: ryjiang <jiangruiyi@gmail.com>
Date: Tue, 13 Aug 2024 15:03:21 +0800
Subject: [PATCH] feat: support set properties on collection creation

Signed-off-by: ryjiang <jiangruiyi@gmail.com>
---
 milvus/grpc/Collection.ts    | 23 ++++++++++++++++-------
 milvus/types/Collection.ts   |  5 ++++-
 test/grpc/Collection.spec.ts | 33 +++++++++++++++++++++++++++++++++
 3 files changed, 53 insertions(+), 8 deletions(-)

diff --git a/milvus/grpc/Collection.ts b/milvus/grpc/Collection.ts
index bf7e1b5f..153bcca4 100644
--- a/milvus/grpc/Collection.ts
+++ b/milvus/grpc/Collection.ts
@@ -86,6 +86,7 @@ export class Collection extends Database {
    * @param {number} [data.num_partitions] - The number of partitions allowed in the new collection.
    * @param {string} [data.consistency_level] - The consistency level of the new collection. Can be "Strong" (Milvus default), "Session", "Bounded", "Eventually", or "Customized".
    * @param {FieldType[]} data.fields - The fields of the new collection. See [FieldType](https://github.com/milvus-io/milvus-sdk-node/blob/main/milvus/types/Collection.ts#L8) for more details.
+   * @param {properties} [data.properties] - An optional object containing key-value pairs of properties for the collection.
    * @param {number} [data.timeout] - An optional duration of time in milliseconds to allow for the RPC. If it is set to undefined, the client keeps waiting until the server responds or error occurs. Default is undefined.
    *
    * @returns {Promise<ResStatus>} The response status of the operation.
@@ -160,17 +161,25 @@ export class Collection extends Database {
     const level =
       ConsistencyLevelEnum[consistency_level] ?? ConsistencyLevelEnum.Bounded;
 
+    // build the request object
+    const req: any = {
+      ...data,
+      schema: schemaBytes,
+      consistency_level: level,
+      enable_dynamic_field:
+        data.enableDynamicField || data.enable_dynamic_field,
+    };
+
+    // if properties is set, parse it to key-value pairs
+    if (data.properties) {
+      req.properties = parseToKeyValue(data.properties);
+    }
+
     // Call the promisify function to create the collection.
     const createPromise = await promisify(
       this.channelPool,
       'CreateCollection',
-      {
-        ...data,
-        schema: schemaBytes,
-        consistency_level: level,
-        enable_dynamic_field:
-          data.enableDynamicField || data.enable_dynamic_field,
-      },
+      req,
       data.timeout || this.timeout
     );
 
diff --git a/milvus/types/Collection.ts b/milvus/types/Collection.ts
index a0d6df23..dc1e79d5 100644
--- a/milvus/types/Collection.ts
+++ b/milvus/types/Collection.ts
@@ -82,6 +82,8 @@ export interface ShowCollectionsReq extends GrpcTimeOut {
   collection_names?: string[];
 }
 
+export type Properties = Record<string, string | number | boolean>;
+
 export interface BaseCreateCollectionReq extends GrpcTimeOut {
   // collection name
   collection_name: string; // required, collection name
@@ -97,6 +99,7 @@ export interface BaseCreateCollectionReq extends GrpcTimeOut {
   partition_key_field?: string; // optional, partition key field
   enable_dynamic_field?: boolean; // optional, enable dynamic field, default is false
   enableDynamicField?: boolean; // optional, alias of enable_dynamic_field
+  properties?: Properties;
 }
 
 export interface CreateCollectionWithFieldsReq extends BaseCreateCollectionReq {
@@ -237,7 +240,7 @@ export interface GetLoadStateResponse extends resStatusResponse {
 }
 
 export interface AlterCollectionReq extends collectionNameReq {
-  properties: Record<string, string | number | boolean>;
+  properties: Properties;
 }
 
 export interface DescribeAliasResponse extends resStatusResponse {
diff --git a/test/grpc/Collection.spec.ts b/test/grpc/Collection.spec.ts
index b6ae5986..33d56e24 100644
--- a/test/grpc/Collection.spec.ts
+++ b/test/grpc/Collection.spec.ts
@@ -22,6 +22,7 @@ const NEW_COLLECTION_NAME = GENERATE_NAME();
 const TEST_CONSISTENCY_LEVEL_COLLECTION_NAME = GENERATE_NAME();
 const LOAD_COLLECTION_NAME = GENERATE_NAME();
 const LOAD_COLLECTION_NAME_SYNC = GENERATE_NAME();
+const COLLECTION_WITH_PROPERTY = GENERATE_NAME();
 const ALIAS = 'my_alias';
 const NON_EXISTENT_COLLECTION_NAME = 'none_existent';
 
@@ -51,6 +52,38 @@ describe(`Collection API`, () => {
     expect(res.error_code).toEqual(ErrorCode.SUCCESS);
   });
 
+  it(`Create Collection with property set should be successful`, async () => {
+    const res = await milvusClient.createCollection({
+      ...genCollectionParams({
+        collectionName: COLLECTION_WITH_PROPERTY,
+        dim: [128],
+      }),
+      properties: {
+        'collection.ttl.seconds': 1000,
+        'mmap.enabled': true,
+      },
+    });
+    expect(res.error_code).toEqual(ErrorCode.SUCCESS);
+
+    const describe = await milvusClient.describeCollection({
+      collection_name: COLLECTION_WITH_PROPERTY,
+    });
+
+    expect(
+      Number(
+        formatKeyValueData(describe.properties, ['collection.ttl.seconds'])[
+          'collection.ttl.seconds'
+        ]
+      )
+    ).toEqual(1000);
+
+    expect(
+      Boolean(
+        formatKeyValueData(describe.properties, ['mmap.enabled'])['mmap.enabled']
+      )
+    ).toEqual(true);
+  });
+
   it(`Should get pk fieldname successfully`, async () => {
     const res = await milvusClient.getPkFieldName({
       collection_name: COLLECTION_NAME,