Add rest client part

TileDB-Inc · Nov 24, 2023 · af6c27b · af6c27b
1 parent 6698bc2
commit af6c27b
Show file tree

Hide file tree

Showing 4 changed files with 80 additions and 1 deletion.
diff --git a/tiledb/sm/array/array.h b/tiledb/sm/array/array.h
@@ -621,6 +621,11 @@ class Array {
   /** Load array directory for non-remote arrays */
   ArrayDirectory& load_array_directory();
 
+  /* Get the REST client */
+  [[nodiscard]] inline shared_ptr<RestClient> rest_client() const {
+    return resources_.rest_client();
+  }
+
  private:
   /* ********************************* */
   /*         PRIVATE ATTRIBUTES        */

diff --git a/tiledb/sm/consolidation_plan/consolidation_plan.cc b/tiledb/sm/consolidation_plan/consolidation_plan.cc
@@ -33,6 +33,7 @@
 #include "tiledb/sm/consolidation_plan/consolidation_plan.h"
 #include "tiledb/common/common.h"
 #include "tiledb/common/logger.h"
+#include "tiledb/sm/rest/rest_client.h"
 
 using namespace tiledb::sm;
 using namespace tiledb::common;
@@ -44,7 +45,20 @@ using namespace tiledb::common;
 ConsolidationPlan::ConsolidationPlan(
     shared_ptr<Array> array, uint64_t fragment_size)
     : desired_fragment_size_(fragment_size) {
-  generate(array);
+  if (array->is_remote()) {
+    auto rest_client = array->rest_client();
+    if (!rest_client) {
+      throw std::runtime_error(
+          "Failed to create a consolidation plan; Remote array"
+          "with no REST client.");
+    }
+    // reach out to the REST client to populate class members
+    fragment_uris_per_node_ = rest_client->post_consolidation_plan_from_rest(
+        array->array_uri(), array->config(), fragment_size);
+    num_nodes_ = fragment_uris_per_node_.size();
+  } else {
+    generate(array);
+  }
 }
 
 ConsolidationPlan::ConsolidationPlan(

diff --git a/tiledb/sm/rest/rest_client.cc b/tiledb/sm/rest/rest_client.cc
@@ -1516,6 +1516,48 @@ Status RestClient::post_vacuum_to_rest(const URI& uri, const Config& config) {
       stats_, url, serialization_type_, &serialized, &returned_data, cache_key);
 }
 
+std::vector<std::vector<std::string>>
+RestClient::post_consolidation_plan_from_rest(
+    const URI& uri, const Config& config, uint64_t fragment_size) {
+  Buffer buff;
+  serialization::serialize_consolidation_plan_request(
+      fragment_size, config, serialization_type_, buff);
+
+  // Wrap in a list
+  BufferList serialized;
+  throw_if_not_ok(serialized.add_buffer(std::move(buff)));
+
+  // Init curl and form the URL
+  Curl curlc(logger_);
+  std::string array_ns, array_uri;
+  throw_if_not_ok(uri.get_rest_components(&array_ns, &array_uri));
+  const std::string cache_key = array_ns + ":" + array_uri;
+  throw_if_not_ok(
+      curlc.init(config_, extra_headers_, &redirect_meta_, &redirect_mtx_));
+  const std::string url = redirect_uri(cache_key) + "/v1/arrays/" + array_ns +
+                          "/" + curlc.url_escape(array_uri) +
+                          "/consolidate/plan";
+
+  // Get the data
+  Buffer returned_data;
+  throw_if_not_ok(curlc.post_data(
+      stats_,
+      url,
+      serialization_type_,
+      &serialized,
+      &returned_data,
+      cache_key));
+  if (returned_data.data() == nullptr || returned_data.size() == 0) {
+    throw Status_RestError(
+        "Error getting query plan from REST; server returned no data.");
+  }
+
+  // Ensure data has a null delimiter for cap'n proto if using JSON
+  throw_if_not_ok(ensure_json_null_delimited_string(&returned_data));
+  return serialization::deserialize_consolidation_plan_response(
+      serialization_type_, returned_data);
+}
+
 #else
 
 RestClient::RestClient() {
@@ -1694,6 +1736,13 @@ Status RestClient::post_vacuum_to_rest(const URI&, const Config&) {
       Status_RestError("Cannot use rest client; serialization not enabled."));
 }
 
+std::vector<std::vector<std::string>>
+RestClient::post_consolidation_plan_from_rest(
+    const URI&, const Config&, uint64_t) {
+  throw StatusException(
+      Status_RestError("Cannot use rest client; serialization not enabled."));
+}
+
 #endif  // TILEDB_SERIALIZATION
 
 }  // namespace sm

diff --git a/tiledb/sm/rest/rest_client.h b/tiledb/sm/rest/rest_client.h
@@ -382,6 +382,17 @@ class RestClient {
    */
   Status post_vacuum_to_rest(const URI& uri, const Config& config);
 
+  /**
+   * Get consolidation plan from the REST server via POST request.
+   *
+   * @param uri Array URI.
+   * @param config Config of the array.
+   * @param fragment_size Maximum fragment size for constructing the plan.
+   * @return The requested consolidation plan
+   */
+  std::vector<std::vector<std::string>> post_consolidation_plan_from_rest(
+      const URI& uri, const Config& config, uint64_t fragment_size);
+
  private:
   /* ********************************* */
   /*        PRIVATE ATTRIBUTES         */