lidavidm commented on a change in pull request #7174:
URL: https://github.com/apache/arrow/pull/7174#discussion_r435502932



##########
File path: cpp/src/arrow/flight/client.cc
##########
@@ -148,7 +142,28 @@ class GrpcClientInterceptorAdapter : public 
grpc::experimental::Interceptor {
   }
 
  private:
+  void ReceivedHeaders(std::multimap<grpc::string_ref, grpc::string_ref>* 
metadata) {
+    if (received_headers_) {

Review comment:
       Yes - I did it that way as we don't normally expose trailers in other 
contexts (e.g. there's no way to send one), so I decided it would be best to 
not mix the two except when we're forced to by grpc-java.
   
   (Or rather, we do expose them with custom error metadata, but that's in a 
limited context.)

##########
File path: cpp/src/arrow/flight/test_integration.cc
##########
@@ -103,10 +103,142 @@ class AuthBasicProtoScenario : public Scenario {
   }
 };
 
+class TestServerMiddleware : public ServerMiddleware {
+ public:
+  explicit TestServerMiddleware(std::string received) : received_(received) {}
+  void SendingHeaders(AddCallHeaders* outgoing_headers) override {
+    outgoing_headers->AddHeader("x-middleware", received_);
+  }
+  void CallCompleted(const Status& status) override {}
+
+  std::string name() const override { return "GrpcTrailersMiddleware"; }
+
+ private:
+  std::string received_;
+};
+
+class TestServerMiddlewareFactory : public ServerMiddlewareFactory {
+ public:
+  Status StartCall(const CallInfo& info, const CallHeaders& incoming_headers,
+                   std::shared_ptr<ServerMiddleware>* middleware) override {
+    const std::pair<CallHeaders::const_iterator, CallHeaders::const_iterator>& 
iter_pair =
+        incoming_headers.equal_range("x-middleware");
+    std::string received = "";
+    if (iter_pair.first != iter_pair.second) {
+      const util::string_view& value = (*iter_pair.first).second;
+      received = std::string(value);
+    }
+    *middleware = std::make_shared<TestServerMiddleware>(received);
+    return Status::OK();
+  }
+};
+
+class TestClientMiddleware : public ClientMiddleware {
+ public:
+  explicit TestClientMiddleware(std::string* received_header)
+      : received_header_(received_header) {}
+
+  void SendingHeaders(AddCallHeaders* outgoing_headers) {
+    outgoing_headers->AddHeader("x-middleware", "expected value");
+  }
+
+  void ReceivedHeaders(const CallHeaders& incoming_headers) {
+    const std::pair<CallHeaders::const_iterator, CallHeaders::const_iterator>& 
iter_pair =
+        incoming_headers.equal_range("x-middleware");
+    if (iter_pair.first != iter_pair.second) {
+      const util::string_view& value = (*iter_pair.first).second;
+      *received_header_ = std::string(value);
+    }
+  }
+
+  void CallCompleted(const Status& status) {}
+
+ private:
+  std::string* received_header_;
+};
+
+class TestClientMiddlewareFactory : public ClientMiddlewareFactory {
+ public:
+  void StartCall(const CallInfo& info, std::unique_ptr<ClientMiddleware>* 
middleware) {
+    *middleware =
+        std::unique_ptr<ClientMiddleware>(new 
TestClientMiddleware(&received_header_));
+  }
+
+  std::string received_header_;
+};
+
+class MiddlewareServer : public FlightServerBase {
+  Status GetFlightInfo(const ServerCallContext& context,
+                       const FlightDescriptor& descriptor,
+                       std::unique_ptr<FlightInfo>* result) override {
+    if (descriptor.type == FlightDescriptor::DescriptorType::CMD &&
+        descriptor.cmd == "success") {
+      // Don't fail
+      std::shared_ptr<Schema> schema = arrow::schema({});
+      Location location;
+      RETURN_NOT_OK(Location::ForGrpcTcp("localhost", 10010, &location));

Review comment:
       Yes, I clarified with a comment.




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
[email protected]


Reply via email to