You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by al...@apache.org on 2025/05/26 10:19:32 UTC
(arrow-rs) branch main updated: Prevent FlightSQL server panics for `do_put` when stream is empty or 1st stream element is an Err (#7492)
This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new 88d59f8492 Prevent FlightSQL server panics for `do_put` when stream is empty or 1st stream element is an Err (#7492)
88d59f8492 is described below
commit 88d59f84925fe6e8dc8dbbe2596235403edb5248
Author: superserious-dev <gi...@superserious.dev>
AuthorDate: Mon May 26 03:19:27 2025 -0700
Prevent FlightSQL server panics for `do_put` when stream is empty or 1st stream element is an Err (#7492)
* Fix bug and add tests to verify
* Remove unnecessary let assignment.
* Add error callback for `do_put`
* cargo fmt
---------
Co-authored-by: Andrew Lamb <an...@nerdnetworks.org>
---
arrow-flight/src/sql/mod.rs | 2 +
arrow-flight/src/sql/server.rs | 54 ++++++++++++++++---
arrow-flight/src/streams.rs | 5 +-
arrow-flight/tests/flight_sql_client.rs | 93 +++++++++++++++++++++++++++++++--
4 files changed, 141 insertions(+), 13 deletions(-)
diff --git a/arrow-flight/src/sql/mod.rs b/arrow-flight/src/sql/mod.rs
index f11b6f5bd9..955f1904a6 100644
--- a/arrow-flight/src/sql/mod.rs
+++ b/arrow-flight/src/sql/mod.rs
@@ -113,6 +113,8 @@ pub mod client;
pub mod metadata;
pub mod server;
+pub use crate::streams::FallibleRequestStream;
+
/// ProstMessageExt are useful utility methods for prost::Message types
pub trait ProstMessageExt: prost::Message + Default {
/// type_url for this Message
diff --git a/arrow-flight/src/sql/server.rs b/arrow-flight/src/sql/server.rs
index 8ab8a16dbb..add7c8db40 100644
--- a/arrow-flight/src/sql/server.rs
+++ b/arrow-flight/src/sql/server.rs
@@ -17,12 +17,9 @@
//! Helper trait [`FlightSqlService`] for implementing a [`FlightService`] that implements FlightSQL.
+use std::fmt::{Display, Formatter};
use std::pin::Pin;
-use futures::{stream::Peekable, Stream, StreamExt};
-use prost::Message;
-use tonic::{Request, Response, Status, Streaming};
-
use super::{
ActionBeginSavepointRequest, ActionBeginSavepointResult, ActionBeginTransactionRequest,
ActionBeginTransactionResult, ActionCancelQueryRequest, ActionCancelQueryResult,
@@ -41,6 +38,9 @@ use crate::{
FlightData, FlightDescriptor, FlightInfo, HandshakeRequest, HandshakeResponse, PutResult,
SchemaResult, Ticket,
};
+use futures::{stream::Peekable, Stream, StreamExt};
+use prost::Message;
+use tonic::{Request, Response, Status, Streaming};
pub(crate) static CREATE_PREPARED_STATEMENT: &str = "CreatePreparedStatement";
pub(crate) static CLOSE_PREPARED_STATEMENT: &str = "ClosePreparedStatement";
@@ -386,6 +386,15 @@ pub trait FlightSqlService: Sync + Send + Sized + 'static {
)))
}
+ /// Implementors may override to handle do_put errors
+ async fn do_put_error_callback(
+ &self,
+ _request: Request<PeekableFlightDataStream>,
+ error: DoPutError,
+ ) -> Result<Response<<Self as FlightService>::DoPutStream>, Status> {
+ Err(Status::unimplemented(format!("Unhandled Error: {}", error)))
+ }
+
/// Execute an update SQL statement.
async fn do_put_statement_update(
&self,
@@ -710,10 +719,21 @@ where
// we wrap this stream in a `Peekable` one, which allows us to peek at
// the first message without discarding it.
let mut request = request.map(PeekableFlightDataStream::new);
- let cmd = Pin::new(request.get_mut()).peek().await.unwrap().clone()?;
+ let mut stream = Pin::new(request.get_mut());
+
+ let peeked_item = stream.peek().await.cloned();
+ let Some(cmd) = peeked_item else {
+ return self
+ .do_put_error_callback(request, DoPutError::MissingCommand)
+ .await;
+ };
- let message =
- Any::decode(&*cmd.flight_descriptor.unwrap().cmd).map_err(decode_error_to_status)?;
+ let Some(flight_descriptor) = cmd?.flight_descriptor else {
+ return self
+ .do_put_error_callback(request, DoPutError::MissingFlightDescriptor)
+ .await;
+ };
+ let message = Any::decode(flight_descriptor.cmd).map_err(decode_error_to_status)?;
match Command::try_from(message).map_err(arrow_error_to_status)? {
Command::CommandStatementUpdate(command) => {
let record_count = self.do_put_statement_update(command, request).await?;
@@ -968,6 +988,26 @@ where
}
}
+/// Unrecoverable errors associated with `do_put` requests
+pub enum DoPutError {
+ /// The first element in the request stream is missing the command
+ MissingCommand,
+ /// The first element in the request stream is missing the flight descriptor
+ MissingFlightDescriptor,
+}
+impl Display for DoPutError {
+ fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+ match self {
+ DoPutError::MissingCommand => {
+ write!(f, "Command is missing.")
+ }
+ DoPutError::MissingFlightDescriptor => {
+ write!(f, "Flight descriptor is missing.")
+ }
+ }
+ }
+}
+
fn decode_error_to_status(err: prost::DecodeError) -> Status {
Status::invalid_argument(format!("{err:?}"))
}
diff --git a/arrow-flight/src/streams.rs b/arrow-flight/src/streams.rs
index ab49612201..0cd3aa41a5 100644
--- a/arrow-flight/src/streams.rs
+++ b/arrow-flight/src/streams.rs
@@ -32,7 +32,7 @@ use std::task::{ready, Poll};
///
/// This can be used to accept a stream of `Result<_>` from a client API and send
/// them to the remote server that wants only the successful results.
-pub(crate) struct FallibleRequestStream<T, E> {
+pub struct FallibleRequestStream<T, E> {
/// sender to notify error
sender: Option<Sender<E>>,
/// fallible stream
@@ -40,7 +40,8 @@ pub(crate) struct FallibleRequestStream<T, E> {
}
impl<T, E> FallibleRequestStream<T, E> {
- pub(crate) fn new(
+ /// Create a FallibleRequestStream
+ pub fn new(
sender: Sender<E>,
fallible_stream: Pin<Box<dyn Stream<Item = std::result::Result<T, E>> + Send + 'static>>,
) -> Self {
diff --git a/arrow-flight/tests/flight_sql_client.rs b/arrow-flight/tests/flight_sql_client.rs
index 349da062a8..f3b7114dba 100644
--- a/arrow-flight/tests/flight_sql_client.rs
+++ b/arrow-flight/tests/flight_sql_client.rs
@@ -22,21 +22,23 @@ use crate::common::utils::make_primitive_batch;
use arrow_array::RecordBatch;
use arrow_flight::decode::FlightRecordBatchStream;
+use arrow_flight::encode::FlightDataEncoderBuilder;
use arrow_flight::error::FlightError;
use arrow_flight::flight_service_server::FlightServiceServer;
use arrow_flight::sql::client::FlightSqlServiceClient;
use arrow_flight::sql::server::{FlightSqlService, PeekableFlightDataStream};
use arrow_flight::sql::{
ActionBeginTransactionRequest, ActionBeginTransactionResult, ActionEndTransactionRequest,
- CommandStatementIngest, EndTransaction, SqlInfo, TableDefinitionOptions, TableExistsOption,
- TableNotExistOption,
+ CommandStatementIngest, EndTransaction, FallibleRequestStream, ProstMessageExt, SqlInfo,
+ TableDefinitionOptions, TableExistsOption, TableNotExistOption,
};
-use arrow_flight::Action;
+use arrow_flight::{Action, FlightData, FlightDescriptor};
use futures::{StreamExt, TryStreamExt};
+use prost::Message;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::Mutex;
-use tonic::{Request, Status};
+use tonic::{IntoStreamingRequest, Request, Status};
use uuid::Uuid;
#[tokio::test]
@@ -116,6 +118,89 @@ pub async fn test_execute_ingest_error() {
);
}
+#[tokio::test]
+pub async fn test_do_put_empty_stream() {
+ // Test for https://github.com/apache/arrow-rs/issues/7329
+
+ let test_server = FlightSqlServiceImpl::new();
+ let fixture = TestFixture::new(test_server.service()).await;
+ let channel = fixture.channel().await;
+ let mut flight_sql_client = FlightSqlServiceClient::new(channel);
+ let cmd = make_ingest_command();
+
+ // Create an empty request stream
+ let input_data = futures::stream::iter(vec![]);
+ let flight_descriptor = FlightDescriptor::new_cmd(cmd.as_any().encode_to_vec());
+ let flight_data_encoder = FlightDataEncoderBuilder::default()
+ .with_flight_descriptor(Some(flight_descriptor))
+ .build(input_data);
+ let flight_data: Vec<FlightData> = Box::pin(flight_data_encoder).try_collect().await.unwrap();
+ let request_stream = futures::stream::iter(flight_data);
+
+ // Execute a `do_put` and verify that the server error contains the expected message
+ let err = flight_sql_client.do_put(request_stream).await.unwrap_err();
+ assert!(err
+ .to_string()
+ .contains("Unhandled Error: Command is missing."),);
+}
+
+#[tokio::test]
+pub async fn test_do_put_first_element_err() {
+ // Test for https://github.com/apache/arrow-rs/issues/7329
+
+ let test_server = FlightSqlServiceImpl::new();
+ let fixture = TestFixture::new(test_server.service()).await;
+ let channel = fixture.channel().await;
+ let mut flight_sql_client = FlightSqlServiceClient::new(channel);
+ let cmd = make_ingest_command();
+
+ let (sender, _receiver) = futures::channel::oneshot::channel();
+
+ // Create a fallible request stream such that the 1st element is a FlightError
+ let input_data = futures::stream::iter(vec![
+ Err(FlightError::NotYetImplemented("random error".to_string())),
+ Ok(make_primitive_batch(5)),
+ ]);
+ let flight_descriptor = FlightDescriptor::new_cmd(cmd.as_any().encode_to_vec());
+ let flight_data_encoder = FlightDataEncoderBuilder::default()
+ .with_flight_descriptor(Some(flight_descriptor))
+ .build(input_data);
+ let flight_data: FallibleRequestStream<FlightData, FlightError> =
+ FallibleRequestStream::new(sender, Box::pin(flight_data_encoder));
+ let request_stream = flight_data.into_streaming_request();
+
+ // Execute a `do_put` and verify that the server error contains the expected message
+ let err = flight_sql_client.do_put(request_stream).await.unwrap_err();
+
+ assert!(err
+ .to_string()
+ .contains("Unhandled Error: Command is missing."),);
+}
+
+#[tokio::test]
+pub async fn test_do_put_missing_flight_descriptor() {
+ // Test for https://github.com/apache/arrow-rs/issues/7329
+
+ let test_server = FlightSqlServiceImpl::new();
+ let fixture = TestFixture::new(test_server.service()).await;
+ let channel = fixture.channel().await;
+ let mut flight_sql_client = FlightSqlServiceClient::new(channel);
+
+ // Create a request stream such that the flight descriptor is missing
+ let stream = futures::stream::iter(vec![Ok(make_primitive_batch(5))]);
+ let flight_data_encoder = FlightDataEncoderBuilder::default()
+ .with_flight_descriptor(None)
+ .build(stream);
+ let flight_data: Vec<FlightData> = Box::pin(flight_data_encoder).try_collect().await.unwrap();
+ let request_stream = futures::stream::iter(flight_data);
+
+ // Execute a `do_put` and verify that the server error contains the expected message
+ let err = flight_sql_client.do_put(request_stream).await.unwrap_err();
+ assert!(err
+ .to_string()
+ .contains("Unhandled Error: Flight descriptor is missing."),);
+}
+
fn make_ingest_command() -> CommandStatementIngest {
CommandStatementIngest {
table_definition_options: Some(TableDefinitionOptions {