zeroshade commented on code in PR #40284:
URL: https://github.com/apache/arrow/pull/40284#discussion_r1508106790
##########
go/arrow/flight/client.go:
##########
@@ -421,3 +424,87 @@ func (c *client) RenewFlightEndpoint(ctx context.Context,
request *RenewFlightEn
}
return &renewedEndpoint, nil
}
+
+func (c *client) SetSessionOptions(ctx context.Context, request
*SetSessionOptionsRequest, opts ...grpc.CallOption) (*SetSessionOptionsResult,
error) {
+ var err error
+ var action flight.Action
+ action.Type = SetSessionOptionsActionType
+ action.Body, err = proto.Marshal(request)
+ if err != nil {
+ return nil, err
+ }
+ stream, err := c.DoAction(ctx, &action, opts...)
+ if err != nil {
+ return nil, err
+ }
+ res, err := stream.Recv()
+ if err != nil {
+ return nil, err
+ }
+ var result SetSessionOptionsResult
+ err = proto.Unmarshal(res.Body, &result)
+ if err != nil {
+ return nil, err
+ }
+ err = ReadUntilEOF(stream)
+ if err != nil {
+ return nil, err
+ }
+ return &result, nil
+}
+
+func (c *client) GetSessionOptions(ctx context.Context, request
*GetSessionOptionsRequest, opts ...grpc.CallOption) (*GetSessionOptionsResult,
error) {
+ var err error
+ var action flight.Action
+ action.Type = GetSessionOptionsActionType
+ action.Body, err = proto.Marshal(request)
+ if err != nil {
+ return nil, err
+ }
+ stream, err := c.DoAction(ctx, &action, opts...)
+ if err != nil {
+ return nil, err
+ }
+ res, err := stream.Recv()
+ if err != nil {
+ return nil, err
+ }
+ var result GetSessionOptionsResult
+ err = proto.Unmarshal(res.Body, &result)
+ if err != nil {
+ return nil, err
+ }
+ err = ReadUntilEOF(stream)
+ if err != nil {
+ return nil, err
+ }
Review Comment:
could we replace these with a helper function that implements this using
generics since the only difference between these methods are the types?
##########
cpp/src/arrow/flight/sql/server_session_middleware.cc:
##########
@@ -80,7 +80,7 @@ class ServerSessionMiddlewareImpl : public
ServerSessionMiddleware {
Status CloseSession() override {
const std::lock_guard<std::shared_mutex> l(mutex_);
- if (static_cast<bool>(session_)) {
+ if (!static_cast<bool>(session_)) {
Review Comment:
was this a test that was failing that this fixes?
##########
go/arrow/flight/session/cookies.go:
##########
@@ -0,0 +1,80 @@
+// 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.
+
+package session
+
+import (
+ "context"
+ "fmt"
+ "net/http"
+
+ "google.golang.org/grpc/metadata"
+)
+
+func GetIncomingCookieByName(ctx context.Context, name string) (http.Cookie,
error) {
+ md, ok := metadata.FromIncomingContext(ctx)
+ if !ok {
+ return http.Cookie{}, fmt.Errorf("no metadata found for
incoming context")
+ }
Review Comment:
should this get tied in with the Cookie middleware somehow?
##########
go/arrow/flight/session/session.go:
##########
@@ -0,0 +1,233 @@
+// 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.
+
+// Package session provides server middleware and reference implementations
for Flight session management.
+//
+// The existing middleware implementation uses cookies, so any client would
need middleware/support for storing and sending those cookies.
+// Both stateful and stateless session cookie implementations are provided.
+// The default stateful implementation persists sessions in-memory, but a
custom SessionStore may be provided for more durable persistence.
+// Stateless session cookies may be used with no additional infrastructure,
but they are not encrypted and should not contain sensitive information.
+package session
+
+import (
+ "context"
+ "errors"
+ "fmt"
+ "net/http"
+ "sync"
+
+ "github.com/apache/arrow/go/v16/arrow/flight"
+ "google.golang.org/grpc"
+ "google.golang.org/grpc/metadata"
+ "google.golang.org/protobuf/proto"
+)
+
+var ErrNoSession error = errors.New("flight: server session not present")
+
+type sessionMiddlewareKey struct{}
+
+// Return a copy of the provided context containing the provided ServerSession
+func NewSessionContext(ctx context.Context, session ServerSession)
context.Context {
+ return context.WithValue(ctx, sessionMiddlewareKey{}, session)
+}
+
+// Retrieve the ServerSession from the provided context if it exists.
+// An error indicates that the session was not found in the context.
+func GetSessionFromContext(ctx context.Context) (ServerSession, error) {
+ session, ok := ctx.Value(sessionMiddlewareKey{}).(ServerSession)
+ if !ok {
+ return nil, ErrNoSession
+ }
+ return session, nil
+}
+
+// Container for named SessionOptionValues
+type ServerSession interface {
+ // An identifier for the session that the server can use to reconstruct
+ // the session state on future requests. It is the responsibility of
+ // each implementation to define the token's semantics.
+ Token() string
+ // Get session option value by name, or nil if it does not exist
+ GetSessionOption(name string) *flight.SessionOptionValue
+ // Get a copy of the session options
+ GetSessionOptions() map[string]*flight.SessionOptionValue
+ // Set session option by name to given value
+ SetSessionOption(name string, value *flight.SessionOptionValue)
+ // Idempotently remove name from this session
+ EraseSessionOption(name string)
+ // Close the session
+ Close() error
+ // Report whether the session has been closed
+ Closed() bool
+}
+
+// Handles session lifecycle management
+type ServerSessionManager interface {
+ // Create a new, empty ServerSession
+ CreateSession(ctx context.Context) (ServerSession, error)
+ // Get the current ServerSession, if one exists
+ GetSession(ctx context.Context) (ServerSession, error)
+ // Cleanup any resources associated with the current ServerSession
+ CloseSession(session ServerSession) error
+}
+
+// Implementation of common session behavior. Intended to be extended
+// by specific session implementations.
+type serverSession struct {
+ closed bool
+
+ options map[string]*flight.SessionOptionValue
+ mu sync.RWMutex
+}
+
+func (session *serverSession) GetSessionOption(name string)
*flight.SessionOptionValue {
+ session.mu.RLock()
+ defer session.mu.RUnlock()
+ value, found := session.options[name]
+ if !found {
+ return nil
+ }
+
+ return value
+}
+
+func (session *serverSession) GetSessionOptions()
map[string]*flight.SessionOptionValue {
+ options := make(map[string]*flight.SessionOptionValue,
len(session.options))
+
+ session.mu.RLock()
+ defer session.mu.RUnlock()
+ for k, v := range session.options {
+ options[k] = proto.Clone(v).(*flight.SessionOptionValue)
+ }
+
+ return options
+}
+
+func (session *serverSession) SetSessionOption(name string, value
*flight.SessionOptionValue) {
+ if value.GetOptionValue() == nil {
+ session.EraseSessionOption(name)
+ return
+ }
+
+ session.mu.Lock()
+ defer session.mu.Unlock()
+ session.options[name] = value
+}
+
+func (session *serverSession) EraseSessionOption(name string) {
+ session.mu.Lock()
+ defer session.mu.Unlock()
+ delete(session.options, name)
+}
+
+func (session *serverSession) Close() error {
+ session.options = nil
+ session.closed = true
+ return nil
+}
+
+func (session *serverSession) Closed() bool {
+ return session.closed
+}
+
+// Create new instance of CustomServerMiddleware implementing server session
persistence.
+//
+// The provided manager can be used to customize session
implementation/behavior.
+// If no manager is provided, a stateful in-memory, goroutine-safe
implementation is used.
+func NewServerSessionMiddleware(manager ServerSessionManager)
*serverSessionMiddleware {
+ // Default manager
Review Comment:
@lidavidm Was a default defined or specified for handling with FlightSQL
sessions?
##########
go/arrow/flight/client.go:
##########
@@ -421,3 +424,87 @@ func (c *client) RenewFlightEndpoint(ctx context.Context,
request *RenewFlightEn
}
return &renewedEndpoint, nil
}
+
+func (c *client) SetSessionOptions(ctx context.Context, request
*SetSessionOptionsRequest, opts ...grpc.CallOption) (*SetSessionOptionsResult,
error) {
+ var err error
+ var action flight.Action
+ action.Type = SetSessionOptionsActionType
+ action.Body, err = proto.Marshal(request)
+ if err != nil {
+ return nil, err
+ }
+ stream, err := c.DoAction(ctx, &action, opts...)
+ if err != nil {
+ return nil, err
+ }
+ res, err := stream.Recv()
+ if err != nil {
+ return nil, err
+ }
+ var result SetSessionOptionsResult
+ err = proto.Unmarshal(res.Body, &result)
+ if err != nil {
+ return nil, err
+ }
+ err = ReadUntilEOF(stream)
+ if err != nil {
+ return nil, err
+ }
+ return &result, nil
Review Comment:
Should we do `return &result, ReadUntilEOF(stream)` ?
##########
go/arrow/flight/session/session.go:
##########
@@ -0,0 +1,240 @@
+// 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.
+
+// Package session provides server middleware and reference implementations
for Flight session management.
+//
+// For more details on the Flight Session Specification, see:
+//
https://arrow.apache.org/docs/format/FlightSql.html#flight-server-session-management
+//
+// [NewServerSessionMiddleware] manages sessions using cookies, so any client
would need its own
+// middleware/support for storing and sending those cookies. The cookies may
be stateful or stateless:
+//
+// - [NewStatefulServerSessionManager] implements stateful cookies.
+//
+// - [NewStatelessServerSessionManager] implements stateless cookies.
+//
+// See details of either implementation for caveats and recommended usage
scenarios.
+package session
+
+import (
+ "context"
+ "errors"
+ "fmt"
+ "net/http"
+ "sync"
+
+ "github.com/apache/arrow/go/v16/arrow/flight"
+ "google.golang.org/grpc"
+ "google.golang.org/grpc/metadata"
+ "google.golang.org/protobuf/proto"
+)
+
+var ErrNoSession error = errors.New("flight: server session not present")
+
+type sessionMiddlewareKey struct{}
+
+// Return a copy of the provided context containing the provided ServerSession
+func NewSessionContext(ctx context.Context, session ServerSession)
context.Context {
Review Comment:
Conventions say it should start with the name of the function:
`NewSessionContext returns a copy of the provided context containing the
provided ServerSession`
##########
go/arrow/flight/server.go:
##########
@@ -54,15 +55,86 @@ type (
Result = flight.Result
CancelFlightInfoResult = flight.CancelFlightInfoResult
CancelStatus = flight.CancelStatus
+ SessionOptionValue = flight.SessionOptionValue
+ SetSessionOptionsRequest = flight.SetSessionOptionsRequest
+ SetSessionOptionsResult = flight.SetSessionOptionsResult
+ SetSessionOptionsResultError = flight.SetSessionOptionsResult_Error
+ GetSessionOptionsRequest = flight.GetSessionOptionsRequest
+ GetSessionOptionsResult = flight.GetSessionOptionsResult
+ CloseSessionRequest = flight.CloseSessionRequest
+ CloseSessionResult = flight.CloseSessionResult
Empty = flight.Empty
)
// Constants for Action types
const (
CancelFlightInfoActionType = "CancelFlightInfo"
RenewFlightEndpointActionType = "RenewFlightEndpoint"
+ SetSessionOptionsActionType = "SetSessionOptions"
+ GetSessionOptionsActionType = "GetSessionOptions"
+ CloseSessionActionType = "CloseSession"
)
+const (
+ // The set option error is unknown. Servers should avoid
+ // using this value (send a NOT_FOUND error if the requested
+ // FlightInfo is not known). Clients can retry the request.
+ SetSessionOptionsResultErrorUnspecified =
flight.SetSessionOptionsResult_UNSPECIFIED
+ // The given session option name is invalid.
+ SetSessionOptionsResultErrorInvalidName =
flight.SetSessionOptionsResult_INVALID_NAME
+ // The session option value or type is invalid.
+ SetSessionOptionsResultErrorInvalidValue =
flight.SetSessionOptionsResult_INVALID_VALUE
+ // The session option cannot be set.
+ SetSessionOptionsResultErrorError = flight.SetSessionOptionsResult_ERROR
+)
+
+const (
+ // The close session status is unknown. Servers should avoid
+ // using this value (send a NOT_FOUND error if the requested
+ // FlightInfo is not known). Clients can retry the request.
+ CloseSessionResultUnspecified = flight.CloseSessionResult_UNSPECIFIED
+ // The session close request is complete.
+ CloseSessionResultClosed = flight.CloseSessionResult_CLOSED
+ // The session close request is in progress. The client may retry the
request.
+ CloseSessionResultClosing = flight.CloseSessionResult_CLOSING
+ // The session is not closeable.
+ CloseSessionResultNotCloseable = flight.CloseSessionResult_NOT_CLOSEABLE
+)
+
+func NewSessionOptionValues(options map[string]any)
(map[string]*flight.SessionOptionValue, error) {
Review Comment:
godoc comment?
##########
go/arrow/flight/session/session.go:
##########
@@ -0,0 +1,240 @@
+// 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.
+
+// Package session provides server middleware and reference implementations
for Flight session management.
+//
+// For more details on the Flight Session Specification, see:
+//
https://arrow.apache.org/docs/format/FlightSql.html#flight-server-session-management
+//
+// [NewServerSessionMiddleware] manages sessions using cookies, so any client
would need its own
+// middleware/support for storing and sending those cookies. The cookies may
be stateful or stateless:
+//
+// - [NewStatefulServerSessionManager] implements stateful cookies.
+//
+// - [NewStatelessServerSessionManager] implements stateless cookies.
+//
+// See details of either implementation for caveats and recommended usage
scenarios.
+package session
+
+import (
+ "context"
+ "errors"
+ "fmt"
+ "net/http"
+ "sync"
+
+ "github.com/apache/arrow/go/v16/arrow/flight"
+ "google.golang.org/grpc"
+ "google.golang.org/grpc/metadata"
+ "google.golang.org/protobuf/proto"
+)
+
+var ErrNoSession error = errors.New("flight: server session not present")
+
+type sessionMiddlewareKey struct{}
+
+// Return a copy of the provided context containing the provided ServerSession
+func NewSessionContext(ctx context.Context, session ServerSession)
context.Context {
+ return context.WithValue(ctx, sessionMiddlewareKey{}, session)
+}
+
+// Retrieve the ServerSession from the provided context if it exists.
+// An error indicates that the session was not found in the context.
+func GetSessionFromContext(ctx context.Context) (ServerSession, error) {
+ session, ok := ctx.Value(sessionMiddlewareKey{}).(ServerSession)
+ if !ok {
+ return nil, ErrNoSession
+ }
+ return session, nil
+}
+
+// Container for named SessionOptionValues
Review Comment:
```suggestion
// ServerSession is a container for named SessionOptionValues
```
##########
go/arrow/flight/server.go:
##########
@@ -54,15 +55,86 @@ type (
Result = flight.Result
CancelFlightInfoResult = flight.CancelFlightInfoResult
CancelStatus = flight.CancelStatus
+ SessionOptionValue = flight.SessionOptionValue
+ SetSessionOptionsRequest = flight.SetSessionOptionsRequest
+ SetSessionOptionsResult = flight.SetSessionOptionsResult
+ SetSessionOptionsResultError = flight.SetSessionOptionsResult_Error
+ GetSessionOptionsRequest = flight.GetSessionOptionsRequest
+ GetSessionOptionsResult = flight.GetSessionOptionsResult
+ CloseSessionRequest = flight.CloseSessionRequest
+ CloseSessionResult = flight.CloseSessionResult
Empty = flight.Empty
)
// Constants for Action types
const (
CancelFlightInfoActionType = "CancelFlightInfo"
RenewFlightEndpointActionType = "RenewFlightEndpoint"
+ SetSessionOptionsActionType = "SetSessionOptions"
+ GetSessionOptionsActionType = "GetSessionOptions"
+ CloseSessionActionType = "CloseSession"
)
+const (
+ // The set option error is unknown. Servers should avoid
+ // using this value (send a NOT_FOUND error if the requested
+ // FlightInfo is not known). Clients can retry the request.
+ SetSessionOptionsResultErrorUnspecified =
flight.SetSessionOptionsResult_UNSPECIFIED
+ // The given session option name is invalid.
+ SetSessionOptionsResultErrorInvalidName =
flight.SetSessionOptionsResult_INVALID_NAME
+ // The session option value or type is invalid.
+ SetSessionOptionsResultErrorInvalidValue =
flight.SetSessionOptionsResult_INVALID_VALUE
+ // The session option cannot be set.
+ SetSessionOptionsResultErrorError = flight.SetSessionOptionsResult_ERROR
+)
+
+const (
+ // The close session status is unknown. Servers should avoid
+ // using this value (send a NOT_FOUND error if the requested
+ // FlightInfo is not known). Clients can retry the request.
+ CloseSessionResultUnspecified = flight.CloseSessionResult_UNSPECIFIED
+ // The session close request is complete.
+ CloseSessionResultClosed = flight.CloseSessionResult_CLOSED
+ // The session close request is in progress. The client may retry the
request.
+ CloseSessionResultClosing = flight.CloseSessionResult_CLOSING
+ // The session is not closeable.
+ CloseSessionResultNotCloseable = flight.CloseSessionResult_NOT_CLOSEABLE
+)
+
+func NewSessionOptionValues(options map[string]any)
(map[string]*flight.SessionOptionValue, error) {
+ sessionOptions := make(map[string]*flight.SessionOptionValue,
len(options))
+ for key, val := range options {
+ optval, err := NewSessionOptionValue(val)
+ if err != nil {
+ return nil, err
+ }
+ sessionOptions[key] = &optval
+ }
+
+ return sessionOptions, nil
+}
+
+func NewSessionOptionValue(value any) (flight.SessionOptionValue, error) {
Review Comment:
Godoc 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.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
