From 086e1acd8dc0c5c76880df0dd2ab68871eb9cdbc Mon Sep 17 00:00:00 2001 From: Ruichen Bao Date: Sat, 25 Jan 2025 17:50:11 +0800 Subject: [PATCH] feat: support RESTful, add CreateImportJobs, ListImportJobs, GetImportJobProgress Signed-off-by: Ruichen Bao --- CMakeLists.txt | 2 + cmake/ThirdPartyPackages.cmake | 40 +++++++++++++ src/impl/BulkImport.cpp | 87 +++++++++++++++++++++++++++++ src/include/milvus/BulkImport.h | 27 +++++++++ src/include/milvus/MilvusClientV2.h | 2 + 5 files changed, 158 insertions(+) create mode 100644 src/impl/BulkImport.cpp create mode 100644 src/include/milvus/BulkImport.h diff --git a/CMakeLists.txt b/CMakeLists.txt index 2be4880..d0b3ffe 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -64,6 +64,8 @@ set_option_category("Thirdparty") define_option_string(MILVUS_WITH_GRPC "Using gRPC from" "module" "package" "module") define_option_string(MILVUS_WITH_ZLIB "Using Zlib from" "module" "package" "module") define_option_string(MILVUS_WITH_NLOHMANN_JSON "nlohmann json from" "module" "package" "module") +define_option_string(MILVUS_WITH_CPP_HTTPLIB "cpp-httplib from" "module" "package" "module") +define_option_string(MILVUS_WITH_BROTLI "Using Brotli from" "module" "package" "module") define_option_string(MILVUS_WITH_GTEST "Using GTest from" "module" "package" "module") diff --git a/cmake/ThirdPartyPackages.cmake b/cmake/ThirdPartyPackages.cmake index fefd8c8..f7074ec 100644 --- a/cmake/ThirdPartyPackages.cmake +++ b/cmake/ThirdPartyPackages.cmake @@ -21,6 +21,8 @@ include(FetchContent) set(GRPC_VERSION 1.49.1) set(NLOHMANN_JSON_VERSION 3.11.3) set(GOOGLETEST_VERSION 1.12.1) +set(CPP_HTTPLIB_VERSION 0.18.5) +set(BROTLI_VERSION 1.1.0) Set(FETCHCONTENT_QUIET FALSE) # grpc @@ -50,6 +52,23 @@ FetchContent_Declare( GIT_PROGRESS TRUE ) +# cpp-httplib +FetchContent_Declare( + cpp-httplib + GIT_REPOSITORY https://github.com/yhirose/cpp-httplib.git + GIT_TAG v${CPP_HTTPLIB_VERSION} + GIT_SHALLOW TRUE + GIT_PROGRESS TRUE +) + +# brotli +FetchContent_Declare( + brotli + GIT_REPOSITORY https://github.com/google/brotli.git + GIT_TAG v${BROTLI_VERSION} + GIT_SHALLOW TRUE + GIT_PROGRESS TRUE +) # grpc if ("${MILVUS_WITH_GRPC}" STREQUAL "package") @@ -88,3 +107,24 @@ else () add_subdirectory(${nlohmann_json_SOURCE_DIR} ${nlohmann_json_BINARY_DIR} EXCLUDE_FROM_ALL) endif () endif () + +# cpp-httplib +if ("${MILVUS_WITH_CPP_HTTPLIB}" STREQUAL "package") + find_package(cpp-httplib REQUIRED) +else () + if (NOT cpp-httplib_POPULATED) + FetchContent_Populate(cpp-httplib) + add_subdirectory(${cpp-httplib_SOURCE_DIR} ${cpp-httplib_BINARY_DIR} EXCLUDE_FROM_ALL) + include_directories(${cpp-httplib_SOURCE_DIR}) + endif () +endif () + +# brotli +if ("${MILVUS_WITH_BROTLI}" STREQUAL "package") + find_package(Brotli REQUIRED) +else () + if (NOT brotli_POPULATED) + FetchContent_Populate(brotli) + add_subdirectory(${brotli_SOURCE_DIR} ${brotli_BINARY_DIR} EXCLUDE_FROM_ALL) + endif () +endif () diff --git a/src/impl/BulkImport.cpp b/src/impl/BulkImport.cpp new file mode 100644 index 0000000..041d927 --- /dev/null +++ b/src/impl/BulkImport.cpp @@ -0,0 +1,87 @@ +#include "milvus/BulkImport.h" + +namespace milvus { + +nlohmann::json +BulkImport::CreateImportJobs(const std::string& url, const std::string& collection_name, + const std::vector& files, const std::string& db_name, + const std::string& api_key, const std::string& partition_name, + const nlohmann::json& options) { + httplib::Client client(url); + + nlohmann::json request_payload = { + {"dbName", db_name}, + {"collectionName", collection_name}, + {"files", nlohmann::json::array({files})}, + }; + + if (!partition_name.empty()) { + request_payload["partitionName"] = partition_name; + } + + if (!options.empty()) { + if (options.contains("timeout") && !options["timeout"].is_null()) { + request_payload["options"] = options; + } + } + + std::string request_url = "/v2/vectordb/jobs/import/create"; + httplib::Headers headers = { + {"Authorization", "Bearer " + api_key}, + }; + std::string body = request_payload.dump(); + + auto res = client.Post(request_url, headers, body, "application/json"); + if (res && res->status == 200) { + return nlohmann::json::parse(res->body); + } else { + return nullptr; + } +} + +nlohmann::json +BulkImport::ListImportJobs(const std::string& url, const std::string& collection_name, const std::string& db_name, + const std::string& api_key) { + httplib::Client client(url); + + nlohmann::json request_payload = { + {"collectionName", collection_name}, + {"dbName", db_name}, + }; + + std::string request_url = "/v2/vectordb/jobs/import/list"; + httplib::Headers headers = { + {"Authorization", "Bearer " + api_key}, + }; + std::string body = request_payload.dump(); + + auto res = client.Post(request_url, headers, body, "application/json"); + if (res && res->status == 200) { + return nlohmann::json::parse(res->body); + } else { + return nullptr; + } +} + +nlohmann::json +BulkImport::GetImportJobProgress(const std::string& url, const std::string& job_id, const std::string& db_name, + const std::string& api_key) { + httplib::Client client(url); + + nlohmann::json payload = {{"dbName", db_name}, {"jobID", job_id}}; + + std::string request_url = "/v2/vectordb/jobs/import/get_progress"; + httplib::Headers headers = { + {"Authorization", "Bearer " + api_key}, + }; + std::string body = payload.dump(); + + auto res = client.Post(request_url, headers, body, "application/json"); + if (res && res->status == 200) { + return nlohmann::json::parse(res->body); + } else { + return nullptr; + } +} + +} // namespace milvus diff --git a/src/include/milvus/BulkImport.h b/src/include/milvus/BulkImport.h new file mode 100644 index 0000000..7d74821 --- /dev/null +++ b/src/include/milvus/BulkImport.h @@ -0,0 +1,27 @@ +#pragma once + +#include + +#include +#include +#include + +namespace milvus { + +class BulkImport { + public: + static nlohmann::json + CreateImportJobs(const std::string& url, const std::string& collection_name, const std::vector& files, + const std::string& db_name = "default", const std::string& api_key = "", + const std::string& partition_name = "", const nlohmann::json& options = nlohmann::json{}); + + static nlohmann::json + ListImportJobs(const std::string& url, const std::string& collection_name, const std::string& db_name = "default", + const std::string& api_key = ""); + + static nlohmann::json + GetImportJobProgress(const std::string& url, const std::string& job_id, const std::string& db_name = "default", + const std::string& api_key = ""); +}; + +} // namespace milvus diff --git a/src/include/milvus/MilvusClientV2.h b/src/include/milvus/MilvusClientV2.h index 9962426..5ed200a 100644 --- a/src/include/milvus/MilvusClientV2.h +++ b/src/include/milvus/MilvusClientV2.h @@ -16,6 +16,8 @@ #pragma once +#include + #include #include "Status.h"