chyezh 86eff6e589
enhance: streaming node client implementation (#34653)
issue: #33285

- add streaming node grpc client wrapper
- add unittest for streaming node grpc client side
- fix binary unsafe bug for message

---------

Signed-off-by: chyezh <chyezh@outlook.com>
2024-07-19 17:37:40 +08:00

36 lines
1010 B
Go

package producer
import (
"github.com/milvus-io/milvus/internal/proto/streamingpb"
"github.com/milvus-io/milvus/pkg/streaming/util/message"
)
// produceGrpcClient is a wrapped producer server of log messages.
type produceGrpcClient struct {
streamingpb.StreamingNodeHandlerService_ProduceClient
}
// SendProduceMessage sends the produce message to server.
func (p *produceGrpcClient) SendProduceMessage(requestID int64, msg message.MutableMessage) error {
return p.Send(&streamingpb.ProduceRequest{
Request: &streamingpb.ProduceRequest_Produce{
Produce: &streamingpb.ProduceMessageRequest{
RequestId: requestID,
Message: &streamingpb.Message{
Payload: msg.Payload(),
Properties: msg.Properties().ToRawMap(),
},
},
},
})
}
// SendClose sends the close request to server.
func (p *produceGrpcClient) SendClose() error {
return p.Send(&streamingpb.ProduceRequest{
Request: &streamingpb.ProduceRequest_Close{
Close: &streamingpb.CloseProducerRequest{},
},
})
}