Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix: support db_name for APIs #356

Merged
merged 1 commit into from
Sep 3, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 21 additions & 8 deletions milvus/grpc/Collection.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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[] = [];
Expand Down Expand Up @@ -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
);

Expand Down
23 changes: 18 additions & 5 deletions milvus/grpc/Data.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
16 changes: 11 additions & 5 deletions milvus/grpc/MilvusIndex.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
102 changes: 100 additions & 2 deletions test/grpc/Database.spec.ts
Original file line number Diff line number Diff line change
@@ -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';
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand Down
Loading