fix: support db_name for APIs (#356)

Signed-off-by: ryjiang <[email protected]>

milvus/grpc/Collection.ts

@@ -250,13 +250,19 @@ export class Collection extends Database {
   async showCollections(
     data?: ShowCollectionsReq
   ): Promise<ShowCollectionsResponse> {
+    const req: any = {
+      type: data ? data.type : ShowCollectionsType.All,
+      collection_names: data?.collection_names || [],
+    };
+
+    if (data?.db_name) {
+      req.db_name = data.db_name;
+    }
+
     const promise = await promisify(
       this.channelPool,
       'ShowCollections',
-      {
-        type: data ? data.type : ShowCollectionsType.All,
-        collection_names: data?.collection_names || [],
-      },
+      req,
       data?.timeout || this.timeout
     );
     const result: CollectionData[] = [];
@@ -296,13 +302,20 @@ export class Collection extends Database {
    */
   async alterCollection(data: AlterCollectionReq): Promise<ResStatus> {
     checkCollectionName(data);
+
+    const req: any = {
+      collection_name: data.collection_name,
+      properties: parseToKeyValue(data.properties),
+    };
+
+    if (data.db_name) {
+      req.db_name = data.db_name;
+    }
+
     const promise = await promisify(
       this.channelPool,
       'AlterCollection',
-      {
-        collection_name: data.collection_name,
-        properties: parseToKeyValue(data.properties),
-      },
+      req,
       data?.timeout || this.timeout
     );
 

milvus/grpc/Data.ts

@@ -127,9 +127,13 @@ export class Data extends Collection {
       throw new Error(ERROR_REASONS.INSERT_CHECK_FIELD_DATA_IS_REQUIRED);
     }
     const { collection_name } = data;
-    const collectionInfo = await this.describeCollection({
-      collection_name,
-    });
+
+    const describeReq = { collection_name };
+    if (data.db_name) {
+      (describeReq as any).db_name = data.db_name;
+    }
+
+    const collectionInfo = await this.describeCollection(describeReq);
 
     if (collectionInfo.status.error_code !== ErrorCode.SUCCESS) {
       throw collectionInfo;
@@ -454,6 +458,10 @@ export class Data extends Collection {
       this.milvusProto
     );
 
+    if (data.db_name) {
+      (request as any).db_name = data.db_name;
+    }
+
     // execute search
     const originSearchResult: SearchRes = await promisify(
       this.channelPool,
@@ -962,11 +970,16 @@ export class Data extends Collection {
   }
 
   async count(data: CountReq): Promise<CountResult> {
-    const queryResult = await this.query({
+    const req: any = {
       collection_name: data.collection_name,
       expr: data.expr || '',
       output_fields: [DEFAULT_COUNT_QUERY_STRING],
-    });
+    };
+
+    if (data.db_name) {
+      req.db_name = data.db_name;
+    }
+    const queryResult = await this.query(req);
 
     return {
       status: queryResult.status,

milvus/grpc/MilvusIndex.ts

@@ -340,14 +340,20 @@ export class Index extends Data {
    */
   async alterIndex(data: AlterIndexReq): Promise<ResStatus> {
     checkCollectionName(data);
+    const req = {
+      collection_name: data.collection_name,
+      index_name: data.index_name,
+      extra_params: parseToKeyValue(data.params),
+    } as any;
+
+    if (data.db_name) {
+      req.db_name = data.db_name;
+    }
+
     const promise = await promisify(
       this.channelPool,
       'AlterIndex',
-      {
-        collection_name: data.collection_name,
-        index_name: data.index_name,
-        extra_params: parseToKeyValue(data.params),
-      },
+      req,
       data.timeout || this.timeout
     );
     return promise;

test/grpc/Database.spec.ts

@@ -1,5 +1,10 @@
 import { MilvusClient, ErrorCode, DEFAULT_DB } from '../../milvus';
-import { IP, genCollectionParams, GENERATE_NAME } from '../tools';
+import {
+  IP,
+  genCollectionParams,
+  GENERATE_NAME,
+  generateInsertData,
+} from '../tools';
 
 let milvusClient = new MilvusClient({ address: IP, logLevel: 'info' });
 const DEFAULT = 'default';
@@ -78,12 +83,27 @@ describe(`Database API`, () => {
 
   it(`using db_name in API should be ok`, async () => {
     // create collection on another db
+    const params = genCollectionParams({
+      collectionName: COLLECTION_NAME2,
+      dim: [4],
+    });
     const createCollection = await milvusClient.createCollection({
-      ...genCollectionParams({ collectionName: COLLECTION_NAME2, dim: [4] }),
+      ...params,
       db_name: DB_NAME2,
     });
     expect(createCollection.error_code).toEqual(ErrorCode.SUCCESS);
 
+    // gen data
+    const vectors = generateInsertData(params.fields, 5);
+
+    // insert
+    const insert = await milvusClient.insert({
+      collection_name: COLLECTION_NAME2,
+      db_name: DB_NAME2,
+      data: vectors,
+    });
+    expect(insert.status.error_code).toEqual(ErrorCode.SUCCESS);
+
     // describe collection
     const describeCollection = await milvusClient.describeCollection({
       collection_name: COLLECTION_NAME2,
@@ -98,21 +118,99 @@ describe(`Database API`, () => {
       ErrorCode.UnexpectedError
     );
 
+    // alterCollection
+    const alterCollection = await milvusClient.alterCollection({
+      collection_name: COLLECTION_NAME2,
+      db_name: DB_NAME2,
+      properties: { 'collection.segment.rowLimit': 10000 },
+    });
+    expect(alterCollection.error_code).toEqual(ErrorCode.SUCCESS);
+    const describeCollectionAfterAlter = await milvusClient.describeCollection({
+      collection_name: COLLECTION_NAME2,
+      db_name: DB_NAME2,
+    });
+    expect(describeCollectionAfterAlter.properties).toEqual([
+      { key: 'collection.segment.rowLimit', value: '10000' },
+    ]);
+
+    // show collections
+    const showCollections = await milvusClient.showCollections({
+      db_name: DB_NAME2,
+    });
+    expect(showCollections.data.length).toBeGreaterThan(0);
+
+    // create partition
+    const createPartition = await milvusClient.createPartition({
+      collection_name: COLLECTION_NAME2,
+      db_name: DB_NAME2,
+      partition_name: 'partition1',
+    });
+    expect(createPartition.error_code).toEqual(ErrorCode.SUCCESS);
+
+    // show partitions
+    const showPartitions = await milvusClient.showPartitions({
+      collection_name: COLLECTION_NAME2,
+      db_name: DB_NAME2,
+    });
+    expect(showPartitions.partition_names).toContain('partition1');
+
+    // getCollectionStatistics
+    const getCollectionStatistics = await milvusClient.getCollectionStatistics({
+      collection_name: COLLECTION_NAME2,
+      db_name: DB_NAME2,
+    });
+    expect(getCollectionStatistics.status.error_code).toEqual(
+      ErrorCode.SUCCESS
+    );
+
     // create index
     const createIndex = await milvusClient.createIndex({
       collection_name: COLLECTION_NAME2,
       db_name: DB_NAME2,
       field_name: 'vector',
+      index_name: 'vector2',
     });
     expect(createIndex.error_code).toEqual(ErrorCode.SUCCESS);
 
+    // alter index
+    const alterIndex = await milvusClient.alterIndex({
+      collection_name: COLLECTION_NAME2,
+      db_name: DB_NAME2,
+      index_name: 'vector2',
+      params: { 'mmap.enabled': true },
+    });
+    expect(alterIndex.error_code).toEqual(ErrorCode.SUCCESS);
+
+    // describe index
+    const describeIndex = await milvusClient.describeIndex({
+      collection_name: COLLECTION_NAME2,
+      db_name: DB_NAME2,
+      index_name: 'vector2',
+    });
+    expect(describeIndex.index_descriptions[0].index_name).toEqual('vector2');
+
     // load collection
     const loadCollection = await milvusClient.loadCollection({
       collection_name: COLLECTION_NAME2,
       db_name: DB_NAME2,
     });
     expect(loadCollection.error_code).toEqual(ErrorCode.SUCCESS);
 
+    // query
+    const query = await milvusClient.count({
+      collection_name: COLLECTION_NAME2,
+      db_name: DB_NAME2,
+    });
+    expect(query.status.error_code).toEqual(ErrorCode.SUCCESS);
+
+    // search
+    const search = await milvusClient.search({
+      collection_name: COLLECTION_NAME2,
+      db_name: DB_NAME2,
+      data: [1, 2, 3, 4],
+    });
+    expect(search.status.error_code).toEqual(ErrorCode.SUCCESS);
+
     // release collection
     const releaseCollection = await milvusClient.releaseCollection({
       collection_name: COLLECTION_NAME2,