From f581ea931c4abc6248b64de1e9a0ef2e230bb06c Mon Sep 17 00:00:00 2001 From: ryjiang Date: Tue, 3 Sep 2024 13:58:45 +0800 Subject: [PATCH] fix: support db_name for APIs Signed-off-by: ryjiang --- milvus/grpc/Collection.ts | 29 ++++++++--- milvus/grpc/Data.ts | 23 +++++++-- milvus/grpc/MilvusIndex.ts | 16 ++++-- test/grpc/Database.spec.ts | 102 ++++++++++++++++++++++++++++++++++++- 4 files changed, 150 insertions(+), 20 deletions(-) diff --git a/milvus/grpc/Collection.ts b/milvus/grpc/Collection.ts index 153bcca4..a1bf553d 100644 --- a/milvus/grpc/Collection.ts +++ b/milvus/grpc/Collection.ts @@ -250,13 +250,19 @@ export class Collection extends Database { async showCollections( data?: ShowCollectionsReq ): Promise { + 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 { 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 ); diff --git a/milvus/grpc/Data.ts b/milvus/grpc/Data.ts index db498b17..298a5b62 100644 --- a/milvus/grpc/Data.ts +++ b/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 { - 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, diff --git a/milvus/grpc/MilvusIndex.ts b/milvus/grpc/MilvusIndex.ts index e1ce19da..48044bf4 100644 --- a/milvus/grpc/MilvusIndex.ts +++ b/milvus/grpc/MilvusIndex.ts @@ -340,14 +340,20 @@ export class Index extends Data { */ async alterIndex(data: AlterIndexReq): Promise { 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; diff --git a/test/grpc/Database.spec.ts b/test/grpc/Database.spec.ts index a2d30e3c..73024419 100644 --- a/test/grpc/Database.spec.ts +++ b/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,14 +118,77 @@ 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, @@ -113,6 +196,21 @@ describe(`Database API`, () => { }); 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,