flight: Instantiate BaseType<> with unique_ptr for FlightInfo and Pol…

…lInfo
apache · Jul 15, 2024 · 47dabff · 47dabff
1 parent 1a292d6
commit 47dabff
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 6 deletions.
diff --git a/cpp/src/arrow/flight/serialization_internal.cc b/cpp/src/arrow/flight/serialization_internal.cc
@@ -268,6 +268,13 @@ Status FromProto(const pb::FlightInfo& pb_info, FlightInfo::Data* info) {
   return Status::OK();
 }
 
+Status FromProto(const pb::FlightInfo& pb_info, std::unique_ptr<FlightInfo>* info) {
+  FlightInfo::Data info_data;
+  RETURN_NOT_OK(FromProto(pb_info, &info_data));
+  *info = std::make_unique<FlightInfo>(std::move(info_data));
+  return Status::OK();
+}
+
 Status FromProto(const pb::BasicAuth& pb_basic_auth, BasicAuth* basic_auth) {
   basic_auth->password = pb_basic_auth.password();
   basic_auth->username = pb_basic_auth.username();
@@ -340,6 +347,13 @@ Status FromProto(const pb::PollInfo& pb_info, PollInfo* info) {
   return Status::OK();
 }
 
+Status FromProto(const pb::PollInfo& pb_info, std::unique_ptr<PollInfo>* info) {
+  PollInfo poll_info;
+  RETURN_NOT_OK(FromProto(pb_info, &poll_info));
+  *info = std::make_unique<PollInfo>(std::move(poll_info));
+  return Status::OK();
+}
+
 Status ToProto(const PollInfo& info, pb::PollInfo* pb_info) {
   if (info.info) {
     RETURN_NOT_OK(ToProto(*info.info, pb_info->mutable_info()));

diff --git a/cpp/src/arrow/flight/serialization_internal.h b/cpp/src/arrow/flight/serialization_internal.h
@@ -61,7 +61,9 @@ Status FromProto(const pb::FlightEndpoint& pb_endpoint, FlightEndpoint* endpoint
 Status FromProto(const pb::RenewFlightEndpointRequest& pb_request,
                  RenewFlightEndpointRequest* request);
 Status FromProto(const pb::FlightInfo& pb_info, FlightInfo::Data* info);
+Status FromProto(const pb::FlightInfo& pb_info, std::unique_ptr<FlightInfo>* info);
 Status FromProto(const pb::PollInfo& pb_info, PollInfo* info);
+Status FromProto(const pb::PollInfo& pb_info, std::unique_ptr<PollInfo>* info);
 Status FromProto(const pb::CancelFlightInfoRequest& pb_request,
                  CancelFlightInfoRequest* request);
 Status FromProto(const pb::SchemaResult& pb_result, std::string* result);

diff --git a/cpp/src/arrow/flight/types.h b/cpp/src/arrow/flight/types.h
@@ -162,18 +162,34 @@ struct ARROW_FLIGHT_EXPORT CertKeyPair {
 
 namespace internal {
 
+template <typename T>
+struct remove_unique_ptr {
+  using type = T;
+};
+
+template <typename T>
+struct remove_unique_ptr<std::unique_ptr<T>> {
+  using type = T;
+};
+
 // Base CRTP type
 template <class T>
 struct BaseType {
  protected:
-  const T& self() const { return static_cast<const T&>(*this); }
-  T& self() { return static_cast<T&>(*this); }
+  using SelfT = typename remove_unique_ptr<T>::type;
+
+  const SelfT& self() const { return static_cast<const SelfT&>(*this); }
+  SelfT& self() { return static_cast<SelfT&>(*this); }
 
  public:
   BaseType() = default;
 
-  friend bool operator==(const T& left, const T& right) { return left.Equals(right); }
-  friend bool operator!=(const T& left, const T& right) { return !left.Equals(right); }
+  friend bool operator==(const SelfT& left, const SelfT& right) {
+    return left.Equals(right);
+  }
+  friend bool operator!=(const SelfT& left, const SelfT& right) {
+    return !left.Equals(right);
+  }
 
   /// \brief Serialize this message to its wire-format representation.
   inline arrow::Result<std::string> SerializeToString() const {
@@ -589,7 +605,8 @@ struct ARROW_FLIGHT_EXPORT SchemaResult : public internal::BaseType<SchemaResult
 
 /// \brief The access coordinates for retrieval of a dataset, returned by
 /// GetFlightInfo
-class ARROW_FLIGHT_EXPORT FlightInfo : public internal::BaseType<FlightInfo> {
+class ARROW_FLIGHT_EXPORT FlightInfo
+    : public internal::BaseType<std::unique_ptr<FlightInfo>> {
  public:
   struct Data {
     std::string schema;
@@ -668,7 +685,8 @@ class ARROW_FLIGHT_EXPORT FlightInfo : public internal::BaseType<FlightInfo> {
 };
 
 /// \brief The information to process a long-running query.
-class ARROW_FLIGHT_EXPORT PollInfo : public internal::BaseType<PollInfo> {
+class ARROW_FLIGHT_EXPORT PollInfo
+    : public internal::BaseType<std::unique_ptr<PollInfo>> {
  public:
   /// The currently available results so far.
   std::unique_ptr<FlightInfo> info = NULLPTR;