lidavidm commented on code in PR #34817:
URL: https://github.com/apache/arrow/pull/34817#discussion_r1245575638
##########
format/FlightSql.proto:
##########
@@ -1842,6 +1842,94 @@ message ActionCancelQueryResult {
CancelResult result = 1;
}
+/*
+ * Request message for the "Close Session" action.
+ */
+message ActionCloseSessionRequest {
+ option (experimental) = true;
+}
+
+/*
+ * The result of closing a session.
+ *
+ * The result should be wrapped in a google.protobuf.Any message.
+ */
+message ActionCloseSessionResult {
+ option (experimental) = true;
+
+ enum CloseSessionResult {
+ // The session close status is unknown. Servers should avoid using
+ // this value (send a NOT_FOUND error if the requested query is
+ // not known). Clients can retry the request.
+ CLOSE_RESULT_UNSPECIFIED = 0;
+ // The session close request is complete. Subsequent requests with
+ // a NOT_FOUND error.
+ CLOSE_RESULT_CLOSED = 1;
+ // The session close request is in progress. The client may retry
+ // the close request.
+ CLOSE_RESULT_CLOSING = 2;
+ // The session is not closeable. The client should not retry the
+ // close request.
+ CLOSE_RESULT_NOT_CLOSEABLE = 3;
+ }
+
+ CloseSessionResult result = 1;
+}
+
+message SessionOption {
+ option (experimental) = true;
+
+ message StringListValue {
+ repeated string values = 1;
+ }
+
+ string option_name = 1;
+ oneof option_value {
+ string string_value = 2;
+ bool bool_value = 3;
+ sfixed32 int32_value = 4;
+ sfixed64 int64_value = 5;
+ float float_value = 6;
+ double double_value = 7;
Review Comment:
Do we need both 32 and 64 bit int, and 32 and 64 bit float?
##########
cpp/src/arrow/flight/sql/types.h:
##########
@@ -44,6 +44,23 @@ using SqlInfoResult =
/// \brief Map SQL info identifier to its value.
using SqlInfoResultMap = std::unordered_map<int32_t, SqlInfoResult>;
+/// \brief Variant supporting all possible types for SetSessionOptions
+using SessionOptionValue =
+ std::variant<std::string, bool, int32_t, int64_t, float, double,
std::vector<std::string>>;
+
+enum struct SessionOptionValueType : size_t {
+ kString, kBool, kInt32, kInt64, kFloat, kDouble, kStringList
+};
+
+struct ARROW_FLIGHT_SQL_EXPORT SessionOption {
+ std::string option_name;
+ SessionOptionValue option_value;
+
+ explicit SessionOption(std::string name, SessionOptionValue val)
+ : option_name{ std::move(name) }, option_value{ std::move(val) } {}
+ SessionOption() {}
Review Comment:
Just `SessionOption() = default;`
##########
format/FlightSql.proto:
##########
@@ -1842,6 +1842,94 @@ message ActionCancelQueryResult {
CancelResult result = 1;
}
+/*
+ * Request message for the "Close Session" action.
+ */
+message ActionCloseSessionRequest {
+ option (experimental) = true;
+}
+
+/*
+ * The result of closing a session.
+ *
+ * The result should be wrapped in a google.protobuf.Any message.
+ */
+message ActionCloseSessionResult {
+ option (experimental) = true;
+
+ enum CloseSessionResult {
+ // The session close status is unknown. Servers should avoid using
+ // this value (send a NOT_FOUND error if the requested query is
+ // not known). Clients can retry the request.
+ CLOSE_RESULT_UNSPECIFIED = 0;
+ // The session close request is complete. Subsequent requests with
+ // a NOT_FOUND error.
+ CLOSE_RESULT_CLOSED = 1;
+ // The session close request is in progress. The client may retry
+ // the close request.
+ CLOSE_RESULT_CLOSING = 2;
+ // The session is not closeable. The client should not retry the
+ // close request.
+ CLOSE_RESULT_NOT_CLOSEABLE = 3;
+ }
+
+ CloseSessionResult result = 1;
+}
+
+message SessionOption {
+ option (experimental) = true;
+
+ message StringListValue {
+ repeated string values = 1;
+ }
+
+ string option_name = 1;
+ oneof option_value {
+ string string_value = 2;
+ bool bool_value = 3;
+ sfixed32 int32_value = 4;
+ sfixed64 int64_value = 5;
+ float float_value = 6;
+ double double_value = 7;
+ StringListValue string_list_value = 8;
+ }
+}
+
+message ActionSetSessionOptionsRequest {
+ option (experimental) = true;
+
+ repeated SessionOption session_options = 1;
+}
+
+message ActionSetSessionOptionsResult {
+ option (experimental) = true;
+
+ enum SetSessionOptionResult {
+ // The status of setting the option is unknown. Servers should avoid using
+ // this value (send a NOT_FOUND error if the requested query is
+ // not known). Clients can retry the request.
+ SET_SESSION_OPTION_RESULT_UNSPECIFIED = 0;
+ // The session option setting completed successfully.
+ SET_SESSION_OPTION_RESULT_OK = 1;
+ // The session cannot be set to the given value.
+ SET_SESSION_OPTION_RESULT_INVALID_VALUE = 2;
+ // The session cannot be set.
+ SET_SESSION_OPTION_RESULT_ERROR = 3;
+ }
+
+ repeated SetSessionOptionResult results = 1;
Review Comment:
This is one result per option?
##########
cpp/src/arrow/flight/sql/server_session_middleware.h:
##########
@@ -0,0 +1,88 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+// Middleware for handling Flight SQL Sessions including session cookie
handling.
+// Currently experimental.
+
+#pragma once
+
+#include "arrow/flight/server_middleware.h"
+#include "arrow/flight/sql/types.h"
+
+namespace arrow {
+namespace flight {
+namespace sql {
+
+class ServerSessionMiddlewareFactory;
+
+static constexpr char const kSessionCookieName[] =
+ "flight_sql_session_id";
+
+class FlightSqlSession {
+ protected:
+ std::map<std::string, SessionOptionValue> map_;
Review Comment:
Does order matter? Can we use unordered_map?
##########
cpp/src/arrow/flight/sql/client.cc:
##########
@@ -802,6 +802,157 @@ ::arrow::Result<CancelResult>
FlightSqlClient::CancelQuery(
return Status::IOError("Server returned unknown result ", result.result());
}
+::arrow::Result<std::vector<SetSessionOptionResult>>
FlightSqlClient::SetSessionOptions(
+ const FlightCallOptions& options,
+ const std::vector<SessionOption>& session_options) {
+ flight_sql_pb::ActionSetSessionOptionsRequest request;
+ for (const SessionOption& in_opt : session_options) {
+ flight_sql_pb::SessionOption* opt = request.add_session_options();
+ const std::string& name = in_opt.option_name;
+ opt->set_option_name(name);
+
+ const SessionOptionValue& value = in_opt.option_value;
+ if (value.index() == std::variant_npos)
+ return Status::Invalid("Undefined SessionOptionValue type ");
+ switch (static_cast<SessionOptionValueType>(value.index())) {
Review Comment:
I think the way you're supposed to do this is `std::visit`, and then we
don't need the extra enum.
##########
format/FlightSql.proto:
##########
Review Comment:
FWIW, for other more recent proposals for features like this we've decided
to make them usable from base Arrow Flight as well. I think it's also
reasonable to do the same here. (There wouldn't be handlers added on the base
Flight server - you would have to look for and handle the action yourself - but
there were handlers added to the Flight SQL server.)
##########
format/FlightSql.proto:
##########
@@ -1842,6 +1842,94 @@ message ActionCancelQueryResult {
CancelResult result = 1;
}
+/*
+ * Request message for the "Close Session" action.
+ */
+message ActionCloseSessionRequest {
Review Comment:
To circle back here, expectations should be clearly documented in the
Protobuf. So here, it should be explicitly stated that the client and server
are using a shared header to store the session state, and that the header
should be the HTTP Cookie/Set-Cookie header.
##########
cpp/src/arrow/flight/sql/server_session_middleware.h:
##########
@@ -0,0 +1,88 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+// Middleware for handling Flight SQL Sessions including session cookie
handling.
+// Currently experimental.
+
+#pragma once
+
+#include "arrow/flight/server_middleware.h"
+#include "arrow/flight/sql/types.h"
+
+namespace arrow {
+namespace flight {
+namespace sql {
+
+class ServerSessionMiddlewareFactory;
+
+static constexpr char const kSessionCookieName[] =
+ "flight_sql_session_id";
+
+class FlightSqlSession {
+ protected:
Review Comment:
Why protected?
##########
cpp/src/arrow/flight/sql/server_session_middleware.h:
##########
@@ -0,0 +1,88 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+// Middleware for handling Flight SQL Sessions including session cookie
handling.
+// Currently experimental.
+
+#pragma once
+
+#include "arrow/flight/server_middleware.h"
+#include "arrow/flight/sql/types.h"
+
+namespace arrow {
+namespace flight {
+namespace sql {
+
+class ServerSessionMiddlewareFactory;
+
+static constexpr char const kSessionCookieName[] =
+ "flight_sql_session_id";
+
+class FlightSqlSession {
+ protected:
+ std::map<std::string, SessionOptionValue> map_;
+ std::shared_mutex map_lock_;
+ public:
+ /// \brief Get session option by key
+ SessionOptionValue GetSessionOption(const std::string&);
+ /// \brief Set session option by key to given value
+ void SetSessionOption(const std::string&, const SessionOptionValue&);
Review Comment:
nit: prefer string_view over const string&?
##########
cpp/src/arrow/flight/sql/server_session_middleware.cc:
##########
@@ -0,0 +1,179 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <shared_mutex>
+#include "arrow/flight/sql/server_session_middleware.h"
+#include <boost/lexical_cast.hpp>
+#include <boost/uuid/uuid.hpp>
+#include <boost/uuid/uuid_generators.hpp>
+#include <boost/uuid/uuid_io.hpp>
+
+namespace arrow {
+namespace flight {
+namespace sql {
+
+/// \brief A factory for ServerSessionMiddleware, itself storing session data.
+class ServerSessionMiddlewareFactory : public ServerMiddlewareFactory {
+ protected:
+ std::map<std::string, std::shared_ptr<FlightSqlSession>> session_store_;
+ std::shared_mutex session_store_lock_;
+ boost::uuids::random_generator uuid_generator_;
+
+ std::vector<std::pair<std::string, std::string>> ParseCookieString(
Review Comment:
Can we do this without duplicating the base cookie middleware logic?
If that means that we don't ship the middleware, but only implement it in
the unit/integration tests, that might be reasonable, so long as it's clear how
servers/clients are expected to implement this (what headers to send/receive,
etc.)
--
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]
