add grpc interceptor option (#990)

pull/991/head
Tony Chen 4 years ago committed by GitHub
parent 29e387a39c
commit 92b6bf54eb
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
  1. 16
      transport/grpc/client.go
  2. 3
      transport/grpc/server.go
  3. 2
      transport/http/client.go

@ -48,6 +48,13 @@ func WithDiscovery(d registry.Discovery) ClientOption {
} }
} }
// WithUnaryInterceptor returns a DialOption that specifies the interceptor for unary RPCs.
func WithUnaryInterceptor(in ...grpc.UnaryClientInterceptor) ClientOption {
return func(o *clientOptions) {
o.ints = in
}
}
// WithOptions with gRPC options. // WithOptions with gRPC options.
func WithOptions(opts ...grpc.DialOption) ClientOption { func WithOptions(opts ...grpc.DialOption) ClientOption {
return func(o *clientOptions) { return func(o *clientOptions) {
@ -61,6 +68,7 @@ type clientOptions struct {
timeout time.Duration timeout time.Duration
middleware middleware.Middleware middleware middleware.Middleware
discovery registry.Discovery discovery registry.Discovery
ints []grpc.UnaryClientInterceptor
grpcOpts []grpc.DialOption grpcOpts []grpc.DialOption
} }
@ -84,9 +92,15 @@ func dial(ctx context.Context, insecure bool, opts ...ClientOption) (*grpc.Clien
for _, o := range opts { for _, o := range opts {
o(&options) o(&options)
} }
var ints = []grpc.UnaryClientInterceptor{
unaryClientInterceptor(options.middleware, options.timeout),
}
if len(options.ints) > 0 {
ints = append(ints, options.ints...)
}
var grpcOpts = []grpc.DialOption{ var grpcOpts = []grpc.DialOption{
grpc.WithBalancerName(roundrobin.Name), grpc.WithBalancerName(roundrobin.Name),
grpc.WithUnaryInterceptor(unaryClientInterceptor(options.middleware, options.timeout)), grpc.WithChainUnaryInterceptor(ints...),
} }
if options.discovery != nil { if options.discovery != nil {
grpcOpts = append(grpcOpts, grpc.WithResolvers(discovery.NewBuilder(options.discovery))) grpcOpts = append(grpcOpts, grpc.WithResolvers(discovery.NewBuilder(options.discovery)))

@ -61,8 +61,7 @@ func Middleware(m ...middleware.Middleware) ServerOption {
} }
} }
// UnaryInterceptor returns a ServerOption that sets the UnaryServerInterceptor for the // UnaryInterceptor returns a ServerOption that sets the UnaryServerInterceptor for the server.
// server.
func UnaryInterceptor(in ...grpc.UnaryServerInterceptor) ServerOption { func UnaryInterceptor(in ...grpc.UnaryServerInterceptor) ServerOption {
return func(s *Server) { return func(s *Server) {
s.ints = in s.ints = in

@ -145,7 +145,7 @@ func NewClient(ctx context.Context, opts ...ClientOption) (*Client, error) {
return nil, err return nil, err
} }
var r *resolver var r *resolver
if target.Endpoint != "" && options.discovery != nil { if options.discovery != nil {
if target.Scheme == "discovery" { if target.Scheme == "discovery" {
if r, err = newResolver(ctx, options.discovery, target); err != nil { if r, err = newResolver(ctx, options.discovery, target); err != nil {
return nil, fmt.Errorf("[http client] new resolver failed!err: %v", options.endpoint) return nil, fmt.Errorf("[http client] new resolver failed!err: %v", options.endpoint)

Loading…
Cancel
Save