milvus/cpp/src/server/grpc_impl/GrpcRequestHandler.cpp

/*******************************************************************************
 * Copyright 上海赜睿信息科技有限公司(Zilliz) - All Rights Reserved
 * Unauthorized copying of this file, via any medium is strictly prohibited.
 * Proprietary and confidential.
 ******************************************************************************/

#include "GrpcRequestHandler.h"
#include "GrpcRequestTask.h"
#include "utils/TimeRecorder.h"

namespace zilliz {
namespace milvus {
namespace server {

::grpc::Status
GrpcRequestHandler::CreateTable(::grpc::ServerContext *context,
                                const ::milvus::grpc::TableSchema *request,
                                ::milvus::grpc::Status *response) {

    BaseTaskPtr task_ptr = CreateTableTask::Create(*request);
    GrpcRequestScheduler::ExecTask(task_ptr, response);
    return ::grpc::Status::OK;
}

::grpc::Status
GrpcRequestHandler::HasTable(::grpc::ServerContext *context,
                             const ::milvus::grpc::TableName *request,
                             ::milvus::grpc::BoolReply *response) {

    bool has_table = false;
    BaseTaskPtr task_ptr = HasTableTask::Create(request->table_name(), has_table);
    ::milvus::grpc::Status grpc_status;
    GrpcRequestScheduler::ExecTask(task_ptr, &grpc_status);
    response->set_bool_reply(has_table);
    response->mutable_status()->set_reason(grpc_status.reason());
    response->mutable_status()->set_error_code(grpc_status.error_code());
    return ::grpc::Status::OK;
}

::grpc::Status
GrpcRequestHandler::DropTable(::grpc::ServerContext* context,
                              const ::milvus::grpc::TableName* request,
                              ::milvus::grpc::Status* response) {

    BaseTaskPtr task_ptr = DropTableTask::Create(request->table_name());
    GrpcRequestScheduler::ExecTask(task_ptr, response);
    return ::grpc::Status::OK;
}

::grpc::Status
GrpcRequestHandler::BuildIndex(::grpc::ServerContext* context,
                               const ::milvus::grpc::TableName* request,
                               ::milvus::grpc::Status* response) {

    BaseTaskPtr task_ptr = BuildIndexTask::Create(request->table_name());
    GrpcRequestScheduler::ExecTask(task_ptr, response);
    return ::grpc::Status::OK;
}

::grpc::Status
GrpcRequestHandler::InsertVector(::grpc::ServerContext* context,
                                 const ::milvus::grpc::InsertInfos* request,
                                 ::milvus::grpc::VectorIds* response) {

    BaseTaskPtr task_ptr = InsertVectorTask::Create(*request, *response);
    ::milvus::grpc::Status grpc_status;
    GrpcRequestScheduler::ExecTask(task_ptr, &grpc_status);
    response->mutable_status()->set_reason(grpc_status.reason());
    response->mutable_status()->set_error_code(grpc_status.error_code());
    return ::grpc::Status::OK;
}

::grpc::Status
GrpcRequestHandler::SearchVector(::grpc::ServerContext* context,
                                 const ::milvus::grpc::SearchVectorInfos* request,
                                 ::grpc::ServerWriter<::milvus::grpc::TopKQueryResult>* writer) {

    std::vector<std::string> file_id_array;
    BaseTaskPtr task_ptr = SearchVectorTask::Create(*request, file_id_array, *writer);
    ::milvus::grpc::Status grpc_status;
    GrpcRequestScheduler::ExecTask(task_ptr, &grpc_status);
    if (grpc_status.error_code() != SERVER_SUCCESS) {
        ::grpc::Status status(::grpc::INVALID_ARGUMENT, grpc_status.reason());
        return status;
    } else {
        return ::grpc::Status::OK;
    }
}

::grpc::Status
GrpcRequestHandler::SearchVectorInFiles(::grpc::ServerContext* context,
                                        const ::milvus::grpc::SearchVectorInFilesInfos* request,
                                        ::grpc::ServerWriter<::milvus::grpc::TopKQueryResult>* writer) {

    std::vector<std::string> file_id_array;
    BaseTaskPtr task_ptr = SearchVectorTask::Create(request->search_vector_infos(), file_id_array, *writer);
    ::milvus::grpc::Status grpc_status;
    GrpcRequestScheduler::ExecTask(task_ptr, &grpc_status);
    if (grpc_status.error_code() != SERVER_SUCCESS) {
        ::grpc::Status status(::grpc::INVALID_ARGUMENT, grpc_status.reason());
        return status;
    } else {
        return ::grpc::Status::OK;
    }
}

::grpc::Status
GrpcRequestHandler::DescribeTable(::grpc::ServerContext* context,
                                  const ::milvus::grpc::TableName* request,
                                  ::milvus::grpc::TableSchema* response) {

    BaseTaskPtr task_ptr = DescribeTableTask::Create(request->table_name(), *response);
    ::milvus::grpc::Status grpc_status;
    GrpcRequestScheduler::ExecTask(task_ptr, &grpc_status);
    response->mutable_table_name()->mutable_status()->set_error_code(grpc_status.error_code());
    response->mutable_table_name()->mutable_status()->set_reason(grpc_status.reason());
    return ::grpc::Status::OK;
}

::grpc::Status
GrpcRequestHandler::GetTableRowCount(::grpc::ServerContext* context,
                                     const ::milvus::grpc::TableName* request,
                                     ::milvus::grpc::TableRowCount* response) {

    int64_t row_count = 0;
    BaseTaskPtr task_ptr = GetTableRowCountTask::Create(request->table_name(), row_count);
    ::milvus::grpc::Status grpc_status;
    GrpcRequestScheduler::ExecTask(task_ptr, &grpc_status);
    response->set_table_row_count(row_count);
    response->mutable_status()->set_reason(grpc_status.reason());
    response->mutable_status()->set_error_code(grpc_status.error_code());
    return ::grpc::Status::OK;
}

::grpc::Status
GrpcRequestHandler::ShowTables(::grpc::ServerContext* context,
                               const ::milvus::grpc::Command* request,
                               ::grpc::ServerWriter<::milvus::grpc::TableName>* writer) {

    BaseTaskPtr task_ptr = ShowTablesTask::Create(*writer);
    ::milvus::grpc::Status grpc_status;
    GrpcRequestScheduler::ExecTask(task_ptr, &grpc_status);
    if (grpc_status.error_code() != SERVER_SUCCESS) {
        ::grpc::Status status(::grpc::UNKNOWN, grpc_status.reason());
        return status;
    } else {
        return ::grpc::Status::OK;
    }
}

::grpc::Status
GrpcRequestHandler::Ping(::grpc::ServerContext* context,
                         const ::milvus::grpc::Command* request,
                         ::milvus::grpc::ServerStatus* response) {

    std::string result;
    BaseTaskPtr task_ptr = PingTask::Create(request->cmd(), result);
    ::milvus::grpc::Status grpc_status;
    GrpcRequestScheduler::ExecTask(task_ptr, &grpc_status);
    response->set_info(result);
    response->mutable_status()->set_reason(grpc_status.reason());
    response->mutable_status()->set_error_code(grpc_status.error_code());
    return ::grpc::Status::OK;
}


}
}
}