raulcd commented on code in PR #50407:
URL: https://github.com/apache/arrow/pull/50407#discussion_r4313224920
##########
cpp/src/arrow/flight/flight_test.cc:
##########
@@ -1668,12 +1668,14 @@ TEST_F(TestCancel, DoExchange) {
stop_source.RequestStop(Status::Cancelled("StopSource"));
ASSERT_OK_AND_ASSIGN(auto do_exchange_result,
client_->DoExchange(options,
FlightDescriptor::Command("")));
+ ASSERT_OK(do_exchange_result.writer->DoneWriting());
Review Comment:
Is there any reason for this test change?
##########
ci/docker/ubuntu-22.04-cpp.dockerfile:
##########
@@ -223,6 +223,8 @@ ENV absl_SOURCE=BUNDLED \
google_cloud_cpp_storage_SOURCE=BUNDLED \
opentelemetry_cpp_SOURCE=BUNDLED \
ORC_SOURCE=BUNDLED \
+ Protobuf_SOURCE=BUNDLED \
+ gRPC_SOURCE=BUNDLED \
Review Comment:
I don't think those changed are relevant for this PR anymore, there
shouldn't be a change for Protobuf bundled/grpc right?
##########
cpp/src/arrow/flight/transport/grpc/grpc_server.cc:
##########
@@ -416,7 +276,6 @@ class GrpcServiceHandler final : public
FlightService::Service {
std::unique_ptr<FlightInfo> info;
auto res = impl_->base()->GetFlightInfo(flight_context, descr, &info);
- addMiddlewareHeaders(context, flight_context);
Review Comment:
I am unsure whether this was intended behavior initially but removing this
is a change in behavior. Removing this would stop updating cookies that are
dependent on the handler being called. I did a minor test (claude did the test
as it was quicker), it passes on main but fails on the PR:
<details>
<summary>Test</summary>
```diff
diff --git a/cpp/src/arrow/flight/flight_test.cc
b/cpp/src/arrow/flight/flight_test.cc
index 352bd7c1df..7d2921ed4e 100644
--- a/cpp/src/arrow/flight/flight_test.cc
+++ b/cpp/src/arrow/flight/flight_test.cc
@@ -546,6 +546,75 @@ class CountingServerMiddlewareFactory : public
ServerMiddlewareFactory {
std::atomic<int> failed_;
};
+// A server middleware that only sends a header once the RPC handler has
+// flagged it, to check that headers are gathered after the handler runs.
+class HandlerStateServerMiddleware : public ServerMiddleware {
+ public:
+ void SendingHeaders(AddCallHeaders* outgoing_headers) override {
+ if (handler_ran) {
+ outgoing_headers->AddHeader("x-handler-ran", "true");
+ }
+ }
+ void CallCompleted(const Status& status) override {}
+
+ std::string name() const override { return
"HandlerStateServerMiddleware"; }
+
+ bool handler_ran = false;
+};
+
+class HandlerStateServerMiddlewareFactory : public ServerMiddlewareFactory {
+ public:
+ Status StartCall(const CallInfo& info, const ServerCallContext& context,
+ std::shared_ptr<ServerMiddleware>* middleware) override {
+ *middleware = std::make_shared<HandlerStateServerMiddleware>();
+ return Status::OK();
+ }
+};
+
+class HandlerStateTestServer : public FlightServerBase {
+ Status GetFlightInfo(const ServerCallContext& context, const
FlightDescriptor& request,
+ std::unique_ptr<FlightInfo>* info) override {
+ auto middleware =
+
static_cast<HandlerStateServerMiddleware*>(context.GetMiddleware("handler_state"));
+ if (middleware == nullptr) {
+ return Status::Invalid("Could not find middleware");
+ }
+ middleware->handler_ran = true;
+ ARROW_ASSIGN_OR_RAISE(auto flight_info,
+ FlightInfo::Make(*arrow::schema({}), request, {},
-1, -1));
+ *info = std::make_unique<FlightInfo>(std::move(flight_info));
+ return Status::OK();
+ }
+};
+
+// A client middleware that records the values of the x-handler-ran header.
+class HandlerStateClientMiddleware : public ClientMiddleware {
+ public:
+ explicit HandlerStateClientMiddleware(std::vector<std::string>* values)
+ : values_(values) {}
+ void SendingHeaders(AddCallHeaders* outgoing_headers) override {}
+ void ReceivedHeaders(const CallHeaders& incoming_headers) override {
+ auto [begin, end] = incoming_headers.equal_range("x-handler-ran");
+ for (auto it = begin; it != end; ++it) {
+ values_->emplace_back(it->second);
+ }
+ }
+ void CallCompleted(const Status& status) override {}
+
+ private:
+ std::vector<std::string>* values_;
+};
+
+class HandlerStateClientMiddlewareFactory : public ClientMiddlewareFactory {
+ public:
+ void StartCall(const CallInfo& info,
+ std::unique_ptr<ClientMiddleware>* middleware) override {
+ *middleware = std::make_unique<HandlerStateClientMiddleware>(&values_);
+ }
+
+ std::vector<std::string> values_;
+};
+
// The current span ID, used to emulate OpenTracing style distributed
// tracing. Only used for communication between application code and
// client middleware.
@@ -831,6 +900,34 @@ class TestCountingServerMiddleware : public
::testing::Test {
std::unique_ptr<FlightServerBase> server_;
};
+class TestHandlerStateServerMiddleware : public ::testing::Test {
+ public:
+ void SetUp() {
+ client_middleware_ =
std::make_shared<HandlerStateClientMiddlewareFactory>();
+ ASSERT_OK(MakeServer<HandlerStateTestServer>(
+ &server_, &client_,
+ [](FlightServerOptions* options) {
+ options->middleware.push_back(
+ {"handler_state",
std::make_shared<HandlerStateServerMiddlewareFactory>()});
+ return Status::OK();
+ },
+ [&](FlightClientOptions* options) {
+ options->middleware.push_back(client_middleware_);
+ return Status::OK();
+ }));
+ }
+
+ void TearDown() {
+ ASSERT_OK(client_->Close());
+ ASSERT_OK(server_->Shutdown());
+ }
+
+ protected:
+ std::shared_ptr<HandlerStateClientMiddlewareFactory> client_middleware_;
+ std::unique_ptr<FlightClient> client_;
+ std::unique_ptr<FlightServerBase> server_;
+};
+
// Setup for this test is 2 servers
// 1. Client makes request to server A with a request ID set
// 2. server A extracts the request ID and makes a request to server B
@@ -1428,6 +1525,13 @@ TEST_F(TestCountingServerMiddleware, Count) {
ASSERT_EQ(1, request_counter_->failed_);
}
+// GH-34865: headers for GetFlightInfo are also gathered after the handler
+// runs, so middleware can send headers depending on the handler's state.
+TEST_F(TestHandlerStateServerMiddleware, GetFlightInfo) {
+ ASSERT_OK(client_->GetFlightInfo(FlightDescriptor::Command("")));
+ ASSERT_THAT(client_middleware_->values_, ::testing::ElementsAre("true"));
+}
+
TEST_F(TestPropagatingMiddleware, Propagate) {
Action action;
std::unique_ptr<Result> result;
```
</details>
```suggestion
helper_.AddMiddlewareHeaders(context, &flight_context);
```
We probably want to add a new test to capture the behavior. @lidavidm I
suspect that's intended, right?
--
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.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
