1use std::sync::Arc;
16
17use api::region::RegionResponse;
18use api::v1::ResponseHeader;
19use api::v1::region::RegionRequest;
20use arc_swap::ArcSwapOption;
21use arrow_flight::Ticket;
22use async_stream::stream;
23use async_trait::async_trait;
24use common_error::ext::BoxedError;
25use common_error::status_code::StatusCode;
26use common_grpc::flight::{FlightDecoder, FlightMessage};
27use common_meta::error::{self as meta_error, Result as MetaResult};
28use common_meta::node_manager::Datanode;
29use common_query::request::QueryRequest;
30use common_recordbatch::error::ExternalSnafu;
31use common_recordbatch::{RecordBatch, RecordBatchStreamWrapper, SendableRecordBatchStream};
32use common_telemetry::error;
33use common_telemetry::tracing_context::TracingContext;
34use prost::Message;
35use query::query_engine::DefaultSerializer;
36use snafu::{OptionExt, ResultExt, location};
37use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan};
38use tokio_stream::StreamExt;
39
40use crate::error::{
41 self, ConvertFlightDataSnafu, FlightGetSnafu, IllegalDatabaseResponseSnafu,
42 IllegalFlightMessagesSnafu, MissingFieldSnafu, Result, ServerSnafu,
43};
44use crate::{Client, Error, metrics};
45
46#[derive(Debug)]
47pub struct RegionRequester {
48 client: Client,
49 send_compression: bool,
50 accept_compression: bool,
51}
52
53#[async_trait]
54impl Datanode for RegionRequester {
55 async fn handle(&self, request: RegionRequest) -> MetaResult<RegionResponse> {
56 self.handle_inner(request).await.map_err(|err| {
57 if err.should_retry() {
58 meta_error::Error::RetryLater {
59 source: BoxedError::new(err),
60 clean_poisons: false,
61 }
62 } else {
63 meta_error::Error::External {
64 source: BoxedError::new(err),
65 location: location!(),
66 }
67 }
68 })
69 }
70
71 async fn handle_query(&self, request: QueryRequest) -> MetaResult<SendableRecordBatchStream> {
72 let plan = DFLogicalSubstraitConvertor
73 .encode(&request.plan, DefaultSerializer)
74 .map_err(BoxedError::new)
75 .context(meta_error::ExternalSnafu)?
76 .to_vec();
77 let request = api::v1::region::QueryRequest {
78 header: request.header,
79 region_id: request.region_id.as_u64(),
80 plan,
81 };
82
83 let ticket = Ticket {
84 ticket: request.encode_to_vec().into(),
85 };
86 self.do_get_inner(ticket)
87 .await
88 .map_err(BoxedError::new)
89 .context(meta_error::ExternalSnafu)
90 }
91}
92
93impl RegionRequester {
94 pub fn new(client: Client, send_compression: bool, accept_compression: bool) -> Self {
95 Self {
96 client,
97 send_compression,
98 accept_compression,
99 }
100 }
101
102 pub async fn do_get_inner(&self, ticket: Ticket) -> Result<SendableRecordBatchStream> {
103 let mut flight_client = self
104 .client
105 .make_flight_client(self.send_compression, self.accept_compression)?;
106 let response = flight_client
107 .mut_inner()
108 .do_get(ticket)
109 .await
110 .or_else(|e| {
111 let tonic_code = e.code();
112 let e: error::Error = e.into();
113 error!(
114 e; "Failed to do Flight get, addr: {}, code: {}",
115 flight_client.addr(),
116 tonic_code
117 );
118 Err(BoxedError::new(e)).with_context(|_| FlightGetSnafu {
119 addr: flight_client.addr().to_string(),
120 tonic_code,
121 })
122 })?;
123
124 let flight_data_stream = response.into_inner();
125 let mut decoder = FlightDecoder::default();
126
127 let mut flight_message_stream = flight_data_stream.map(move |flight_data| {
128 flight_data
129 .map_err(Error::from)
130 .and_then(|data| decoder.try_decode(&data).context(ConvertFlightDataSnafu))?
131 .context(IllegalFlightMessagesSnafu {
132 reason: "none message",
133 })
134 });
135
136 let Some(first_flight_message) = flight_message_stream.next().await else {
137 return IllegalFlightMessagesSnafu {
138 reason: "Expect the response not to be empty",
139 }
140 .fail();
141 };
142 let FlightMessage::Schema(schema) = first_flight_message? else {
143 return IllegalFlightMessagesSnafu {
144 reason: "Expect schema to be the first flight message",
145 }
146 .fail();
147 };
148
149 let metrics = Arc::new(ArcSwapOption::from(None));
150 let metrics_ref = metrics.clone();
151
152 let tracing_context = TracingContext::from_current_span();
153
154 let schema = Arc::new(
155 datatypes::schema::Schema::try_from(schema).context(error::ConvertSchemaSnafu)?,
156 );
157 let schema_cloned = schema.clone();
158 let stream = Box::pin(stream!({
159 let _span = tracing_context.attach(common_telemetry::tracing::info_span!(
160 "poll_flight_data_stream"
161 ));
162
163 let mut buffered_message: Option<FlightMessage> = None;
164 let mut stream_ended = false;
165
166 while !stream_ended {
167 let flight_message_item = if let Some(msg) = buffered_message.take() {
169 Some(Ok(msg))
170 } else {
171 flight_message_stream.next().await
172 };
173
174 let flight_message = match flight_message_item {
175 Some(Ok(message)) => message,
176 Some(Err(e)) => {
177 yield Err(BoxedError::new(e)).context(ExternalSnafu);
178 break;
179 }
180 None => break,
181 };
182
183 match flight_message {
184 FlightMessage::RecordBatch(record_batch) => {
185 let result_to_yield = RecordBatch::try_from_df_record_batch(
186 schema_cloned.clone(),
187 record_batch,
188 );
189
190 if let Some(next_flight_message_result) = flight_message_stream.next().await
192 {
193 match next_flight_message_result {
194 Ok(FlightMessage::Metrics(s)) => {
195 let m = serde_json::from_str(&s).ok().map(Arc::new);
196 metrics_ref.swap(m);
197 }
198 Ok(FlightMessage::RecordBatch(rb)) => {
199 buffered_message = Some(FlightMessage::RecordBatch(rb));
202 }
203 Ok(_) => {
204 yield IllegalFlightMessagesSnafu {
205 reason: "A RecordBatch message can only be succeeded by a Metrics message or another RecordBatch message"
206 }
207 .fail()
208 .map_err(BoxedError::new)
209 .context(ExternalSnafu);
210 break;
211 }
212 Err(e) => {
213 yield Err(BoxedError::new(e)).context(ExternalSnafu);
214 break;
215 }
216 }
217 } else {
218 stream_ended = true;
220 }
221
222 yield result_to_yield;
223 }
224 FlightMessage::Metrics(s) => {
225 let m = serde_json::from_str(&s).ok().map(Arc::new);
227 metrics_ref.swap(m);
228 break;
229 }
230 _ => {
231 yield IllegalFlightMessagesSnafu {
232 reason: "A Schema message must be succeeded exclusively by a set of RecordBatch messages"
233 }
234 .fail()
235 .map_err(BoxedError::new)
236 .context(ExternalSnafu);
237 break;
238 }
239 }
240 }
241 }));
242 let record_batch_stream = RecordBatchStreamWrapper {
243 schema,
244 stream,
245 output_ordering: None,
246 metrics,
247 };
248 Ok(Box::pin(record_batch_stream))
249 }
250
251 async fn handle_inner(&self, request: RegionRequest) -> Result<RegionResponse> {
252 let request_type = request
253 .body
254 .as_ref()
255 .with_context(|| MissingFieldSnafu { field: "body" })?
256 .as_ref()
257 .to_string();
258 let _timer = metrics::METRIC_REGION_REQUEST_GRPC
259 .with_label_values(&[request_type.as_str()])
260 .start_timer();
261
262 let (addr, mut client) = self.client.raw_region_client()?;
263
264 let response = client
265 .handle(request)
266 .await
267 .map_err(|e| {
268 let code = e.code();
269 error::Error::RegionServer {
271 addr,
272 code,
273 source: BoxedError::new(error::Error::from(e)),
274 location: location!(),
275 }
276 })?
277 .into_inner();
278
279 check_response_header(&response.header)?;
280
281 Ok(RegionResponse::from_region_response(response))
282 }
283
284 pub async fn handle(&self, request: RegionRequest) -> Result<RegionResponse> {
285 self.handle_inner(request).await
286 }
287}
288
289pub fn check_response_header(header: &Option<ResponseHeader>) -> Result<()> {
290 let status = header
291 .as_ref()
292 .and_then(|header| header.status.as_ref())
293 .context(IllegalDatabaseResponseSnafu {
294 err_msg: "either response header or status is missing",
295 })?;
296
297 if StatusCode::is_success(status.status_code) {
298 Ok(())
299 } else {
300 let code =
301 StatusCode::from_u32(status.status_code).context(IllegalDatabaseResponseSnafu {
302 err_msg: format!("unknown server status: {:?}", status),
303 })?;
304 ServerSnafu {
305 code,
306 msg: status.err_msg.clone(),
307 }
308 .fail()
309 }
310}
311
312#[cfg(test)]
313mod test {
314 use api::v1::Status as PbStatus;
315
316 use super::*;
317 use crate::Error::{IllegalDatabaseResponse, Server};
318
319 #[test]
320 fn test_check_response_header() {
321 let result = check_response_header(&None);
322 assert!(matches!(
323 result.unwrap_err(),
324 IllegalDatabaseResponse { .. }
325 ));
326
327 let result = check_response_header(&Some(ResponseHeader { status: None }));
328 assert!(matches!(
329 result.unwrap_err(),
330 IllegalDatabaseResponse { .. }
331 ));
332
333 let result = check_response_header(&Some(ResponseHeader {
334 status: Some(PbStatus {
335 status_code: StatusCode::Success as u32,
336 err_msg: String::default(),
337 }),
338 }));
339 assert!(result.is_ok());
340
341 let result = check_response_header(&Some(ResponseHeader {
342 status: Some(PbStatus {
343 status_code: u32::MAX,
344 err_msg: String::default(),
345 }),
346 }));
347 assert!(matches!(
348 result.unwrap_err(),
349 IllegalDatabaseResponse { .. }
350 ));
351
352 let result = check_response_header(&Some(ResponseHeader {
353 status: Some(PbStatus {
354 status_code: StatusCode::Internal as u32,
355 err_msg: "blabla".to_string(),
356 }),
357 }));
358 let Server { code, msg, .. } = result.unwrap_err() else {
359 unreachable!()
360 };
361 assert_eq!(code, StatusCode::Internal);
362 assert_eq!(msg, "blabla");
363 }
364}