go/kratos
1
0
Fork 0
mirror of https://github.com/go-kratos/kratos.git synced 2025-01-26 03:52:12 +02:00
Code Issues Releases Activity

feat: add grpc client stream interceptor opts (#2610)

Browse Source 
Co-authored-by: yao.li <yao.li@cloudminds.com>
This commit is contained in:
liyaopinner 2023-01-16 10:53:26 +08:00 committed by GitHub
parent bebea0c103
commit 35800916dc
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 11 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
transport/grpc/client.go
View File

@ -82,6 +82,13 @@ func WithUnaryInterceptor(in ...grpc.UnaryClientInterceptor) ClientOption {
}
}
// WithStreamInterceptor returns a DialOption that specifies the interceptor for streaming RPCs.
func WithStreamInterceptor(in ...grpc.StreamClientInterceptor) ClientOption {
return func(o *clientOptions) {
o.streamInts = in
}
}
// WithOptions with gRPC options.
func WithOptions(opts ...grpc.DialOption) ClientOption {
return func(o *clientOptions) {
@ -111,6 +118,7 @@ type clientOptions struct {
discovery registry.Discovery
middleware []middleware.Middleware
ints []grpc.UnaryClientInterceptor
streamInts []grpc.StreamClientInterceptor
grpcOpts []grpc.DialOption
balancerName string
filters []selector.NodeFilter
@ -145,6 +153,9 @@ func dial(ctx context.Context, insecure bool, opts ...ClientOption) (*grpc.Clien
grpc.WithDefaultServiceConfig(fmt.Sprintf(`{"loadBalancingConfig": [{"%s":{}}]}`, options.balancerName)),
grpc.WithChainUnaryInterceptor(ints...),
}
if len(options.streamInts) > 0 {
grpcOpts = append(grpcOpts, grpc.WithChainStreamInterceptor(options.streamInts...))
}
if options.discovery != nil {
grpcOpts = append(grpcOpts,
grpc.WithResolvers(