astelm
/
grpcurl
mirror of https://github.com/fullstorydev/grpcurl.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

增加callOptions

This commit is contained in:
woniu317 2021-12-21 09:48:50 +08:00
parent 8ee6c9423b
commit 4f02cf92a9
1 changed files with 15 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

30
invoke.go
View File

@ -51,7 +51,7 @@ type RequestMessageSupplier func() ([]byte, error)
// //
// Deprecated: use InvokeRPC instead. // Deprecated: use InvokeRPC instead.
func InvokeRpc(ctx context.Context, source DescriptorSource, cc *grpc.ClientConn, methodName string, func InvokeRpc(ctx context.Context, source DescriptorSource, cc *grpc.ClientConn, methodName string,
headers []string, handler InvocationEventHandler, requestData RequestMessageSupplier) error { headers []string, handler InvocationEventHandler, requestData RequestMessageSupplier, opts ...grpc.CallOption) error {
return InvokeRPC(ctx, source, cc, methodName, headers, handler, func(m proto.Message) error { return InvokeRPC(ctx, source, cc, methodName, headers, handler, func(m proto.Message) error {
// New function is almost identical, but the request supplier function works differently. // New function is almost identical, but the request supplier function works differently.
@ -61,7 +61,7 @@ func InvokeRpc(ctx context.Context, source DescriptorSource, cc *grpc.ClientConn
return err return err
} }
return jsonpb.Unmarshal(bytes.NewReader(data), m) return jsonpb.Unmarshal(bytes.NewReader(data), m)
}) }, opts...)
} }
// RequestSupplier is a function that is called to populate messages for a gRPC operation. The // RequestSupplier is a function that is called to populate messages for a gRPC operation. The
@ -85,7 +85,7 @@ type RequestSupplier func(proto.Message) error
// than the one invoking event callbacks. (This only happens for bi-directional streaming RPCs, where // than the one invoking event callbacks. (This only happens for bi-directional streaming RPCs, where
// one goroutine sends request messages and another consumes the response messages). // one goroutine sends request messages and another consumes the response messages).
func InvokeRPC(ctx context.Context, source DescriptorSource, ch grpcdynamic.Channel, methodName string, func InvokeRPC(ctx context.Context, source DescriptorSource, ch grpcdynamic.Channel, methodName string,
headers []string, handler InvocationEventHandler, requestData RequestSupplier) error { headers []string, handler InvocationEventHandler, requestData RequestSupplier, opts ...grpc.CallOption) error {
md := MetadataFromHeaders(headers) md := MetadataFromHeaders(headers)
@ -140,18 +140,18 @@ func InvokeRPC(ctx context.Context, source DescriptorSource, ch grpcdynamic.Chan
defer cancel() defer cancel()
if mtd.IsClientStreaming() && mtd.IsServerStreaming() { if mtd.IsClientStreaming() && mtd.IsServerStreaming() {
return invokeBidi(ctx, stub, mtd, handler, requestData, req) return invokeBidi(ctx, stub, mtd, handler, requestData, req, opts...)
} else if mtd.IsClientStreaming() { } else if mtd.IsClientStreaming() {
return invokeClientStream(ctx, stub, mtd, handler, requestData, req) return invokeClientStream(ctx, stub, mtd, handler, requestData, req, opts...)
} else if mtd.IsServerStreaming() { } else if mtd.IsServerStreaming() {
return invokeServerStream(ctx, stub, mtd, handler, requestData, req) return invokeServerStream(ctx, stub, mtd, handler, requestData, req, opts...)
} else { } else {
return invokeUnary(ctx, stub, mtd, handler, requestData, req) return invokeUnary(ctx, stub, mtd, handler, requestData, req, opts...)
} }
} }
func invokeUnary(ctx context.Context, stub grpcdynamic.Stub, md *desc.MethodDescriptor, handler InvocationEventHandler, func invokeUnary(ctx context.Context, stub grpcdynamic.Stub, md *desc.MethodDescriptor, handler InvocationEventHandler,
requestData RequestSupplier, req proto.Message) error { requestData RequestSupplier, req proto.Message, opts ...grpc.CallOption) error {
err := requestData(req) err := requestData(req)
if err != nil && err != io.EOF { if err != nil && err != io.EOF {
@ -170,7 +170,7 @@ func invokeUnary(ctx context.Context, stub grpcdynamic.Stub, md *desc.MethodDesc
// Now we can actually invoke the RPC! // Now we can actually invoke the RPC!
var respHeaders metadata.MD var respHeaders metadata.MD
var respTrailers metadata.MD var respTrailers metadata.MD
resp, err := stub.InvokeRpc(ctx, md, req, grpc.Trailer(&respTrailers), grpc.Header(&respHeaders)) resp, err := stub.InvokeRpc(ctx, md, req, append([]grpc.CallOption{grpc.Trailer(&respTrailers), grpc.Header(&respHeaders)}, opts...)...)
stat, ok := status.FromError(err) stat, ok := status.FromError(err)
if !ok { if !ok {
@ -191,10 +191,10 @@ func invokeUnary(ctx context.Context, stub grpcdynamic.Stub, md *desc.MethodDesc
} }
func invokeClientStream(ctx context.Context, stub grpcdynamic.Stub, md *desc.MethodDescriptor, handler InvocationEventHandler, func invokeClientStream(ctx context.Context, stub grpcdynamic.Stub, md *desc.MethodDescriptor, handler InvocationEventHandler,
requestData RequestSupplier, req proto.Message) error { requestData RequestSupplier, req proto.Message, opts ...grpc.CallOption) error {
// invoke the RPC! // invoke the RPC!
str, err := stub.InvokeRpcClientStream(ctx, md) str, err := stub.InvokeRpcClientStream(ctx, md, opts...)
// Upload each request message in the stream // Upload each request message in the stream
var resp proto.Message var resp proto.Message
@ -245,7 +245,7 @@ func invokeClientStream(ctx context.Context, stub grpcdynamic.Stub, md *desc.Met
} }
func invokeServerStream(ctx context.Context, stub grpcdynamic.Stub, md *desc.MethodDescriptor, handler InvocationEventHandler, func invokeServerStream(ctx context.Context, stub grpcdynamic.Stub, md *desc.MethodDescriptor, handler InvocationEventHandler,
requestData RequestSupplier, req proto.Message) error { requestData RequestSupplier, req proto.Message, opts ...grpc.CallOption) error {
err := requestData(req) err := requestData(req)
if err != nil && err != io.EOF { if err != nil && err != io.EOF {
@ -262,7 +262,7 @@ func invokeServerStream(ctx context.Context, stub grpcdynamic.Stub, md *desc.Met
} }
// Now we can actually invoke the RPC! // Now we can actually invoke the RPC!
str, err := stub.InvokeRpcServerStream(ctx, md, req) str, err := stub.InvokeRpcServerStream(ctx, md, req, opts...)
if respHeaders, err := str.Header(); err == nil { if respHeaders, err := str.Header(); err == nil {
handler.OnReceiveHeaders(respHeaders) handler.OnReceiveHeaders(respHeaders)
@ -294,13 +294,13 @@ func invokeServerStream(ctx context.Context, stub grpcdynamic.Stub, md *desc.Met
} }
func invokeBidi(ctx context.Context, stub grpcdynamic.Stub, md *desc.MethodDescriptor, handler InvocationEventHandler, func invokeBidi(ctx context.Context, stub grpcdynamic.Stub, md *desc.MethodDescriptor, handler InvocationEventHandler,
requestData RequestSupplier, req proto.Message) error { requestData RequestSupplier, req proto.Message, callOptions ...grpc.CallOption) error {
ctx, cancel := context.WithCancel(ctx) ctx, cancel := context.WithCancel(ctx)
defer cancel() defer cancel()
// invoke the RPC! // invoke the RPC!
str, err := stub.InvokeRpcBidiStream(ctx, md) str, err := stub.InvokeRpcBidiStream(ctx, md, callOptions...)
var wg sync.WaitGroup var wg sync.WaitGroup
var sendErr atomic.Value var sendErr atomic.Value