1717
1818#include < mutex>
1919
20+ #include " arrow/flight/server.h"
2021#include " arrow/flight/sql/server_session_middleware.h"
2122#include " arrow/flight/sql/server_session_middleware_factory.h"
2223
@@ -134,12 +135,12 @@ ServerSessionMiddlewareFactory::ParseCookieString(const std::string_view& s) {
134135}
135136
136137Status ServerSessionMiddlewareFactory::StartCall (
137- const CallInfo&, const CallHeaders& incoming_headers ,
138+ const CallInfo&, const ServerCallContext& context ,
138139 std::shared_ptr<ServerMiddleware>* middleware) {
139140 std::string session_id;
140141
141142 const std::pair<CallHeaders::const_iterator, CallHeaders::const_iterator>&
142- headers_it_pr = incoming_headers.equal_range (" cookie" );
143+ headers_it_pr = context. incoming_headers () .equal_range (" cookie" );
143144 for (auto itr = headers_it_pr.first ; itr != headers_it_pr.second ; ++itr) {
144145 const std::string_view& cookie_header = itr->second ;
145146 const std::vector<std::pair<std::string, std::string>> cookies =
@@ -158,16 +159,16 @@ Status ServerSessionMiddlewareFactory::StartCall(
158159 // No cookie was found
159160 // Temporary workaround until middleware handling fixed
160161 auto [id, s] = CreateNewSession ();
161- *middleware = std::make_shared<ServerSessionMiddlewareImpl>(this , incoming_headers,
162- std::move (s), id, false );
162+ *middleware = std::make_shared<ServerSessionMiddlewareImpl>(
163+ this , context. incoming_headers (), std::move (s), id, false );
163164 } else {
164165 const std::shared_lock<std::shared_mutex> l (session_store_lock_);
165166 if (auto it = session_store_.find (session_id); it == session_store_.end ()) {
166167 return Status::Invalid (" Invalid or expired " , kSessionCookieName , " cookie." );
167168 } else {
168169 auto session = it->second ;
169170 *middleware = std::make_shared<ServerSessionMiddlewareImpl>(
170- this , incoming_headers, std::move (session), session_id);
171+ this , context. incoming_headers () , std::move (session), session_id);
171172 }
172173 }
173174
0 commit comments