Line data Source code
1 : use std::io::{Read, Write};
2 :
3 : use bytes::{Buf, BufMut};
4 : use tonic::{codec::{Codec, DecodeBuf, Decoder, EncodeBuf, Encoder}, Status};
5 :
6 : type GenericGrpcClientData = Vec<u8>;
7 :
8 : /// A tonic codec that passes raw message bytes through without protobuf encoding
9 : #[derive(Debug, Clone, Default)]
10 : struct PassthroughCodec();
11 :
12 : /// A gRPC client that calls arbitrary methods with raw, already encoded message bytes
13 : pub struct GenericGrpcClient {
14 : /// The underlying tonic client on the gRPC channel
15 : inner: tonic::client::Grpc<tonic::transport::Channel>,
16 : }
17 :
18 : impl GenericGrpcClient {
19 12 : pub fn new(inner: tonic::transport::Channel) -> Self {
20 12 : let inner = tonic::client::Grpc::new(inner);
21 12 : Self { inner }
22 12 : }
23 :
24 4 : pub async fn call_unary(
25 4 : &mut self,
26 4 : request: tonic::Request<GenericGrpcClientData>,
27 4 : path: tonic::codegen::http::uri::PathAndQuery,
28 4 : ) -> std::result::Result<
29 4 : tonic::Response<GenericGrpcClientData>,
30 4 : tonic::Status,
31 4 : > {
32 4 : self.inner
33 4 : .ready()
34 4 : .await
35 4 : .map_err(|e| {
36 0 : tonic::Status::new(
37 0 : tonic::Code::Unknown,
38 0 : format!("Service was not ready: {}", e),
39 : )
40 0 : })?;
41 4 : let codec = PassthroughCodec::default();
42 4 : self.inner.unary(request, path, codec).await
43 4 : }
44 :
45 6 : pub async fn call_server_streaming(
46 6 : &mut self,
47 6 : request: tonic::Request<GenericGrpcClientData>,
48 6 : path: tonic::codegen::http::uri::PathAndQuery,
49 6 : ) -> std::result::Result<
50 6 : tonic::Response<tonic::codec::Streaming<GenericGrpcClientData>>,
51 6 : tonic::Status,
52 6 : > {
53 6 : self.inner
54 6 : .ready()
55 6 : .await
56 6 : .map_err(|e| {
57 0 : tonic::Status::new(
58 0 : tonic::Code::Unknown,
59 0 : format!("Service was not ready: {}", e),
60 : )
61 0 : })?;
62 6 : let codec = PassthroughCodec::default();
63 6 : self.inner.server_streaming(request, path, codec).await
64 6 : }
65 : }
66 :
67 : /// The encoder of the [`PassthroughCodec`], writing message bytes as is
68 : #[derive(Debug, Clone, Default)]
69 : struct PassthroughEncoder();
70 : /// The decoder of the [`PassthroughCodec`], reading message bytes as is
71 : #[derive(Debug, Clone, Default)]
72 : struct PassthroughDecoder();
73 :
74 : /// By default, Tonic uses a prost codec.
75 : /// We use a dedicated passthrough codec that bypasses prost and writes parameter bytes as-is.
76 : impl Codec for PassthroughCodec {
77 : type Encode = GenericGrpcClientData;
78 : type Decode = GenericGrpcClientData;
79 :
80 : type Encoder = PassthroughEncoder;
81 : type Decoder = PassthroughDecoder;
82 :
83 10 : fn encoder(&mut self) -> Self::Encoder {
84 10 : PassthroughEncoder()
85 10 : }
86 10 : fn decoder(&mut self) -> Self::Decoder {
87 10 : PassthroughDecoder()
88 10 : }
89 : }
90 : impl Encoder for PassthroughEncoder {
91 : type Item = GenericGrpcClientData;
92 : type Error = Status;
93 10 : fn encode(&mut self, item: Self::Item, buf: &mut EncodeBuf<'_>) -> Result<(), Self::Error> {
94 10 : buf.writer().write(&item)?;
95 10 : Ok(())
96 10 : }
97 : }
98 : impl Decoder for PassthroughDecoder {
99 : type Item = GenericGrpcClientData;
100 : type Error = Status;
101 19 : fn decode(&mut self, buf: &mut DecodeBuf<'_>) -> Result<Option<Self::Item>, Self::Error> {
102 19 : let mut buffer = Vec::new();
103 19 : buf.reader().read_to_end(&mut buffer)?;
104 19 : Ok(Some(buffer))
105 19 : }
106 : }
107 :
108 : #[cfg(test)]
109 : mod test {
110 : use super::*;
111 : use anyhow::Result;
112 : use prost::Message;
113 : use tokio_stream::StreamExt;
114 :
115 : use crate::{proto::dashql_test::{TestServerStreamingRequest, TestServerStreamingResponse, TestUnaryRequest, TestUnaryResponse}, test::grpc_service_mock::{spawn_grpc_test_service_mock, GrpcServiceMock}};
116 :
117 : #[tokio::test]
118 1 : async fn test_unary() -> Result<()> {
119 : // Spawn a test service mock
120 1 : let (mock, mut setup_unary, mut _setup_server_streaming) = GrpcServiceMock::new();
121 1 : let (addr, shutdown) = spawn_grpc_test_service_mock(mock).await;
122 :
123 : // Respond single unary response
124 1 : let unary_call = tokio::spawn(async move {
125 1 : let (param, result_sender) = setup_unary.recv().await.unwrap();
126 1 : result_sender.send(Ok(TestUnaryResponse {
127 1 : data: "response data".to_string()
128 1 : })).await.unwrap();
129 1 : param
130 1 : });
131 :
132 : // Setup the gRPC client
133 1 : let channel = tonic::transport::Endpoint::new(format!("http://{}", addr))?.connect().await?;
134 1 : let mut client = GenericGrpcClient::new(channel);
135 :
136 : // Call the unary test function
137 1 : let unary_path = tonic::codegen::http::uri::PathAndQuery::from_static(
138 : "/dashql.test.TestService/TestUnary",
139 : );
140 1 : let unary_param = TestUnaryRequest {
141 1 : data: "request data".to_string()
142 1 : };
143 1 : let unary_req = tonic::Request::new(unary_param.encode_to_vec());
144 1 : let unary_result = client.call_unary(unary_req, unary_path).await?;
145 :
146 : // Check received parameter
147 1 : let received_param = unary_call.await?;
148 1 : assert_eq!(received_param.data, "request data");
149 :
150 : // Check received buffer
151 1 : let received_buffer = unary_result.into_inner();
152 1 : let received_message = TestUnaryResponse::decode(received_buffer.as_slice())?;
153 1 : assert_eq!(received_message.data, "response data");
154 :
155 1 : shutdown.send(()).unwrap();
156 2 : Ok(())
157 1 : }
158 :
159 : #[tokio::test]
160 1 : async fn test_server_streaming() -> Result<()> {
161 : // Spawn a test service mock
162 1 : let (mock, mut _setup_unary, mut setup_server_streaming) = GrpcServiceMock::new();
163 1 : let (addr, shutdown) = spawn_grpc_test_service_mock(mock).await;
164 :
165 : // Respond single streaming response
166 1 : let streaming_call = tokio::spawn(async move {
167 1 : let (param, result_sender) = setup_server_streaming.recv().await.unwrap();
168 1 : result_sender.send(Ok(TestServerStreamingResponse {
169 1 : data: "response data".to_string()
170 1 : })).await.unwrap();
171 1 : param
172 1 : });
173 :
174 : // Setup the gRPC client
175 1 : let channel = tonic::transport::Endpoint::new(format!("http://{}", addr))?.connect().await?;
176 1 : let mut client = GenericGrpcClient::new(channel);
177 :
178 : // Call the unary test function
179 1 : let streaming_path = tonic::codegen::http::uri::PathAndQuery::from_static(
180 : "/dashql.test.TestService/TestServerStreaming",
181 : );
182 1 : let streaming_param = TestServerStreamingRequest {
183 1 : data: "request data".to_string()
184 1 : };
185 1 : let streaming_req = tonic::Request::new(streaming_param.encode_to_vec());
186 1 : let streaming_response = client.call_server_streaming(streaming_req, streaming_path).await?;
187 :
188 : // Check received parameter
189 1 : let received_param = streaming_call.await?;
190 1 : assert_eq!(received_param.data, "request data");
191 :
192 : // Check received buffer
193 1 : let mut stream = streaming_response.into_inner();
194 2 : while let Some(response) = stream.next().await {
195 1 : let received_message = TestUnaryResponse::decode(response.unwrap().as_slice())?;
196 1 : assert_eq!(received_message.data, "response data");
197 : }
198 :
199 1 : shutdown.send(()).unwrap();
200 2 : Ok(())
201 1 : }
202 : }
|
