adds gRPC ClientAfter handler and multiple calls to Client Before/After funcs
Bas van Beek
7 years ago
21 | 21 | dec DecodeResponseFunc |
22 | 22 | grpcReply reflect.Type |
23 | 23 | before []RequestFunc |
24 | after []ClientResponseFunc | |
24 | 25 | } |
25 | 26 | |
26 | 27 | // NewClient constructs a usable Client for a single remote endpoint. |
53 | 54 | ).Interface(), |
54 | 55 | ), |
55 | 56 | before: []RequestFunc{}, |
57 | after: []ClientResponseFunc{}, | |
56 | 58 | } |
57 | 59 | for _, option := range options { |
58 | 60 | option(c) |
66 | 68 | // ClientBefore sets the RequestFuncs that are applied to the outgoing gRPC |
67 | 69 | // request before it's invoked. |
68 | 70 | func ClientBefore(before ...RequestFunc) ClientOption { |
69 | return func(c *Client) { c.before = before } | |
71 | return func(c *Client) { c.before = append(c.before, before...) } | |
72 | } | |
73 | ||
74 | // ClientAfter sets the ClientResponseFuncs that are applied to the incoming | |
75 | // gRPC response prior to it being decoded. This is useful for obtaining | |
76 | // response metadata and adding onto the context prior to decoding. | |
77 | func ClientAfter(after ...ClientResponseFunc) ClientOption { | |
78 | return func(c *Client) { c.after = append(c.after, after...) } | |
70 | 79 | } |
71 | 80 | |
72 | 81 | // Endpoint returns a usable endpoint that will invoke the gRPC specified by the |
87 | 96 | } |
88 | 97 | ctx = metadata.NewContext(ctx, *md) |
89 | 98 | |
99 | var header, trailer metadata.MD | |
90 | 100 | grpcReply := reflect.New(c.grpcReply).Interface() |
91 | if err = grpc.Invoke(ctx, c.method, req, grpcReply, c.client); err != nil { | |
101 | if err = grpc.Invoke( | |
102 | ctx, c.method, req, grpcReply, c.client, | |
103 | grpc.Header(&header), grpc.Trailer(&trailer), | |
104 | ); err != nil { | |
92 | 105 | return nil, err |
106 | } | |
107 | ||
108 | for _, f := range c.after { | |
109 | ctx = f(ctx, &header, &trailer) | |
93 | 110 | } |
94 | 111 | |
95 | 112 | response, err := c.dec(ctx, grpcReply) |
11 | 11 | binHdrSuffix = "-bin" |
12 | 12 | ) |
13 | 13 | |
14 | // RequestFunc may take information from an gRPC request and put it into a | |
15 | // request context. In Servers, BeforeFuncs are executed prior to invoking the | |
16 | // endpoint. In Clients, BeforeFuncs are executed after creating the request | |
14 | // RequestFunc may take information from a gRPC request and put it into a | |
15 | // request context. In Servers, RequestFuncs are executed prior to invoking the | |
16 | // endpoint. In Clients, RequestFuncs are executed after creating the request | |
17 | 17 | // but prior to invoking the gRPC client. |
18 | 18 | type RequestFunc func(context.Context, *metadata.MD) context.Context |
19 | 19 | |
20 | // ResponseFunc may take information from a request context and use it to | |
20 | // ServerResponseFunc may take information from a request context and use it to | |
21 | 21 | // manipulate the gRPC metadata header. ResponseFuncs are only executed in |
22 | 22 | // servers, after invoking the endpoint but prior to writing a response. |
23 | type ResponseFunc func(context.Context, *metadata.MD) | |
23 | type ServerResponseFunc func(context.Context, *metadata.MD) | |
24 | ||
25 | // ClientResponseFunc may take information from a gRPC metadata header and/or | |
26 | // trailer and make the responses available for consumption. ClientResponseFuncs | |
27 | // are only executed in clients, after a request has been made, but prior to it | |
28 | // being decoded. | |
29 | type ClientResponseFunc func(ctx context.Context, header *metadata.MD, trailer *metadata.MD) context.Context | |
24 | 30 | |
25 | 31 | // SetResponseHeader returns a ResponseFunc that sets the specified metadata |
26 | 32 | // key-value pair. |
27 | func SetResponseHeader(key, val string) ResponseFunc { | |
33 | func SetResponseHeader(key, val string) ServerResponseFunc { | |
28 | 34 | return func(_ context.Context, md *metadata.MD) { |
29 | 35 | key, val := EncodeKeyValue(key, val) |
30 | 36 | (*md)[key] = append((*md)[key], val) |
23 | 23 | dec DecodeRequestFunc |
24 | 24 | enc EncodeResponseFunc |
25 | 25 | before []RequestFunc |
26 | after []ResponseFunc | |
26 | after []ServerResponseFunc | |
27 | 27 | logger log.Logger |
28 | 28 | } |
29 | 29 | |
63 | 63 | |
64 | 64 | // ServerAfter functions are executed on the HTTP response writer after the |
65 | 65 | // endpoint is invoked, but before anything is written to the client. |
66 | func ServerAfter(after ...ResponseFunc) ServerOption { | |
66 | func ServerAfter(after ...ServerResponseFunc) ServerOption { | |
67 | 67 | return func(s *Server) { s.after = append(s.after, after...) } |
68 | 68 | } |
69 | 69 |
61 | 61 | // ClientBefore sets the RequestFuncs that are applied to the outgoing HTTP |
62 | 62 | // request before it's invoked. |
63 | 63 | func ClientBefore(before ...RequestFunc) ClientOption { |
64 | return func(c *Client) { c.before = before } | |
64 | return func(c *Client) { c.before = append(c.before, before...) } | |
65 | 65 | } |
66 | 66 | |
67 | 67 | // ClientAfter sets the ClientResponseFuncs applied to the incoming HTTP |
68 | 68 | // request prior to it being decoded. This is useful for obtaining anything off |
69 | 69 | // of the response and adding onto the context prior to decoding. |
70 | 70 | func ClientAfter(after ...ClientResponseFunc) ClientOption { |
71 | return func(c *Client) { c.after = after } | |
71 | return func(c *Client) { c.after = append(c.after, after...) } | |
72 | 72 | } |
73 | 73 | |
74 | 74 | // BufferedStream sets whether the Response.Body is left open, allowing it |