-
Notifications
You must be signed in to change notification settings - Fork 3.4k
/
server.go
405 lines (363 loc) · 16.1 KB
/
server.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
// 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 flight
import (
"context"
"fmt"
"net"
"os"
"os/signal"
"github.com/apache/arrow/go/v16/arrow/flight/gen/flight"
"google.golang.org/grpc"
)
type (
FlightServer = flight.FlightServiceServer
FlightService_HandshakeServer = flight.FlightService_HandshakeServer
HandshakeResponse = flight.HandshakeResponse
HandshakeRequest = flight.HandshakeRequest
FlightService_ListFlightsServer = flight.FlightService_ListFlightsServer
FlightService_DoGetServer = flight.FlightService_DoGetServer
FlightService_DoPutServer = flight.FlightService_DoPutServer
FlightService_DoExchangeServer = flight.FlightService_DoExchangeServer
FlightService_DoActionServer = flight.FlightService_DoActionServer
FlightService_ListActionsServer = flight.FlightService_ListActionsServer
Criteria = flight.Criteria
FlightDescriptor = flight.FlightDescriptor
FlightEndpoint = flight.FlightEndpoint
Location = flight.Location
FlightInfo = flight.FlightInfo
PollInfo = flight.PollInfo
FlightData = flight.FlightData
PutResult = flight.PutResult
Ticket = flight.Ticket
SchemaResult = flight.SchemaResult
Action = flight.Action
ActionType = flight.ActionType
CancelFlightInfoRequest = flight.CancelFlightInfoRequest
RenewFlightEndpointRequest = flight.RenewFlightEndpointRequest
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
)
// NewSessionOptionValues returns a map with the same keys as the input map, but with all values converted
// to SessionOptionValues. If any values fail conversion, an error will be returned.
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
}
// NewSessionOptionValue takes any value and constructs a SessionOptionValue suitable for setting session values.
// An error will be returned if the value is not one of the types supported by SessionOptionValue.
func NewSessionOptionValue(value any) (flight.SessionOptionValue, error) {
if value == nil {
return flight.SessionOptionValue{}, nil
}
switch val := value.(type) {
case string:
return flight.SessionOptionValue{OptionValue: &flight.SessionOptionValue_StringValue{StringValue: val}}, nil
case bool:
return flight.SessionOptionValue{OptionValue: &flight.SessionOptionValue_BoolValue{BoolValue: val}}, nil
case int64:
return flight.SessionOptionValue{OptionValue: &flight.SessionOptionValue_Int64Value{Int64Value: val}}, nil
case float64:
return flight.SessionOptionValue{OptionValue: &flight.SessionOptionValue_DoubleValue{DoubleValue: val}}, nil
case []string:
return flight.SessionOptionValue{OptionValue: &flight.SessionOptionValue_StringListValue_{StringListValue: &flight.SessionOptionValue_StringListValue{Values: val}}}, nil
default:
return flight.SessionOptionValue{}, fmt.Errorf("invalid option type %[1]T for value %[1]v", val)
}
}
// Constants for CancelStatus
const (
// The cancellation 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.
CancelStatusUnspecified = flight.CancelStatus_CANCEL_STATUS_UNSPECIFIED
// The cancellation request is complete. Subsequent requests
// with the same payload may return CancelStatusCancelled or a
// arrow.ErrNotFound error.
CancelStatusCancelled = flight.CancelStatus_CANCEL_STATUS_CANCELLED
// The cancellation request is in progress. The client may
// retry the cancellation request.
CancelStatusCancelling = flight.CancelStatus_CANCEL_STATUS_CANCELLING
// The FlightInfo is not cancellable. The client should not
// retry the cancellation request.
CancelStatusNotCancellable = flight.CancelStatus_CANCEL_STATUS_NOT_CANCELLABLE
)
// Constants for Location
const (
// LocationReuseConnection is a special location that tells clients
// they may fetch the data from the same service that they obtained
// the FlightEndpoint response from.
LocationReuseConnection = "arrow-flight-reuse-connection://?"
)
// RegisterFlightServiceServer registers an existing flight server onto an
// existing grpc server, or anything that is a grpc service registrar.
func RegisterFlightServiceServer(s *grpc.Server, srv FlightServer) {
flight.RegisterFlightServiceServer(s, srv)
}
// From https://github.com/grpc/grpc-go/blob/4c776ec01572d55249df309251900554b46adb41/reflection/serverreflection.go#L69-L83
// This interface is inlined to make this arrow library compatible with
// grpc < 1.45 .
// See "google.golang.org/grpc/reflection" 's reflection.ServiceInfoProvider .
// serviceInfoProvider is an interface used to retrieve metadata about the
// services to expose.
//
// The reflection service is only interested in the service names, but the
// signature is this way so that *grpc.Server implements it. So it is okay
// for a custom implementation to return zero values for the
// grpc.ServiceInfo values in the map.
//
// # Experimental
//
// Notice: This type is EXPERIMENTAL and may be changed or removed in a
// later release.
type serviceInfoProvider interface {
GetServiceInfo() map[string]grpc.ServiceInfo
}
// Server is an interface for hiding some of the grpc specifics to make
// it slightly easier to manage a flight service, slightly modeled after
// the C++ implementation
type Server interface {
// Init takes in the address to bind to and creates the listener. If both this
// and InitListener are called, then whichever was called last will be used.
Init(addr string) error
// InitListener initializes with an already created listener rather than
// creating a new one like Init does. If both this and Init are called,
// whichever was called last is what will be used as they both set a listener
// into the server.
InitListener(lis net.Listener)
// Addr will return the address that was bound to for the service to listen on
Addr() net.Addr
// SetShutdownOnSignals sets notifications on the given signals to call GracefulStop
// on the grpc service if any of those signals are received
SetShutdownOnSignals(sig ...os.Signal)
// Serve blocks until accepting a connection fails with a fatal error. It will return
// a non-nil error unless it stopped due to calling Shutdown or receiving one of the
// signals set in SetShutdownOnSignals
Serve() error
// Shutdown will call GracefulStop on the grpc server so that it stops accepting connections
// and will wait until current methods complete
Shutdown()
// RegisterFlightService sets up the handler for the Flight Endpoints as per
// normal Grpc setups
RegisterFlightService(FlightServer)
// ServiceRegistrar wraps a single method that supports service registration.
// For example, it may be used to register health check provided by grpc-go.
grpc.ServiceRegistrar
// serviceInfoProvider is an interface used to retrieve metadata about the services to expose.
// If reflection is enabled on the server, all the endpoints can be invoked using grpcurl.
serviceInfoProvider
}
// BaseFlightServer is the base flight server implementation and must be
// embedded in any server implementation to ensure forward compatibility
// with any modifications of the spec without compiler errors.
type BaseFlightServer struct {
flight.UnimplementedFlightServiceServer
authHandler ServerAuthHandler
}
func (s *BaseFlightServer) GetAuthHandler() ServerAuthHandler { return s.authHandler }
func (s *BaseFlightServer) SetAuthHandler(handler ServerAuthHandler) {
s.authHandler = handler
}
func (s *BaseFlightServer) Handshake(stream flight.FlightService_HandshakeServer) error {
if s.authHandler == nil {
return nil
}
return s.authHandler.Authenticate(&serverAuthConn{stream})
}
// CustomerServerMiddleware is a helper interface for more easily defining custom
// grpc middleware without having to expose or understand all the grpc bells and whistles.
type CustomServerMiddleware interface {
// StartCall will be called with the current context of the call, grpc.SetHeader can be used to add outgoing headers
// if the returned context is non-nil, then it will be used as the new context being passed through the calls
StartCall(ctx context.Context) context.Context
// CallCompleted is a callback which is called with the return from the handler
// it will be nil if everything was successful or will be the error about to be returned
// to grpc
CallCompleted(ctx context.Context, err error)
}
// CreateServerMiddlware constructs a ServerMiddleware object for the passed in custom
// middleware, generating both the Unary and Stream interceptors from the interface.
func CreateServerMiddleware(middleware CustomServerMiddleware) ServerMiddleware {
return ServerMiddleware{
Unary: func(ctx context.Context, req interface{}, _ *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (ret interface{}, err error) {
nctx := middleware.StartCall(ctx)
if nctx != nil {
ctx = nctx
}
ret, err = handler(ctx, req)
middleware.CallCompleted(ctx, err)
return
},
Stream: func(srv interface{}, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
ctx := middleware.StartCall(stream.Context())
if ctx != nil {
stream = &wrappedStream{ServerStream: stream, ctx: ctx}
}
err := handler(srv, stream)
middleware.CallCompleted(stream.Context(), err)
return err
},
}
}
type ServerMiddleware struct {
Stream grpc.StreamServerInterceptor
Unary grpc.UnaryServerInterceptor
}
type server struct {
lis net.Listener
sigChannel <-chan os.Signal
done chan bool
server *grpc.Server
}
// NewFlightServer takes any grpc Server options desired, such as TLS certs and so
// on which will just be passed through to the underlying grpc server.
//
// Alternatively, a grpc server can be created normally without this helper as the
// grpc server generated code is still being exported. This only exists to allow
// the utility of the helpers
//
// Deprecated: prefer to use NewServerWithMiddleware, due to auth handler middleware
// this function will be problematic if any of the grpc options specify other middleware.
func NewFlightServer(opt ...grpc.ServerOption) Server {
opt = append([]grpc.ServerOption{
grpc.ChainStreamInterceptor(serverAuthStreamInterceptor),
grpc.ChainUnaryInterceptor(serverAuthUnaryInterceptor),
}, opt...)
return &server{
server: grpc.NewServer(opt...),
}
}
// NewServerWithMiddleware takes a slice of middleware which will be used
// by grpc and chained, the first middleware will be the outer most with the last
// middleware being the inner most wrapper around the actual call. It also takes
// any grpc Server options desired, such as TLS certs and so on which will just
// be passed through to the underlying grpc server.
//
// Because of the usage of `ChainStreamInterceptor` and `ChainUnaryInterceptor` do
// not specify any middleware using the grpc options, use the ServerMiddleware slice
// instead as the auth middleware will be added for handling the case that a service
// handler is registered that uses the ServerAuthHandler.
//
// Alternatively, a grpc server can be created normally without this helper as the
// grpc server generated code is still being exported. This only exists to allow
// the utility of the helpers.
func NewServerWithMiddleware(middleware []ServerMiddleware, opts ...grpc.ServerOption) Server {
unary := make([]grpc.UnaryServerInterceptor, 1, len(middleware)+1)
unary[0] = serverAuthUnaryInterceptor
stream := make([]grpc.StreamServerInterceptor, 1, len(middleware)+1)
stream[0] = serverAuthStreamInterceptor
if len(middleware) > 0 {
for _, m := range middleware {
if m.Unary != nil {
unary = append(unary, m.Unary)
}
if m.Stream != nil {
stream = append(stream, m.Stream)
}
}
}
opts = append(opts, grpc.ChainUnaryInterceptor(unary...), grpc.ChainStreamInterceptor(stream...))
return &server{server: grpc.NewServer(opts...)}
}
func (s *server) Init(addr string) (err error) {
s.lis, err = net.Listen("tcp", addr)
return
}
func (s *server) InitListener(lis net.Listener) {
s.lis = lis
}
func (s *server) Addr() net.Addr {
return s.lis.Addr()
}
func (s *server) SetShutdownOnSignals(sig ...os.Signal) {
c := make(chan os.Signal, 1)
signal.Notify(c, sig...)
s.sigChannel = c
}
func (s *server) Serve() error {
s.done = make(chan bool)
go func() {
select {
case <-s.sigChannel:
s.server.GracefulStop()
case <-s.done:
}
}()
err := s.server.Serve(s.lis)
close(s.done)
return err
}
func (s *server) RegisterFlightService(svc FlightServer) {
flight.RegisterFlightServiceServer(s.server, svc)
}
func (s *server) Shutdown() {
s.server.GracefulStop()
}
func (s *server) RegisterService(sd *grpc.ServiceDesc, ss interface{}) {
s.server.RegisterService(sd, ss)
}
func (s *server) GetServiceInfo() map[string]grpc.ServiceInfo {
return s.server.GetServiceInfo()
}