/
client.go
104 lines (85 loc) · 2.57 KB
/
client.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
// Package client is generic grpc client-side.
package client
import (
"context"
"google.golang.org/grpc"
"google.golang.org/grpc/credentials"
"google.golang.org/grpc/credentials/insecure"
"google.golang.org/grpc/resolver"
)
// Option client option func
type Option func(*options)
type options struct {
builders []resolver.Builder
isLoadBalance bool
credentials credentials.TransportCredentials
unaryInterceptors []grpc.UnaryClientInterceptor
streamInterceptors []grpc.StreamClientInterceptor
}
func defaultOptions() *options {
return &options{}
}
func (o *options) apply(opts ...Option) {
for _, opt := range opts {
opt(o)
}
}
// WithServiceDiscover set service discover
func WithServiceDiscover(builders ...resolver.Builder) Option {
return func(o *options) {
o.builders = builders
}
}
// WithLoadBalance set load balance
func WithLoadBalance() Option {
return func(o *options) {
o.isLoadBalance = true
}
}
// WithSecure set secure
func WithSecure(credential credentials.TransportCredentials) Option {
return func(o *options) {
o.credentials = credential
}
}
// WithUnaryInterceptor set unary interceptor
func WithUnaryInterceptor(interceptors ...grpc.UnaryClientInterceptor) Option {
return func(o *options) {
o.unaryInterceptors = interceptors
}
}
// WithStreamInterceptor set stream interceptor
func WithStreamInterceptor(interceptors ...grpc.StreamClientInterceptor) Option {
return func(o *options) {
o.streamInterceptors = interceptors
}
}
// Dial to grpc server
func Dial(ctx context.Context, endpoint string, opts ...Option) (*grpc.ClientConn, error) {
o := defaultOptions()
o.apply(opts...)
var dialOptions []grpc.DialOption
// service discovery
if len(o.builders) > 0 {
dialOptions = append(dialOptions, grpc.WithResolvers(o.builders...))
}
// load balance option
if o.isLoadBalance {
dialOptions = append(dialOptions, grpc.WithDefaultServiceConfig(`{"loadBalancingConfig": [{"round_robin":{}}]}`))
}
// secure option
if o.credentials == nil {
dialOptions = append(dialOptions, grpc.WithTransportCredentials(insecure.NewCredentials()))
} else {
dialOptions = append(dialOptions, grpc.WithTransportCredentials(o.credentials))
}
// custom unary interceptor option
if len(o.unaryInterceptors) > 0 {
dialOptions = append(dialOptions, grpc.WithChainUnaryInterceptor(o.unaryInterceptors...))
}
// custom stream interceptor option
if len(o.streamInterceptors) > 0 {
dialOptions = append(dialOptions, grpc.WithChainStreamInterceptor(o.streamInterceptors...))
}
return grpc.DialContext(ctx, endpoint, dialOptions...)
}