-
Notifications
You must be signed in to change notification settings - Fork 1
/
grpc.go
145 lines (119 loc) · 3.72 KB
/
grpc.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
package zapi
import (
"context"
"crypto/tls"
"io"
"github.com/golang/protobuf/ptypes/empty"
"golang.org/x/oauth2"
"google.golang.org/grpc"
_ "google.golang.org/grpc/balancer/grpclb" // register the grpclb balancer
"google.golang.org/grpc/credentials"
"google.golang.org/grpc/credentials/oauth"
"google.golang.org/grpc/metadata"
"zvelo.io/go-zapi/internal/zvelo"
msg "zvelo.io/msg/msgpb"
)
// A GRPCv1Client implements msg.APIv1Client as well as an io.Closer that, if
// closed, will close the underlying gRPC connection.
type GRPCv1Client interface {
msg.APIv1Client
io.Closer
}
type grpcV1Client struct {
options *options
client msg.APIv1Client
io.Closer
}
// A GRPCv1Dialer is used to simplify connecting to zveloAPI with the correct
// options. grpc DialOptions will override the defaults.
type GRPCv1Dialer interface {
Dial(context.Context, ...grpc.DialOption) (GRPCv1Client, error)
}
type grpcV1Dialer struct {
options *options
}
func (d grpcV1Dialer) Dial(ctx context.Context, opts ...grpc.DialOption) (GRPCv1Client, error) {
dialOpts := append([]grpc.DialOption{}, opts...)
if d.options.withoutTLS {
dialOpts = append(dialOpts, grpc.WithInsecure())
} else {
// #nosec
dialOpts = append(dialOpts, grpc.WithTransportCredentials(credentials.NewTLS(&tls.Config{
InsecureSkipVerify: d.options.tlsInsecureSkipVerify,
})))
}
if d.options.TokenSource != nil {
dialOpts = append(dialOpts,
grpc.WithPerRPCCredentials(oauth.TokenSource{
TokenSource: d.options,
}),
)
}
conn, err := grpc.DialContext(ctx, d.options.grpcTarget, dialOpts...)
if err != nil {
return nil, err
}
return grpcV1Client{
Closer: conn,
client: msg.NewAPIv1Client(conn),
options: d.options,
}, nil
}
// NewGRPCv1 returns a properly configured GRPCv1Dialer
func NewGRPCv1(ts oauth2.TokenSource, opts ...Option) GRPCv1Dialer {
o := defaults(ts)
for _, opt := range opts {
opt(o)
}
return grpcV1Dialer{options: o}
}
func (c grpcV1Client) Query(ctx context.Context, in *msg.QueryRequests, opts ...grpc.CallOption) (*msg.QueryReplies, error) {
zvelo.DebugContextOut(ctx, c.options.debug)
header, trailer, opts := grpcMD(opts...)
resp, err := c.client.Query(ctx, in, opts...)
zvelo.DebugMD(c.options.debug, *header, *trailer)
return resp, err
}
func (c grpcV1Client) Result(ctx context.Context, in *msg.RequestID, opts ...grpc.CallOption) (*msg.QueryResult, error) {
zvelo.DebugContextOut(ctx, c.options.debug)
header, trailer, opts := grpcMD(opts...)
resp, err := c.client.Result(ctx, in, opts...)
zvelo.DebugMD(c.options.debug, *header, *trailer)
return resp, err
}
func (c grpcV1Client) Suggest(ctx context.Context, in *msg.Suggestion, opts ...grpc.CallOption) (*empty.Empty, error) {
zvelo.DebugContextOut(ctx, c.options.debug)
header, trailer, opts := grpcMD(opts...)
resp, err := c.client.Suggest(ctx, in, opts...)
zvelo.DebugMD(c.options.debug, *header, *trailer)
return resp, err
}
func (c grpcV1Client) Stream(ctx context.Context, in *empty.Empty, opts ...grpc.CallOption) (msg.APIv1_StreamClient, error) {
if in == nil {
in = &empty.Empty{}
}
zvelo.DebugContextOut(ctx, c.options.debug)
header, trailer, opts := grpcMD(opts...)
resp, err := c.client.Stream(ctx, in, opts...)
zvelo.DebugMD(c.options.debug, *header, *trailer)
return resp, err
}
func grpcMD(in ...grpc.CallOption) (header, trailer *metadata.MD, opts []grpc.CallOption) {
for _, o := range in {
if m, ok := o.(grpc.HeaderCallOption); ok {
header = m.HeaderAddr
}
if m, ok := o.(grpc.TrailerCallOption); ok {
trailer = m.TrailerAddr
}
}
if header == nil {
header = &metadata.MD{}
opts = append(opts, grpc.Header(header))
}
if trailer == nil {
trailer = &metadata.MD{}
opts = append(opts, grpc.Trailer(trailer))
}
return
}