DX-61034: Arrow changes from apacheGH-34865 PR

Relates to DX-64185
stevelorddremio · Sep 19, 2023 · 2a03f16 · 2a03f16
1 parent d3954fe
commit 2a03f16
Show file tree

Hide file tree

Showing 35 changed files with 829 additions and 26 deletions.
diff --git a/cpp/src/arrow/flight/sql/CMakeLists.txt b/cpp/src/arrow/flight/sql/CMakeLists.txt
@@ -41,7 +41,8 @@ set(ARROW_FLIGHT_SQL_SRCS
     sql_info_internal.cc
     column_metadata.cc
     client.cc
-    protocol_internal.cc)
+    protocol_internal.cc
+    server_session_middleware.cc)
 
 add_arrow_lib(arrow_flight_sql
               CMAKE_PACKAGE_NAME

diff --git a/cpp/src/arrow/flight/sql/client.cc b/cpp/src/arrow/flight/sql/client.cc
@@ -31,6 +31,13 @@
 #include "arrow/result.h"
 #include "arrow/util/logging.h"
 
+// Lambda helper & CTAD
+template<class... Ts>
+struct overloaded : Ts... { using Ts::operator()...; };
+template<class... Ts> // CTAD will not be needed for >=C++20
+overloaded(Ts...) -> overloaded<Ts...>;
+
+namespace pb = arrow::flight::protocol;
 namespace flight_sql_pb = arrow::flight::protocol::sql;
 
 namespace arrow {
@@ -802,6 +809,149 @@ ::arrow::Result<CancelResult> FlightSqlClient::CancelQuery(
   return Status::IOError("Server returned unknown result ", result.result());
 }
 
+::arrow::Result<std::map<std::string, SetSessionOptionResult>>
+FlightSqlClient::SetSessionOptions(
+    const FlightCallOptions& options,
+    const std::map<std::string, SessionOptionValue>& session_options) {
+  pb::ActionSetSessionOptionsRequest request;
+  auto* options_map = request.mutable_session_options();
+
+  for (const auto & [name, opt_value] : session_options) {
+    pb::SessionOptionValue pb_opt_value;
+
+    if (opt_value.index() == std::variant_npos)
+      return Status::Invalid("Undefined SessionOptionValue type ");
+
+    std::visit(overloaded{
+      // TODO move this somewhere common that can have Proto-involved code
+      [&](std::string v) { pb_opt_value.set_string_value(v); },
+      [&](bool v) { pb_opt_value.set_bool_value(v); },
+      [&](int32_t v) { pb_opt_value.set_int32_value(v); },
+      [&](int64_t v) { pb_opt_value.set_int64_value(v); },
+      [&](float v) { pb_opt_value.set_float_value(v); },
+      [&](double v) { pb_opt_value.set_double_value(v); },
+      [&](std::vector<std::string> v) {
+        auto* string_list_value = pb_opt_value.mutable_string_list_value();
+        for (const std::string& s : v)
+          string_list_value->add_values(s);
+      }
+    }, opt_value);
+    (*options_map)[name] = std::move(pb_opt_value);
+  }
+
+  std::unique_ptr<ResultStream> results;
+  ARROW_ASSIGN_OR_RAISE(auto action, PackAction("SetSessionOptions", request));
+  ARROW_RETURN_NOT_OK(DoAction(options, action, &results));
+
+  pb::ActionSetSessionOptionsResult pb_result;
+  ARROW_RETURN_NOT_OK(ReadResult(results.get(), &pb_result));
+  ARROW_RETURN_NOT_OK(DrainResultStream(results.get()));
+  std::map<std::string, SetSessionOptionResult> result;
+  for (const auto & [result_key, result_value] : pb_result.results()) {
+    switch (result_value) {
+      case pb::ActionSetSessionOptionsResult::SET_SESSION_OPTION_RESULT_UNSPECIFIED:
+        result[result_key] = SetSessionOptionResult::kUnspecified;
+        break;
+      case pb::ActionSetSessionOptionsResult
+          ::SET_SESSION_OPTION_RESULT_OK:
+        result[result_key] = SetSessionOptionResult::kOk;
+        break;
+      case pb::ActionSetSessionOptionsResult
+          ::SET_SESSION_OPTION_RESULT_INVALID_VALUE:
+        result[result_key] = SetSessionOptionResult::kInvalidResult;
+        break;
+      case pb::ActionSetSessionOptionsResult::SET_SESSION_OPTION_RESULT_ERROR:
+        result[result_key] = SetSessionOptionResult::kError;
+        break;
+      default:
+        return Status::IOError("Invalid SetSessionOptionResult value for key "
+                               + result_key);
+    }
+  }
+
+  return result;
+}
+
+::arrow::Result<std::map<std::string, SessionOptionValue>>
+FlightSqlClient::GetSessionOptions (
+    const FlightCallOptions& options) {
+  pb::ActionGetSessionOptionsRequest request;
+
+  std::unique_ptr<ResultStream> results;
+  ARROW_ASSIGN_OR_RAISE(auto action, PackAction("GetSessionOptions", request));
+  ARROW_RETURN_NOT_OK(DoAction(options, action, &results));
+
+  pb::ActionGetSessionOptionsResult pb_result;
+  ARROW_RETURN_NOT_OK(ReadResult(results.get(), &pb_result));
+  ARROW_RETURN_NOT_OK(DrainResultStream(results.get()));
+
+  std::map<std::string, SessionOptionValue> result;
+  if (pb_result.session_options_size() > 0) {
+    for (auto& [pb_opt_name, pb_opt_val] : pb_result.session_options()) {
+      SessionOptionValue val;
+      switch (pb_opt_val.option_value_case()) {
+        case pb::SessionOptionValue::OPTION_VALUE_NOT_SET:
+          return Status::Invalid("Unset option_value for name '" + pb_opt_name + "'");
+        case pb::SessionOptionValue::kStringValue:
+          val = pb_opt_val.string_value();
+          break;
+        case pb::SessionOptionValue::kBoolValue:
+          val = pb_opt_val.bool_value();
+          break;
+        case pb::SessionOptionValue::kInt32Value:
+          val = pb_opt_val.int32_value();
+          break;
+        case pb::SessionOptionValue::kInt64Value:
+          val = pb_opt_val.int64_value();
+          break;
+        case pb::SessionOptionValue::kFloatValue:
+          val = pb_opt_val.float_value();
+          break;
+        case pb::SessionOptionValue::kDoubleValue:
+          val = pb_opt_val.double_value();
+          break;
+        case pb::SessionOptionValue::kStringListValue:
+          val.emplace<std::vector<std::string>>();
+          std::get<std::vector<std::string>>(val)
+              .reserve(pb_opt_val.string_list_value().values_size());
+          for (const std::string& s : pb_opt_val.string_list_value().values())
+            std::get<std::vector<std::string>>(val).push_back(s);
+          break;
+      }
+      result[pb_opt_name] = std::move(val);
+    }
+  }
+
+  return result;
+}
+
+::arrow::Result<CloseSessionResult> FlightSqlClient::CloseSession(
+    const FlightCallOptions& options) {
+  pb::ActionCloseSessionRequest request;
+
+  std::unique_ptr<ResultStream> results;
+  ARROW_ASSIGN_OR_RAISE(auto action, PackAction("CloseSession", request));
+  ARROW_RETURN_NOT_OK(DoAction(options, action, &results));
+
+  pb::ActionCloseSessionResult result;
+  ARROW_RETURN_NOT_OK(ReadResult(results.get(), &result));
+  ARROW_RETURN_NOT_OK(DrainResultStream(results.get()));
+  switch (result.result()) {
+    case pb::ActionCloseSessionResult::CLOSE_RESULT_UNSPECIFIED:
+      return CloseSessionResult::kUnspecified;
+    case pb::ActionCloseSessionResult::CLOSE_RESULT_CLOSED:
+      return CloseSessionResult::kClosed;
+    case pb::ActionCloseSessionResult::CLOSE_RESULT_CLOSING:
+      return CloseSessionResult::kClosing;
+    case pb::ActionCloseSessionResult::CLOSE_RESULT_NOT_CLOSEABLE:
+      return CloseSessionResult::kNotClosable;
+    default:
+      break;
+  }
+
+  return Status::IOError("Server returned unknown result ", result.result());
+}
+
 Status FlightSqlClient::Close() { return impl_->Close(); }
 
 std::ostream& operator<<(std::ostream& os, CancelResult result) {

diff --git a/cpp/src/arrow/flight/sql/client.h b/cpp/src/arrow/flight/sql/client.h
@@ -19,6 +19,7 @@
 
 #include <cstdint>
 #include <memory>
+#include <map>
 #include <string>
 
 #include "arrow/flight/client.h"
@@ -329,6 +330,25 @@ class ARROW_FLIGHT_SQL_EXPORT FlightSqlClient {
   /// \param[in] info         The FlightInfo of the query to cancel.
   ::arrow::Result<CancelResult> CancelQuery(const FlightCallOptions& options,
                                             const FlightInfo& info);
+
+  /// \brief Sets session options.
+  ///
+  /// \param[in] options            RPC-layer hints for this call.
+  /// \param[in] session_options    The session options to set.
+  ::arrow::Result<std::map<std::string, SetSessionOptionResult>> SetSessionOptions(
+      const FlightCallOptions& options,
+      const std::map<std::string, SessionOptionValue>& session_options);
+
+  /// \brief Gets current session options.
+  ///
+  /// \param[in] options            RPC-layer hints for this call.
+  ::arrow::Result<std::map<std::string, SessionOptionValue>> GetSessionOptions(
+      const FlightCallOptions& options);
+
+  /// \brief Explicitly closes the session if applicable.
+  ///
+  /// \param[in] options      RPC-layer hints for this call.
+  ::arrow::Result<CloseSessionResult> CloseSession(const FlightCallOptions& options);
 
   /// \brief Explicitly shut down and clean up the client.
   Status Close();

diff --git a/cpp/src/arrow/flight/sql/protocol_internal.h b/cpp/src/arrow/flight/sql/protocol_internal.h
@@ -24,3 +24,4 @@
 #include "arrow/flight/sql/visibility.h"
 
 #include "arrow/flight/sql/FlightSql.pb.h"  // IWYU pragma: export
+#include "arrow/flight/Flight.pb.h"