Merge pull request #5 from zilliztech/uxd-637

[UXD-637] Add partitions api

server/src/app.module.ts

@@ -10,6 +10,7 @@ import { UsersModule } from './users/users.module';
 import { AuthModule } from './auth/auth.module';
 import { ServeStaticModule } from '@nestjs/serve-static';
 import { join } from 'path';
+import { PartitionsModule } from './partitions/partitions.module';
 
 @Module({
   imports: [
@@ -21,6 +22,7 @@ import { join } from 'path';
     CollectionsModule,
     UsersModule,
     AuthModule,
+    PartitionsModule,
   ],
   controllers: [AppController],
   providers: [

server/src/collections/collections.service.ts

@@ -9,7 +9,9 @@ import {
   LoadCollectionReq,
   ReleaseLoadCollectionReq,
 } from '@zilliz/milvus-sdk-node-dev/dist/milvus/types'; // todo: need improve like export types in root file.
-import { throwErrorFromSDK } from 'src/utils/Error';
+import { throwErrorFromSDK } from '../utils/Error';
+import { findKeyValue } from '../utils/Helper';
+import { ROW_COUNT } from '../utils/Const';
 @Injectable()
 export class CollectionsService {
   constructor(private milvusService: MilvusService) {}
@@ -25,13 +27,9 @@ export class CollectionsService {
   }
 
   async createCollection(data: CreateCollectionReq) {
-    try {
-      const res = await this.milvusClient.createCollection(data);
-      throwErrorFromSDK(res);
-      return res;
-    } catch (error) {
-      throw new Error(error);
-    }
+    const res = await this.milvusClient.createCollection(data);
+    throwErrorFromSDK(res);
+    return res;
   }
 
   async describeCollection(data: DescribeCollectionReq) {
@@ -73,6 +71,7 @@ export class CollectionsService {
    */
   async getIndexStatus(data: GetIndexStateReq) {
     const res = await this.milvusClient.getIndexState(data);
+    throwErrorFromSDK(res.status);
     return res;
   }
 
@@ -89,12 +88,10 @@ export class CollectionsService {
         });
         data.push({
           collection_name: name,
-          // schema: collectionInfo.schema,
+          schema: collectionInfo.schema,
           description: collectionInfo.schema.description,
           autoID: collectionInfo.schema.autoID,
-          rowCount: collectionStatistics.stats.find(
-            (v) => v.key === 'row_count',
-          ).value,
+          rowCount: findKeyValue(collectionStatistics.stats, ROW_COUNT),
           // id: collectionInfo.collectionId
         });
       }

server/src/partitions/dto.ts

@@ -0,0 +1,51 @@
+import {
+  IsNotEmpty,
+  IsString,
+  IsEnum,
+  IsArray,
+  ArrayNotEmpty,
+} from 'class-validator';
+
+export enum ManageType {
+  DELETE = 'delete',
+  CREATE = 'create',
+}
+export class GetPartitionsInfo {
+  @IsString()
+  @IsNotEmpty({
+    message: 'collection_name is empty',
+  })
+  readonly collection_name: string;
+}
+
+export class ManagePartition {
+  @IsString()
+  @IsNotEmpty({
+    message: 'collection_name is empty',
+  })
+  readonly collection_name: string;
+
+  @IsString()
+  @IsNotEmpty({
+    message: 'partition_name is empty',
+  })
+  readonly partition_name: string;
+
+  @IsEnum(ManageType, { message: 'Type allow delete and create' })
+  readonly type: ManageType;
+}
+
+export class LoadPartitions {
+  @IsString()
+  @IsNotEmpty({
+    message: 'collection_name is empty',
+  })
+  readonly collection_name: string;
+
+  @IsArray()
+  @ArrayNotEmpty()
+  @IsNotEmpty({
+    message: 'partition_names is empty',
+  })
+  readonly partition_names: string[];
+}

server/src/partitions/partitions.controller.spec.ts

@@ -0,0 +1,18 @@
+import { Test, TestingModule } from '@nestjs/testing';
+import { PartitionsController } from './partitions.controller';
+
+describe('PartitionsController', () => {
+  let controller: PartitionsController;
+
+  beforeEach(async () => {
+    const module: TestingModule = await Test.createTestingModule({
+      controllers: [PartitionsController],
+    }).compile();
+
+    controller = module.get<PartitionsController>(PartitionsController);
+  });
+
+  it('should be defined', () => {
+    expect(controller).toBeDefined();
+  });
+});

server/src/partitions/partitions.controller.ts

@@ -0,0 +1,50 @@
+import {
+  Body,
+  Controller,
+  Get,
+  Post,
+  Put,
+  Query,
+  UsePipes,
+  ValidationPipe,
+} from '@nestjs/common';
+import {
+  GetPartitionsInfo,
+  LoadPartitions,
+  ManagePartition,
+  ManageType,
+} from './dto';
+import { PartitionsService } from './partitions.service';
+
+@Controller('partitions')
+export class PartitionsController {
+  constructor(private partitionsService: PartitionsService) {}
+
+  @Get()
+  @UsePipes(new ValidationPipe())
+  async getPartitions(@Query() query: GetPartitionsInfo) {
+    return await this.partitionsService.getPatitionsInfo(query);
+  }
+
+  @Post()
+  @UsePipes(new ValidationPipe())
+  async managePartition(@Body() body: ManagePartition) {
+    const { type, ...params } = body;
+
+    return type.toLocaleLowerCase() === ManageType.CREATE
+      ? await this.partitionsService.createParition(params)
+      : await this.partitionsService.deleteParition(params);
+  }
+
+  @Put('load')
+  @UsePipes(new ValidationPipe())
+  async loadPartition(@Body() body: LoadPartitions) {
+    return await this.partitionsService.loadPartitions(body);
+  }
+
+  @Put('release')
+  @UsePipes(new ValidationPipe())
+  async releasePartitions(@Body() body: LoadPartitions) {
+    return await this.partitionsService.loadPartitions(body);
+  }
+}

server/src/partitions/partitions.module.ts

@@ -0,0 +1,11 @@
+import { Module } from '@nestjs/common';
+import { PartitionsService } from './partitions.service';
+import { PartitionsController } from './partitions.controller';
+import { MilvusModule } from 'src/milvus/milvus.module';
+
+@Module({
+  imports: [MilvusModule],
+  providers: [PartitionsService],
+  controllers: [PartitionsController],
+})
+export class PartitionsModule {}

server/src/partitions/partitions.service.spec.ts

@@ -0,0 +1,18 @@
+import { Test, TestingModule } from '@nestjs/testing';
+import { PartitionsService } from './partitions.service';
+
+describe('PartitionsService', () => {
+  let service: PartitionsService;
+
+  beforeEach(async () => {
+    const module: TestingModule = await Test.createTestingModule({
+      providers: [PartitionsService],
+    }).compile();
+
+    service = module.get<PartitionsService>(PartitionsService);
+  });
+
+  it('should be defined', () => {
+    expect(service).toBeDefined();
+  });
+});

server/src/partitions/partitions.service.ts

@@ -0,0 +1,77 @@
+import { Injectable } from '@nestjs/common';
+import { MilvusService } from '../milvus/milvus.service';
+import {
+  CreatePartitionReq,
+  DropPartitionReq,
+  GetPartitionStatisticsReq,
+  LoadPartitionsReq,
+  ReleasePartitionsReq,
+  ShowPartitionsReq,
+} from '@zilliz/milvus-sdk-node-dev/dist/milvus/types'; // todo: need improve like export types in root file.
+import { throwErrorFromSDK } from 'src/utils/Error';
+import { findKeyValue } from '../utils/Helper';
+import { ROW_COUNT } from '../utils/Const';
+
+@Injectable()
+export class PartitionsService {
+  constructor(private milvusService: MilvusService) {}
+
+  get milvusClient() {
+    return this.milvusService.milvusClientGetter;
+  }
+
+  async getPatitionsInfo(data: ShowPartitionsReq) {
+    const result = [];
+    const res = await this.getPartitions(data);
+    if (res.partition_names && res.partition_names.length) {
+      for (const [index, name] of res.partition_names.entries()) {
+        const statistics = await this.getPartitionStatistics({
+          ...data,
+          partition_name: name,
+        });
+        result.push({
+          name,
+          id: res.partitionIDs[index],
+          rowCount: findKeyValue(statistics.stats, ROW_COUNT),
+        });
+      }
+    }
+    return result;
+  }
+
+  async getPartitions(data: ShowPartitionsReq) {
+    const res = await this.milvusClient.showPartitions(data);
+    throwErrorFromSDK(res.status);
+    return res;
+  }
+
+  async createParition(data: CreatePartitionReq) {
+    const res = await this.milvusClient.createPartition(data);
+    throwErrorFromSDK(res);
+    return res;
+  }
+
+  async deleteParition(data: DropPartitionReq) {
+    const res = await this.milvusClient.dropPartition(data);
+    throwErrorFromSDK(res);
+    return res;
+  }
+
+  async getPartitionStatistics(data: GetPartitionStatisticsReq) {
+    const res = await this.milvusClient.getPartitionStatistics(data);
+    throwErrorFromSDK(res.status);
+    return res;
+  }
+
+  async loadPartitions(data: LoadPartitionsReq) {
+    const res = await this.milvusClient.loadPartitions(data);
+    throwErrorFromSDK(res);
+    return res;
+  }
+
+  async releasePartitions(data: ReleasePartitionsReq) {
+    const res = await this.milvusClient.releasePartitions(data);
+    throwErrorFromSDK(res);
+    return res;
+  }
+}

server/src/utils/Const.ts

@@ -0,0 +1 @@
+export const ROW_COUNT = 'row_count';

server/src/utils/Helper.ts

@@ -0,0 +1,4 @@
+import { KeyValuePair } from '@zilliz/milvus-sdk-node-dev/dist/milvus/types/Common';
+
+export const findKeyValue = (obj: KeyValuePair[], key: string) =>
+  obj.find((v) => v.key === key)?.value;