diff --git a/src/flow/src/batching_mode/frontend_client.rs b/src/flow/src/batching_mode/frontend_client.rs index 10ce3b6820..e932933ff8 100644 --- a/src/flow/src/batching_mode/frontend_client.rs +++ b/src/flow/src/batching_mode/frontend_client.rs @@ -183,13 +183,21 @@ impl DatabaseWithPeer { /// Try sending a "SELECT 1" to the database async fn try_select_one(&self) -> Result<(), Error> { // notice here use `sql` for `SELECT 1` return 1 row - let _ = self + let output = self .database .sql("SELECT 1") .await .with_context(|_| InvalidRequestSnafu { context: format!("Failed to handle `SELECT 1` request at {:?}", self.peer), })?; + + if let OutputData::Stream(stream) = output.data { + common_recordbatch::util::collect(stream) + .await + .map_err(BoxedError::new) + .context(ExternalSnafu)?; + } + Ok(()) } } @@ -698,9 +706,11 @@ mod tests { use api::v1::query_request::Query; use arrow_flight::flight_service_server::FlightServiceServer; use arrow_flight::{FlightData, Ticket}; + use common_grpc::flight::FlightEncoder; use common_query::{Output, OutputData}; use common_recordbatch::adapter::RecordBatchMetrics; use common_recordbatch::{OrderOption, RecordBatch, RecordBatchStream}; + use datatypes::arrow::datatypes::Schema as ArrowSchema; use datatypes::prelude::{ConcreteDataType, VectorRef}; use datatypes::schema::{ColumnSchema, Schema}; use datatypes::vectors::Int32Vector; @@ -775,6 +785,14 @@ mod tests { #[derive(Debug)] struct SlowFlight; + struct DelayedEofFlight { + schema_sent: Mutex>>, + release: Arc, + } + + #[derive(Debug)] + struct LateStreamErrorFlight; + struct WaitForConcurrentFlight { barrier: Arc, } @@ -870,6 +888,45 @@ mod tests { } } + #[async_trait::async_trait] + impl FlightCraft for DelayedEofFlight { + async fn do_get( + &self, + _request: TonicRequest, + ) -> std::result::Result>, Status> { + let schema = FlightEncoder::default().encode_schema(&ArrowSchema::empty()); + let schema_sent = self.schema_sent.lock().unwrap().take(); + let schema_stream = futures::stream::once(async move { + if let Some(schema_sent) = schema_sent { + let _ = schema_sent.send(()); + } + Ok(schema) + }); + let release = self.release.clone(); + let delayed_eof = futures::stream::unfold(release, |release| async move { + release.notified().await; + None::<(std::result::Result, _)> + }); + + Ok(TonicResponse::new(Box::pin( + schema_stream.chain(delayed_eof), + ))) + } + } + + #[async_trait::async_trait] + impl FlightCraft for LateStreamErrorFlight { + async fn do_get( + &self, + _request: TonicRequest, + ) -> std::result::Result>, Status> { + let schema = FlightEncoder::default().encode_schema(&ArrowSchema::empty()); + let stream = + futures::stream::iter([Ok(schema), Err(Status::unavailable("late stream error"))]); + Ok(TonicResponse::new(Box::pin(stream))) + } + } + #[async_trait::async_trait] impl FlightCraft for WaitForConcurrentFlight { async fn do_get( @@ -1056,6 +1113,71 @@ mod tests { assert!(format!("{err:?}").contains("Invalid value for flow.return_region_seq")); } + #[tokio::test] + async fn test_try_select_one_waits_for_stream_eof() { + let (schema_sent, schema_sent_rx) = tokio::sync::oneshot::channel(); + let release = Arc::new(tokio::sync::Notify::new()); + let (addr, server) = start_flight_server(DelayedEofFlight { + schema_sent: Mutex::new(Some(schema_sent)), + release: release.clone(), + }) + .await; + let database = Database::new( + DEFAULT_CATALOG_NAME, + DEFAULT_SCHEMA_NAME, + Client::with_urls([addr.as_str()]), + ); + let db = DatabaseWithPeer::new( + database, + Peer { + id: 1, + addr: addr.clone(), + }, + ); + let mut probe = tokio::spawn(async move { db.try_select_one().await }); + + timeout(Duration::from_secs(1), schema_sent_rx) + .await + .expect("server should send the schema") + .expect("schema signal should be sent"); + assert!( + timeout(Duration::from_millis(100), &mut probe) + .await + .is_err(), + "SELECT 1 must wait for the delayed stream tail and EOF" + ); + + release.notify_one(); + timeout(Duration::from_secs(1), &mut probe) + .await + .expect("SELECT 1 should complete after EOF") + .expect("probe task should not panic") + .expect("SELECT 1 should succeed after EOF"); + server.abort(); + } + + #[tokio::test] + async fn test_try_select_one_propagates_late_stream_error() { + let (addr, server) = start_flight_server(LateStreamErrorFlight).await; + let database = Database::new( + DEFAULT_CATALOG_NAME, + DEFAULT_SCHEMA_NAME, + Client::with_urls([addr.as_str()]), + ); + let db = DatabaseWithPeer::new( + database, + Peer { + id: 1, + addr: addr.clone(), + }, + ); + + let err = db.try_select_one().await.unwrap_err(); + server.abort(); + + assert!(format!("{err:?}").contains("late stream error")); + } + #[tokio::test] async fn test_check_all_frontends_without_auth_fails_fast_on_unauthenticated_frontend() { let (addr, server) = start_flight_server(RejectUnauthenticatedFlight).await;