1use std::sync::Arc;
16use std::time::Duration;
17
18use api::region::RegionResponse;
19use api::v1::ResponseHeader;
20use api::v1::region::{
21 RegionRequest, RegionRequestHeader, RemoteDynFilterRequest, RemoteDynFilterUnregister,
22 RemoteDynFilterUpdate, region_request, remote_dyn_filter_request,
23};
24use arc_swap::ArcSwapOption;
25use arrow_flight::Ticket;
26use async_stream::stream;
27use async_trait::async_trait;
28use common_error::ext::{BoxedError, ErrorExt};
29use common_error::status_code::StatusCode;
30use common_grpc::flight::{FlightDecoder, FlightMessage};
31use common_meta::error::{self as meta_error, Result as MetaResult};
32use common_meta::node_manager::Datanode;
33use common_query::request::QueryRequest;
34use common_recordbatch::error::ExternalSnafu;
35use common_recordbatch::{RecordBatch, RecordBatchStreamWrapper, SendableRecordBatchStream};
36use common_telemetry::error;
37use common_telemetry::tracing::Span;
38use common_telemetry::tracing_context::TracingContext;
39use futures_util::Stream;
40use prost::Message;
41use query::query_engine::DefaultSerializer;
42use snafu::{OptionExt, ResultExt, location};
43use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan};
44use tokio_stream::StreamExt;
45use tonic::codec::CompressionEncoding;
46
47use crate::error::{
48 self, FlightGetSnafu, IllegalDatabaseResponseSnafu, IllegalFlightMessagesSnafu,
49 MissingFieldSnafu, Result, ServerSnafu,
50};
51use crate::flight::{FlightMessageKind, FlightMessageReader, decode_flight_data};
52use crate::{Client, metrics};
53
54const FLIGHT_DO_GET_TIMEOUT: Duration = Duration::from_secs(10);
55
56#[derive(Debug)]
57pub struct RegionRequester {
58 client: Client,
59 send_compression: bool,
60 accept_compression: bool,
61}
62
63#[async_trait]
64impl Datanode for RegionRequester {
65 async fn handle(&self, request: RegionRequest) -> MetaResult<RegionResponse> {
66 self.handle_inner(request).await.map_err(|err| {
67 if err.should_retry() {
68 meta_error::Error::RetryLater {
69 source: BoxedError::new(err),
70 clean_poisons: false,
71 }
72 } else {
73 meta_error::Error::External {
74 source: BoxedError::new(err),
75 location: location!(),
76 }
77 }
78 })
79 }
80
81 async fn handle_query(&self, request: QueryRequest) -> MetaResult<SendableRecordBatchStream> {
82 let plan = DFLogicalSubstraitConvertor
83 .encode(&request.plan, DefaultSerializer)
84 .map_err(BoxedError::new)
85 .context(meta_error::ExternalSnafu)?
86 .to_vec();
87 let request = api::v1::region::QueryRequest {
88 header: request.header,
89 region_id: request.region_id.as_u64(),
90 plan,
91 };
92
93 let ticket = Ticket {
94 ticket: request.encode_to_vec().into(),
95 };
96 self.do_get_inner(ticket)
97 .await
98 .map_err(BoxedError::new)
99 .context(meta_error::ExternalSnafu)
100 }
101}
102
103impl RegionRequester {
104 pub fn new(client: Client, send_compression: bool, accept_compression: bool) -> Self {
105 Self {
106 client,
107 send_compression,
108 accept_compression,
109 }
110 }
111
112 pub async fn do_get_inner(&self, ticket: Ticket) -> Result<SendableRecordBatchStream> {
113 let mut flight_client = self
114 .client
115 .make_flight_client(self.send_compression, self.accept_compression)?;
116 let addr = flight_client.addr().to_string();
118 let mut request = tonic::Request::new(ticket);
119 request.set_timeout(FLIGHT_DO_GET_TIMEOUT);
120 let response = flight_client
121 .mut_inner()
122 .do_get(request)
123 .await
124 .or_else(|e| {
125 let tonic_code = e.code();
126 let e: error::Error = e.into();
127 error!(
128 e; "Failed to do Flight get, addr: {}, code: {}",
129 addr,
130 tonic_code
131 );
132 Err(BoxedError::new(e)).with_context(|_| FlightGetSnafu {
133 addr: addr.clone(),
134 tonic_code,
135 })
136 })?;
137
138 let flight_data_stream = response.into_inner();
139 let mut decoder = FlightDecoder::default();
140
141 let flight_message_stream = flight_data_stream
142 .filter_map(move |flight_data| decode_flight_data(&mut decoder, flight_data));
143
144 recordbatches_from_flight_message_stream(addr, flight_message_stream).await
145 }
146
147 async fn handle_inner(&self, request: RegionRequest) -> Result<RegionResponse> {
148 let request_body = request
149 .body
150 .as_ref()
151 .with_context(|| MissingFieldSnafu { field: "body" })?;
152 let is_insert = matches!(
153 request_body,
154 region_request::Body::Inserts(_) | region_request::Body::BulkInsert(_)
155 );
156 let request_type = request_body.as_ref().to_string();
157 let _timer = metrics::METRIC_REGION_REQUEST_GRPC
158 .with_label_values(&[request_type.as_str()])
159 .start_timer();
160
161 let (addr, mut client) = self.client.raw_region_client()?;
162 if is_insert {
163 if self.send_compression {
164 client = client.send_compressed(CompressionEncoding::Zstd);
165 }
166 if self.accept_compression {
167 client = client.accept_compressed(CompressionEncoding::Zstd);
168 }
169 }
170
171 let response = client
172 .handle(request)
173 .await
174 .map_err(|e| {
175 let code = e.code();
176 error::Error::RegionServer {
178 addr,
179 code,
180 source: BoxedError::new(error::Error::from(e)),
181 location: location!(),
182 }
183 })?
184 .into_inner();
185
186 check_response_header(&response.header)?;
187
188 Ok(RegionResponse::from_region_response(response))
189 }
190
191 pub async fn handle(&self, request: RegionRequest) -> Result<RegionResponse> {
192 self.handle_inner(request).await
193 }
194
195 pub async fn handle_remote_dyn_filter_update(
196 &self,
197 query_id: impl Into<String>,
198 update: RemoteDynFilterUpdate,
199 ) -> Result<RegionResponse> {
200 self.handle_inner(build_remote_dyn_filter_update_request(query_id, update))
201 .await
202 }
203
204 pub async fn handle_remote_dyn_filter_unregister(
205 &self,
206 query_id: impl Into<String>,
207 unregister: RemoteDynFilterUnregister,
208 ) -> Result<RegionResponse> {
209 self.handle_inner(build_remote_dyn_filter_unregister_request(
210 query_id, unregister,
211 ))
212 .await
213 }
214}
215
216async fn recordbatches_from_flight_message_stream<S>(
217 addr: String,
218 flight_message_stream: S,
219) -> Result<SendableRecordBatchStream>
220where
221 S: Stream<Item = Result<FlightMessage>> + Send + Unpin + 'static,
222{
223 let mut reader = FlightMessageReader::new(addr.clone(), flight_message_stream);
224 let FlightMessage::Schema(schema) = reader
225 .read_first()
226 .await
227 .map_err(|error| flight_stream_error(reader.remote_addr(), error))?
228 else {
229 return IllegalFlightMessagesSnafu {
230 reason: "Expect schema to be the first flight message",
231 }
232 .fail()
233 .map_err(|error| flight_stream_error(reader.remote_addr(), error));
234 };
235
236 let metrics = Arc::new(ArcSwapOption::from(None));
237 let metrics_ref = metrics.clone();
238
239 let tracing_context = TracingContext::from_current_span();
240
241 let schema =
242 Arc::new(datatypes::schema::Schema::try_from(schema).context(error::ConvertSchemaSnafu)?);
243 let schema_cloned = schema.clone();
244 let stream_addr = addr;
245 let stream = Box::pin(stream!({
246 let _span = tracing_context.attach(common_telemetry::tracing::info_span!(
247 "poll_flight_data_stream"
248 ));
249
250 let mut stream_ended = false;
251
252 while !stream_ended {
253 let flight_message = match reader.read_next().await {
254 Ok(Some(message)) => message,
255 Ok(None) => break,
256 Err(error) => {
257 yield Err(BoxedError::new(flight_stream_error(&stream_addr, error)))
258 .context(ExternalSnafu);
259 break;
260 }
261 };
262
263 match flight_message {
264 FlightMessage::RecordBatch(record_batch) => {
265 let result_to_yield =
266 RecordBatch::from_df_record_batch(schema_cloned.clone(), record_batch);
267
268 match reader.peek_next_message_kind().await {
270 Ok(Some(FlightMessageKind::Metrics)) => {
271 let metrics_message = match reader.read_next().await {
272 Ok(Some(FlightMessage::Metrics(metrics))) => metrics,
273 Ok(Some(_) | None) => {
274 yield IllegalFlightMessagesSnafu {
275 reason: "Flight stream changed after peek",
276 }
277 .fail()
278 .map_err(BoxedError::new)
279 .context(ExternalSnafu);
280 break;
281 }
282 Err(error) => {
283 yield Err(BoxedError::new(flight_stream_error(
284 &stream_addr,
285 error,
286 )))
287 .context(ExternalSnafu);
288 break;
289 }
290 };
291 let metrics = serde_json::from_str(&metrics_message).ok().map(Arc::new);
292 metrics_ref.swap(metrics);
293 }
294 Ok(Some(FlightMessageKind::RecordBatch)) => {}
295 Ok(Some(FlightMessageKind::Schema | FlightMessageKind::AffectedRows)) => {
296 yield IllegalFlightMessagesSnafu {
297 reason: "A RecordBatch message can only be succeeded by a Metrics message or another RecordBatch message"
298 }
299 .fail()
300 .map_err(BoxedError::new)
301 .context(ExternalSnafu);
302 break;
303 }
304 Ok(None) => stream_ended = true,
305 Err(error) => {
306 yield Err(BoxedError::new(flight_stream_error(&stream_addr, error)))
307 .context(ExternalSnafu);
308 break;
309 }
310 }
311
312 yield Ok(result_to_yield);
313 }
314 FlightMessage::Metrics(s) => {
315 let m = serde_json::from_str(&s).ok().map(Arc::new);
317 metrics_ref.swap(m);
318 continue;
319 }
320 _ => {
321 yield IllegalFlightMessagesSnafu {
322 reason: "A Schema message must be succeeded exclusively by a set of RecordBatch messages"
323 }
324 .fail()
325 .map_err(BoxedError::new)
326 .context(ExternalSnafu);
327 break;
328 }
329 }
330 }
331 }));
332 let record_batch_stream = RecordBatchStreamWrapper {
333 schema,
334 stream,
335 output_ordering: None,
336 metrics,
337 span: Span::current(),
338 };
339 Ok(Box::pin(record_batch_stream))
340}
341
342fn flight_stream_error(addr: &str, error: error::Error) -> error::Error {
343 let tonic_code = error.tonic_code().unwrap_or(tonic::Code::Unknown);
344 if error.status_code().should_log_error() {
345 error!(
346 error; "Failed to receive Flight data, addr: {}, code: {}",
347 addr,
348 tonic_code
349 );
350 }
351
352 error::Error::FlightGet {
353 addr: addr.to_string(),
354 tonic_code,
355 source: BoxedError::new(error),
356 }
357}
358
359pub fn build_remote_dyn_filter_update_request(
360 query_id: impl Into<String>,
361 update: RemoteDynFilterUpdate,
362) -> RegionRequest {
363 build_remote_dyn_filter_request(
364 query_id.into(),
365 remote_dyn_filter_request::Action::Update(update),
366 )
367}
368
369pub fn build_remote_dyn_filter_unregister_request(
370 query_id: impl Into<String>,
371 unregister: RemoteDynFilterUnregister,
372) -> RegionRequest {
373 build_remote_dyn_filter_request(
374 query_id.into(),
375 remote_dyn_filter_request::Action::Unregister(unregister),
376 )
377}
378
379fn build_remote_dyn_filter_request(
380 query_id: String,
381 action: remote_dyn_filter_request::Action,
382) -> RegionRequest {
383 RegionRequest {
384 header: Some(RegionRequestHeader {
385 tracing_context: TracingContext::from_current_span().to_w3c(),
386 ..Default::default()
387 }),
388 body: Some(region_request::Body::RemoteDynFilter(
389 RemoteDynFilterRequest {
390 query_id,
391 action: Some(action),
392 },
393 )),
394 }
395}
396
397pub fn check_response_header(header: &Option<ResponseHeader>) -> Result<()> {
398 let status = header
399 .as_ref()
400 .and_then(|header| header.status.as_ref())
401 .context(IllegalDatabaseResponseSnafu {
402 err_msg: "either response header or status is missing",
403 })?;
404
405 if StatusCode::is_success(status.status_code) {
406 Ok(())
407 } else {
408 let code =
409 StatusCode::from_u32(status.status_code).context(IllegalDatabaseResponseSnafu {
410 err_msg: format!("unknown server status: {:?}", status),
411 })?;
412 ServerSnafu {
413 code,
414 msg: status.err_msg.clone(),
415 }
416 .fail()
417 }
418}
419
420#[cfg(test)]
421#[allow(deprecated)]
422mod test {
423 use api::v1::Status as PbStatus;
424 use api::v1::region::region_server::{Region, RegionServer};
425 use api::v1::region::{
426 BulkInsertRequest, RegionResponse as PbRegionResponse, RemoteDynFilterUnregister,
427 RemoteDynFilterUpdate, region_request, remote_dyn_filter_request,
428 };
429 use common_recordbatch::adapter::RecordBatchMetrics;
430 use datatypes::arrow::array::Int32Array;
431 use datatypes::prelude::{ConcreteDataType, VectorRef};
432 use datatypes::schema::{ColumnSchema, Schema};
433 use datatypes::vectors::Int32Vector;
434 use futures_util::stream;
435 use tokio::net::TcpListener;
436 use tokio_stream::wrappers::TcpListenerStream;
437 use tonic::codec::CompressionEncoding;
438 use tonic::{Request, Response, Status};
439
440 use super::*;
441 use crate::Error::{self, IllegalDatabaseResponse, Server};
442
443 #[derive(Clone)]
444 struct CompressionRecordingRegionService {
445 zstd_headers: Arc<std::sync::Mutex<Vec<(bool, bool)>>>,
446 }
447
448 #[tonic::async_trait]
449 impl Region for CompressionRecordingRegionService {
450 async fn handle(
451 &self,
452 request: Request<RegionRequest>,
453 ) -> std::result::Result<Response<PbRegionResponse>, Status> {
454 let metadata = request.metadata();
455 let sends_zstd = metadata
456 .get("grpc-encoding")
457 .and_then(|v| v.to_str().ok())
458 .is_some_and(|value| value == "zstd");
459 let accepts_zstd = metadata
460 .get("grpc-accept-encoding")
461 .and_then(|value| value.to_str().ok())
462 .is_some_and(|value| value.split(',').any(|encoding| encoding == "zstd"));
463 self.zstd_headers
464 .lock()
465 .unwrap()
466 .push((sends_zstd, accepts_zstd));
467
468 Ok(Response::new(PbRegionResponse {
469 header: Some(ResponseHeader {
470 status: Some(PbStatus {
471 status_code: StatusCode::Success as u32,
472 ..Default::default()
473 }),
474 }),
475 ..Default::default()
476 }))
477 }
478 }
479
480 #[tokio::test]
481 async fn test_inserts_and_bulk_insert_use_transport_compression() {
482 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
483 let addr = listener.local_addr().unwrap();
484 let zstd_headers = Arc::new(std::sync::Mutex::new(Vec::new()));
485 let service = CompressionRecordingRegionService {
486 zstd_headers: zstd_headers.clone(),
487 };
488 let server = tokio::spawn(async move {
489 tonic::transport::Server::builder()
490 .add_service(
491 RegionServer::new(service)
492 .accept_compressed(CompressionEncoding::Zstd)
493 .send_compressed(CompressionEncoding::Zstd),
494 )
495 .serve_with_incoming(TcpListenerStream::new(listener))
496 .await
497 .unwrap();
498 });
499 let client = Client::with_urls([addr.to_string()]);
500 let requester = RegionRequester::new(client.clone(), true, true);
501 let send_only_requester = RegionRequester::new(client.clone(), true, false);
502 let accept_only_requester = RegionRequester::new(client.clone(), false, true);
503 let disabled_requester = RegionRequester::new(client, false, false);
504 let inserts = || RegionRequest {
505 body: Some(region_request::Body::Inserts(Default::default())),
506 ..Default::default()
507 };
508 let bulk_insert = || RegionRequest {
509 body: Some(region_request::Body::BulkInsert(
510 BulkInsertRequest::default(),
511 )),
512 ..Default::default()
513 };
514
515 requester.handle(inserts()).await.unwrap();
516 send_only_requester.handle(inserts()).await.unwrap();
517 accept_only_requester.handle(inserts()).await.unwrap();
518 requester.handle(bulk_insert()).await.unwrap();
519 disabled_requester.handle(bulk_insert()).await.unwrap();
520 requester
521 .handle(build_remote_dyn_filter_unregister_request(
522 "query-1",
523 RemoteDynFilterUnregister {
524 filter_id: "filter-1".to_string(),
525 },
526 ))
527 .await
528 .unwrap();
529
530 assert_eq!(
531 vec![
532 (true, true),
533 (true, false),
534 (false, true),
535 (true, true),
536 (false, false),
537 (false, false)
538 ],
539 *zstd_headers.lock().unwrap()
540 );
541 server.abort();
542 }
543
544 #[test]
545 fn test_flight_stream_error_preserves_peer_address() {
546 let error = flight_stream_error(
547 "127.0.0.1:4001",
548 tonic::Status::unavailable("datanode unavailable").into(),
549 );
550
551 assert!(matches!(
552 error,
553 error::Error::FlightGet {
554 addr,
555 tonic_code: tonic::Code::Unavailable,
556 ..
557 } if addr == "127.0.0.1:4001"
558 ));
559 }
560
561 #[tokio::test]
562 async fn test_empty_flight_stream_preserves_peer_address() {
563 let Err(error) = recordbatches_from_flight_message_stream(
564 "127.0.0.1:4001".to_string(),
565 stream::empty::<Result<FlightMessage>>(),
566 )
567 .await
568 else {
569 panic!("expected empty Flight stream to fail");
570 };
571
572 assert!(matches!(
573 error,
574 error::Error::FlightGet {
575 addr,
576 tonic_code: tonic::Code::Unknown,
577 ..
578 } if addr == "127.0.0.1:4001"
579 ));
580 }
581
582 fn test_schema() -> Arc<Schema> {
583 Arc::new(Schema::new(vec![ColumnSchema::new(
584 "v",
585 ConcreteDataType::int32_datatype(),
586 false,
587 )]))
588 }
589
590 fn test_metrics_json() -> String {
591 serde_json::to_string(&RecordBatchMetrics {
592 elapsed_compute: 7,
593 ..Default::default()
594 })
595 .unwrap()
596 }
597
598 #[test]
599 fn test_check_response_header() {
600 let result = check_response_header(&None);
601 assert!(matches!(
602 result.unwrap_err(),
603 IllegalDatabaseResponse { .. }
604 ));
605
606 let result = check_response_header(&Some(ResponseHeader { status: None }));
607 assert!(matches!(
608 result.unwrap_err(),
609 IllegalDatabaseResponse { .. }
610 ));
611
612 let result = check_response_header(&Some(ResponseHeader {
613 status: Some(PbStatus {
614 status_code: StatusCode::Success as u32,
615 err_msg: String::default(),
616 }),
617 }));
618 assert!(result.is_ok());
619
620 let result = check_response_header(&Some(ResponseHeader {
621 status: Some(PbStatus {
622 status_code: u32::MAX,
623 err_msg: String::default(),
624 }),
625 }));
626 assert!(matches!(
627 result.unwrap_err(),
628 IllegalDatabaseResponse { .. }
629 ));
630
631 let result = check_response_header(&Some(ResponseHeader {
632 status: Some(PbStatus {
633 status_code: StatusCode::Internal as u32,
634 err_msg: "blabla".to_string(),
635 }),
636 }));
637 let Server { code, msg, .. } = result.unwrap_err() else {
638 unreachable!()
639 };
640 assert_eq!(code, StatusCode::Internal);
641 assert_eq!(msg, "blabla");
642 }
643
644 #[test]
645 fn test_build_remote_dyn_filter_request_sets_header_and_body() {
646 let request = build_remote_dyn_filter_update_request(
647 "query-1",
648 RemoteDynFilterUpdate {
649 filter_id: "filter-1".to_string(),
650 payload: vec![1, 2, 3],
651 generation: 7,
652 is_complete: false,
653 },
654 );
655
656 request.header.expect("remote dyn filter header must exist");
657
658 let body = request.body.expect("remote dyn filter body must exist");
659 let region_request::Body::RemoteDynFilter(remote_request) = body else {
660 panic!("expected remote dyn filter request body");
661 };
662
663 assert_eq!(remote_request.query_id, "query-1");
664 assert!(matches!(
665 remote_request.action,
666 Some(remote_dyn_filter_request::Action::Update(_))
667 ));
668 }
669
670 #[test]
671 fn test_build_remote_dyn_filter_unregister_request_sets_header_and_body() {
672 let request = build_remote_dyn_filter_unregister_request(
673 "query-1",
674 RemoteDynFilterUnregister {
675 filter_id: "filter-9".to_string(),
676 },
677 );
678
679 request.header.expect("remote dyn filter header must exist");
680
681 let body = request.body.expect("remote dyn filter body must exist");
682 let region_request::Body::RemoteDynFilter(remote_request) = body else {
683 panic!("expected remote dyn filter request body");
684 };
685
686 assert_eq!(remote_request.query_id, "query-1");
687 assert!(matches!(
688 remote_request.action,
689 Some(remote_dyn_filter_request::Action::Unregister(_))
690 ));
691 }
692
693 #[tokio::test]
694 async fn test_record_batch_stream_continues_after_pre_batch_metrics() {
695 let schema = test_schema();
696 let batch = RecordBatch::new(
697 schema.clone(),
698 vec![Arc::new(Int32Vector::from_slice([1])) as VectorRef],
699 )
700 .unwrap();
701
702 let mut recordbatches = recordbatches_from_flight_message_stream(
703 "test-peer".to_string(),
704 stream::iter(vec![
705 Ok(FlightMessage::Schema(schema.arrow_schema().clone())),
706 Ok(FlightMessage::Metrics(test_metrics_json())),
707 Ok(FlightMessage::RecordBatch(batch.into_df_record_batch())),
708 ]),
709 )
710 .await
711 .unwrap();
712
713 let batch = recordbatches.next().await.unwrap().unwrap();
714 assert_eq!(batch.num_rows(), 1);
715 assert!(recordbatches.next().await.is_none());
716
717 let metrics = recordbatches.metrics().unwrap();
718 assert_eq!(metrics.elapsed_compute, 7);
719 }
720
721 #[tokio::test]
722 async fn test_record_batch_stream_updates_following_metrics_before_yielding_batch() {
723 let schema = test_schema();
724 let batch = RecordBatch::new(
725 schema.clone(),
726 vec![Arc::new(Int32Vector::from_slice([1])) as VectorRef],
727 )
728 .unwrap();
729 let mut recordbatches = recordbatches_from_flight_message_stream(
730 "test-peer".to_string(),
731 stream::iter(vec![
732 Ok(FlightMessage::Schema(schema.arrow_schema().clone())),
733 Ok(FlightMessage::RecordBatch(batch.into_df_record_batch())),
734 Ok(FlightMessage::Metrics(test_metrics_json())),
735 ]),
736 )
737 .await
738 .unwrap();
739
740 let batch = recordbatches.next().await.unwrap().unwrap();
741 assert_eq!(batch.num_rows(), 1);
742
743 let metrics = recordbatches.metrics().unwrap();
744 assert_eq!(metrics.elapsed_compute, 7);
745 assert!(recordbatches.next().await.is_none());
746 }
747
748 #[tokio::test]
749 async fn test_record_batch_stream_preserves_peeked_record_batch() {
750 let schema = test_schema();
751 let first_batch = RecordBatch::new(
752 schema.clone(),
753 vec![Arc::new(Int32Vector::from_slice([1])) as VectorRef],
754 )
755 .unwrap();
756 let second_batch = RecordBatch::new(
757 schema.clone(),
758 vec![Arc::new(Int32Vector::from_slice([2])) as VectorRef],
759 )
760 .unwrap();
761
762 let mut recordbatches = recordbatches_from_flight_message_stream(
763 "test-peer".to_string(),
764 stream::iter(vec![
765 Ok(FlightMessage::Schema(schema.arrow_schema().clone())),
766 Ok(FlightMessage::RecordBatch(
767 first_batch.into_df_record_batch(),
768 )),
769 Ok(FlightMessage::RecordBatch(
770 second_batch.into_df_record_batch(),
771 )),
772 ]),
773 )
774 .await
775 .unwrap();
776
777 let first_batch = recordbatches.next().await.unwrap().unwrap();
778 let second_batch = recordbatches.next().await.unwrap().unwrap();
779 assert_eq!(
780 first_batch
781 .column(0)
782 .as_any()
783 .downcast_ref::<Int32Array>()
784 .unwrap()
785 .value(0),
786 1
787 );
788 assert_eq!(
789 second_batch
790 .column(0)
791 .as_any()
792 .downcast_ref::<Int32Array>()
793 .unwrap()
794 .value(0),
795 2
796 );
797 assert!(recordbatches.next().await.is_none());
798 }
799
800 #[tokio::test]
801 async fn test_record_batch_stream_captures_final_metrics_after_record_batch() {
802 let schema = test_schema();
803 let batch = RecordBatch::new(
804 schema.clone(),
805 vec![Arc::new(Int32Vector::from_slice([1])) as VectorRef],
806 )
807 .unwrap();
808 let final_metrics = serde_json::to_string(&RecordBatchMetrics {
809 elapsed_compute: 99,
810 ..Default::default()
811 })
812 .unwrap();
813 let mut recordbatches = recordbatches_from_flight_message_stream(
814 "test-peer".to_string(),
815 stream::iter(vec![
816 Ok(FlightMessage::Schema(schema.arrow_schema().clone())),
817 Ok(FlightMessage::RecordBatch(batch.into_df_record_batch())),
818 Ok(FlightMessage::Metrics(final_metrics)),
819 ]),
820 )
821 .await
822 .unwrap();
823
824 assert_eq!(recordbatches.next().await.unwrap().unwrap().num_rows(), 1);
825 assert!(recordbatches.next().await.is_none());
826 assert_eq!(recordbatches.metrics().unwrap().elapsed_compute, 99);
827 }
828
829 #[tokio::test]
830 async fn test_record_batch_stream_exposes_error_after_pre_batch_metrics() {
831 let schema = test_schema();
832 let mut recordbatches = recordbatches_from_flight_message_stream(
833 "test-peer".to_string(),
834 stream::iter(vec![
835 Ok(FlightMessage::Schema(schema.arrow_schema().clone())),
836 Ok(FlightMessage::Metrics(test_metrics_json())),
837 Err(Error::from(Status::internal("boom after metrics"))),
838 ]),
839 )
840 .await
841 .unwrap();
842
843 let err = recordbatches.next().await.unwrap().unwrap_err();
844 assert_eq!("External error", err.to_string());
845 assert!(
846 format!("{err:?}").contains("boom after metrics"),
847 "unexpected error: {err:?}"
848 );
849 assert!(recordbatches.next().await.is_none());
850 }
851}