attu/express/src/partitions/partitions.controller.ts

import { NextFunction, Request, Response, Router } from "express";
import { dtoValidationMiddleware } from "../middlewares/validation";
import { PartitionsService } from "./partitions.service";
import { milvusService } from "../milvus";

import {
  GetPartitionsInfoDto,
  ManagePartitionDto,
  LoadPartitionsDto,
} from "./dto";

export class PartitionController {
  private router: Router;
  private partitionsService: PartitionsService;

  constructor() {
    this.partitionsService = new PartitionsService(milvusService);
    this.router = Router();
  }

  generateRoutes() {
    this.router.get(
      "/",
      dtoValidationMiddleware(GetPartitionsInfoDto),
      this.getPatitionsInfo.bind(this)
    );

    this.router.post(
      "/",
      dtoValidationMiddleware(ManagePartitionDto),
      this.managePartition.bind(this)
    );

    this.router.post(
      "/load",
      dtoValidationMiddleware(LoadPartitionsDto),
      this.loadPartition.bind(this)
    );

    this.router.post(
      "/release",
      dtoValidationMiddleware(LoadPartitionsDto),
      this.releasePartition.bind(this)
    );

    return this.router;
  }

  async getPatitionsInfo(req: Request, res: Response, next: NextFunction) {
    const collectionName = "" + req.query?.collection_name;
    try {
      const result = await this.partitionsService.getPatitionsInfo({
        collection_name: collectionName,
      });
      res.send(result);
    } catch (error) {
      next(error);
    }
  }

  async managePartition(req: Request, res: Response, next: NextFunction) {
    const { type, ...params } = req.body;
    try {
      const result =
        type.toLocaleLowerCase() === "create"
          ? await this.partitionsService.createParition(params)
          : await this.partitionsService.deleteParition(params);
      res.send(result);
    } catch (error) {
      next(error);
    }
  }

  async loadPartition(req: Request, res: Response, next: NextFunction) {
    const data = req.body;
    try {
      const result = await this.partitionsService.loadPartitions(data);
      res.send(result);
    } catch (error) {
      next(error);
    }
  }

  async releasePartition(req: Request, res: Response, next: NextFunction) {
    const data = req.body;
    try {
      const result = await this.partitionsService.releasePartitions(data);
      res.send(result);
    } catch (error) {
      next(error);
    }
  }
}