grpcall/response.go
2023-05-30 17:30:06 +08:00

74 lines
1.9 KiB
Go

package grpcall
import (
"context"
"fmt"
"github.com/jhump/protoreflect/desc"
"google.golang.org/grpc/metadata"
)
type Response struct {
method *desc.MethodDescriptor
header *metadata.MD
data string
send chan string
recv chan string
done chan error
cancel context.CancelFunc
}
func (this Response) Header() metadata.MD {
return *this.header
}
func (this Response) Data() (string, error) {
if this.method.IsClientStreaming() || this.method.IsServerStreaming() {
return "", fmt.Errorf("%q is %s, cannot use unary data", this.method.GetFullyQualifiedName(), this.GetCategory())
}
return this.data, nil
}
// 发送数据流
func (this Response) Send() (chan<- string, error) {
if !this.method.IsClientStreaming() || !this.method.IsServerStreaming() {
return nil, fmt.Errorf("%q is %s, cannot use send streaming func", this.method.GetFullyQualifiedName(), this.GetCategory())
}
return this.send, nil
}
// 接收数据流
func (this Response) Recv() (<-chan string, error) {
if !this.method.IsServerStreaming() {
return nil, fmt.Errorf("%q is %s, cannot use streaming func", this.method.GetFullyQualifiedName(), this.GetCategory())
}
return this.recv, nil
}
// 数据是否响应结束
func (this Response) Done() (<-chan error, error) {
if !this.method.IsServerStreaming() {
return nil, fmt.Errorf("%q is %s, cannot use streaming func", this.method.GetFullyQualifiedName(), this.GetCategory())
}
return this.done, nil
}
// 取消接收数据
func (this Response) Cancel() {
if this.method.IsServerStreaming() {
this.cancel()
}
}
func (this Response) GetCategory() string {
if this.method.IsClientStreaming() && this.method.IsServerStreaming() {
return "bidi-streaming"
} else if this.method.IsClientStreaming() {
return "client-streaming"
} else if this.method.IsServerStreaming() {
return "server-streaming"
} else {
return "unary"
}
}