grpcall/response.go

package grpcall

import (
	"context"
	"encoding/json"
	"fmt"
	"io"

	"github.com/jhump/protoreflect/desc"
	"google.golang.org/grpc/metadata"
)

type Response struct {
	method        *desc.MethodDescriptor
	header        *metadata.MD
	recv          *io.PipeReader
	send          chan []byte
	sendCompleted chan error
	cancel        context.CancelFunc
}

func (this Response) Header() metadata.MD {
	return *this.header
}

func (this Response) Recv() io.Reader {
	return this.recv
}

func (this Response) RecvDecoder() *json.Decoder {
	return json.NewDecoder(this.Recv())
}

// 发送数据流
func (this Response) Send(data string) error {
	if !this.method.IsClientStreaming() || !this.method.IsServerStreaming() {
		return fmt.Errorf("%q is %s, not streaming gRPC", this.method.GetFullyQualifiedName(), this.GetCategory())
	}

	this.send <- []byte(data)

	return <-this.sendCompleted
}

// 取消接收数据
func (this Response) Close() error {
	if this.method.IsServerStreaming() {
		this.cancel()
	}

	return this.recv.Close()
}

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"
	}
}