Repository: apache/arrow-rs (Rust)

Issue:
FlightSqlClient `parameter_binding` is write-only
**Is your feature request related to a problem or challenge? Please describe what you are trying to do.**

I would like to run parameterized queries from the FlightSql client to integration test what users will be doing with the JDBC client. Unfortunately, I see that the `parameter_binding` field is written to, but never read-from - so these values must not be getting to the server where I need to implement support.

**Describe the solution you'd like**

1. Test with the JDBC client to see how these values are passed
2. Implement the necessary changes to the arrow-rs FlightSqlClient to reproduce the JDBC behavior

**Describe alternatives you've considered**

Continue to leave this important functionality unimplemented.
DoPut FlightSQL handler inadvertently consumes schema at start of Request<Streaming<FlightData>>
**Describe the bug**
When executing `DoPutPreparedStatementQuery` requests to bind parameters to a prepared statement, using the `DoPut` Flight RPC, one may receive an error like this:
```
rpc error: code = Internal desc = Received RecordBatch prior to Schema
```

This is because the implementation of `FlightService` for FlightSQL servers consumes the first flight batch by calling `.message()` on the `Request<Streaming<FlightData>>` handle:
```rust
async fn do_put(
    &self,
    mut request: Request<Streaming<FlightData>>,
) -> Result<Response<Self::DoPutStream>, Status> {
    let cmd = request.get_mut().message().await?.unwrap();
```
It does this to inspect the type of the `FlightDescriptor`, which tells it which `DoPut` handler to dispatch to in the `FlightSqlService`. If the schema was encoded in the very first message together with the `FlightDescriptor`, calling `.message()` consumes both the `FlightDescriptor` and the schema, so that a `FlightSqlService` handler such as `do_put_prepared_statement_query` no longer can read the schema from the `Request<Streaming<FlightData>>`. 

**To Reproduce**
Implement a `FlightSqlService` that handles any kind of `DoPut` request and decodes the flight stream, such as the following:
```rust
impl FlightSqlService for FlightSqlServer {
    async fn do_put_prepared_statement_query(
        &self,
        query: CommandPreparedStatementQuery,
        request: Request<Streaming<FlightData>>,
    ) -> Result<Response<<Self as FlightService>::DoPutStream>> {
        info!("do_put_prepared_statement_query");
        let parameters = FlightRecordBatchStream::new_from_flight_data(
            request.into_inner().map_err(|e| e.into()),
        )
        .try_collect::<Vec<_>>()
        .await?;

        todo!()
    }
}
```

then execute a `DoPut` using a raw `do_put` request on the `FlightSqlServiceClient`, sending the following batch:
```rust
let cmd = CommandPreparedStatementQuery {
    prepared_statement_handle: handle.clone(),
};

let flight_stream = FlightDataEncoderBuilder::new()
    .with_flight_descriptor(Some(FlightDescriptor::new_cmd(
        cmd.as_any().encode_to_vec(),
    )))
    .with_schema(batch.schema())
    .build(stream::once(async { Ok(batch) }));
client
    .do_put(flight_stream.map(Result::unwrap))
    .await
    .context("submit DoPut request")?
    .try_collect::<Vec<_>>()
    .await
    .context("read DoPut response")?;
```

**Expected behavior**
The server should be able to properly handle flight streams sent from a compliant FlightSQL client. 

**Additional context**
This bug was uncovered while implementing and testing prepared statements for a Rust FlightSQL server. The same bug was observed when submitting queries from both Rust and Go FlightSQL clients (the server is a Rust FlightSQL server). 



Fix (patch):
diff --git a/arrow-flight/examples/flight_sql_server.rs b/arrow-flight/examples/flight_sql_server.rs
index 1e99957390d8..d1aeae6f0a6c 100644
--- a/arrow-flight/examples/flight_sql_server.rs
+++ b/arrow-flight/examples/flight_sql_server.rs
@@ -15,6 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use arrow_flight::sql::server::PeekableFlightDataStream;
 use base64::prelude::BASE64_STANDARD;
 use base64::Engine;
 use futures::{stream, Stream, TryStreamExt};
@@ -602,7 +603,7 @@ impl FlightSqlService for FlightSqlServiceImpl {
     async fn do_put_statement_update(
         &self,
         _ticket: CommandStatementUpdate,
-        _request: Request<Streaming<FlightData>>,
+        _request: Request<PeekableFlightDataStream>,
     ) -> Result<i64, Status> {
         Ok(FAKE_UPDATE_RESULT)
     }
@@ -610,7 +611,7 @@ impl FlightSqlService for FlightSqlServiceImpl {
     async fn do_put_substrait_plan(
         &self,
         _ticket: CommandStatementSubstraitPlan,
-        _request: Request<Streaming<FlightData>>,
+        _request: Request<PeekableFlightDataStream>,
     ) -> Result<i64, Status> {
         Err(Status::unimplemented(
             "do_put_substrait_plan not implemented",
@@ -620,7 +621,7 @@ impl FlightSqlService for FlightSqlServiceImpl {
     async fn do_put_prepared_statement_query(
         &self,
         _query: CommandPreparedStatementQuery,
-        _request: Request<Streaming<FlightData>>,
+        _request: Request<PeekableFlightDataStream>,
     ) -> Result<Response<<Self as FlightService>::DoPutStream>, Status> {
         Err(Status::unimplemented(
             "do_put_prepared_statement_query not implemented",
@@ -630,7 +631,7 @@ impl FlightSqlService for FlightSqlServiceImpl {
     async fn do_put_prepared_statement_update(
         &self,
         _query: CommandPreparedStatementUpdate,
-        _request: Request<Streaming<FlightData>>,
+        _request: Request<PeekableFlightDataStream>,
     ) -> Result<i64, Status> {
         Err(Status::unimplemented(
             "do_put_prepared_statement_update not implemented",
diff --git a/arrow-flight/src/bin/flight_sql_client.rs b/arrow-flight/src/bin/flight_sql_client.rs
index 20c8062f899e..d7b02414c5cc 100644
--- a/arrow-flight/src/bin/flight_sql_client.rs
+++ b/arrow-flight/src/bin/flight_sql_client.rs
@@ -15,15 +15,16 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use std::{sync::Arc, time::Duration};
+use std::{error::Error, sync::Arc, time::Duration};
 
-use arrow_array::RecordBatch;
-use arrow_cast::pretty::pretty_format_batches;
+use arrow_array::{ArrayRef, Datum, RecordBatch, StringArray};
+use arrow_cast::{cast_with_options, pretty::pretty_format_batches, CastOptions};
 use arrow_flight::{
     sql::client::FlightSqlServiceClient, utils::flight_data_to_batches, FlightData,
+    FlightInfo,
 };
 use arrow_schema::{ArrowError, Schema};
-use clap::Parser;
+use clap::{Parser, Subcommand};
 use futures::TryStreamExt;
 use tonic::transport::{Channel, ClientTlsConfig, Endpoint};
 use tracing_log::log::info;
@@ -98,8 +99,20 @@ struct Args {
     #[clap(flatten)]
     client_args: ClientArgs,
 
-    /// SQL query.
-    query: String,
+    #[clap(subcommand)]
+    cmd: Command,
+}
+
+#[derive(Debug, Subcommand)]
+enum Command {
+    StatementQuery {
+        query: String,
+    },
+    PreparedStatementQuery {
+        query: String,
+        #[clap(short, value_parser = parse_key_val)]
+        params: Vec<(String, String)>,
+    },
 }
 
 #[tokio::main]
@@ -108,12 +121,50 @@ async fn main() {
     setup_logging();
     let mut client = setup_client(args.client_args).await.expect("setup client");
 
-    let info = client
-        .execute(args.query, None)
+    let flight_info = match args.cmd {
+        Command::StatementQuery { query } => client
+            .execute(query, None)
+            .await
+            .expect("execute statement"),
+        Command::PreparedStatementQuery { query, params } => {
+            let mut prepared_stmt = client
+                .prepare(query, None)
+                .await
+                .expect("prepare statement");
+
+            if !params.is_empty() {
+                prepared_stmt
+                    .set_parameters(
+                        construct_record_batch_from_params(
+                            &params,
+                            prepared_stmt
+                                .parameter_schema()
+                                .expect("get parameter schema"),
+                        )
+                        .expect("construct parameters"),
+                    )
+                    .expect("bind parameters")
+            }
+
+            prepared_stmt
+                .execute()
+                .await
+                .expect("execute prepared statement")
+        }
+    };
+
+    let batches = execute_flight(&mut client, flight_info)
         .await
-        .expect("prepare statement");
-    info!("got flight info");
+        .expect("read flight data");
 
+    let res = pretty_format_batches(batches.as_slice()).expect("format results");
+    println!("{res}");
+}
+
+async fn execute_flight(
+    client: &mut FlightSqlServiceClient<Channel>,
+    info: FlightInfo,
+) -> Result<Vec<RecordBatch>, ArrowError> {
     let schema = Arc::new(Schema::try_from(info.clone()).expect("valid schema"));
     let mut batches = Vec::with_capacity(info.endpoint.len() + 1);
     batches.push(RecordBatch::new_empty(schema));
@@ -134,8 +185,27 @@ async fn main() {
     }
     info!("received data");
 
-    let res = pretty_format_batches(batches.as_slice()).expect("format results");
-    println!("{res}");
+    Ok(batches)
+}
+
+fn construct_record_batch_from_params(
+    params: &[(String, String)],
+    parameter_schema: &Schema,
+) -> Result<RecordBatch, ArrowError> {
+    let mut items = Vec::<(&String, ArrayRef)>::new();
+
+    for (name, value) in params {
+        let field = parameter_schema.field_with_name(name)?;
+        let value_as_array = StringArray::new_scalar(value);
+        let casted = cast_with_options(
+            value_as_array.get().0,
+            field.data_type(),
+            &CastOptions::default(),
+        )?;
+        items.push((name, casted))
+    }
+
+    RecordBatch::try_from_iter(items)
 }
 
 fn setup_logging() {
@@ -203,3 +273,13 @@ async fn setup_client(
 
     Ok(client)
 }
+
+/// Parse a single key-value pair
+fn parse_key_val(
+    s: &str,
+) -> Result<(String, String), Box<dyn Error + Send + Sync + 'static>> {
+    let pos = s
+        .find('=')
+        .ok_or_else(|| format!("invalid KEY=value: no `=` found in `{s}`"))?;
+    Ok((s[..pos].parse()?, s[pos + 1..].parse()?))
+}
diff --git a/arrow-flight/src/sql/client.rs b/arrow-flight/src/sql/client.rs
index 4b1f38ebcbb7..2d382cf2ca20 100644
--- a/arrow-flight/src/sql/client.rs
+++ b/arrow-flight/src/sql/client.rs
@@ -24,6 +24,8 @@ use std::collections::HashMap;
 use std::str::FromStr;
 use tonic::metadata::AsciiMetadataKey;
 
+use crate::encode::FlightDataEncoderBuilder;
+use crate::error::FlightError;
 use crate::flight_service_client::FlightServiceClient;
 use crate::sql::server::{CLOSE_PREPARED_STATEMENT, CREATE_PREPARED_STATEMENT};
 use crate::sql::{
@@ -32,8 +34,8 @@ use crate::sql::{
     CommandGetCrossReference, CommandGetDbSchemas, CommandGetExportedKeys,
     CommandGetImportedKeys, CommandGetPrimaryKeys, CommandGetSqlInfo,
     CommandGetTableTypes, CommandGetTables, CommandGetXdbcTypeInfo,
-    CommandPreparedStatementQuery, CommandStatementQuery, CommandStatementUpdate,
-    DoPutUpdateResult, ProstMessageExt, SqlInfo,
+    CommandPreparedStatementQuery, CommandPreparedStatementUpdate, CommandStatementQuery,
+    CommandStatementUpdate, DoPutUpdateResult, ProstMessageExt, SqlInfo,
 };
 use crate::{
     Action, FlightData, FlightDescriptor, FlightInfo, HandshakeRequest,
@@ -439,9 +441,12 @@ impl PreparedStatement<Channel> {
 
     /// Executes the prepared statement query on the server.
     pub async fn execute(&mut self) -> Result<FlightInfo, ArrowError> {
+        self.write_bind_params().await?;
+
         let cmd = CommandPreparedStatementQuery {
             prepared_statement_handle: self.handle.clone(),
         };
+
         let result = self
             .flight_sql_client
             .get_flight_info_for_command(cmd)
@@ -451,7 +456,9 @@ impl PreparedStatement<Channel> {
 
     /// Executes the prepared statement update query on the server.
     pub async fn execute_update(&mut self) -> Result<i64, ArrowError> {
-        let cmd = CommandPreparedStatementQuery {
+        self.write_bind_params().await?;
+
+        let cmd = CommandPreparedStatementUpdate {
             prepared_statement_handle: self.handle.clone(),
         };
         let descriptor = FlightDescriptor::new_cmd(cmd.as_any().encode_to_vec());
@@ -492,6 +499,36 @@ impl PreparedStatement<Channel> {
         Ok(())
     }
 
+    /// Submit parameters to the server, if any have been set on this prepared statement instance
+    async fn write_bind_params(&mut self) -> Result<(), ArrowError> {
+        if let Some(ref params_batch) = self.parameter_binding {
+            let cmd = CommandPreparedStatementQuery {
+                prepared_statement_handle: self.handle.clone(),
+            };
+
+            let descriptor = FlightDescriptor::new_cmd(cmd.as_any().encode_to_vec());
+            let flight_stream_builder = FlightDataEncoderBuilder::new()
+                .with_flight_descriptor(Some(descriptor))
+                .with_schema(params_batch.schema());
+            let flight_data = flight_stream_builder
+                .build(futures::stream::iter(
+                    self.parameter_binding.clone().map(Ok),
+                ))
+                .try_collect::<Vec<_>>()
+                .await
+                .map_err(flight_error_to_arrow_error)?;
+
+            self.flight_sql_client
+                .do_put(stream::iter(flight_data))
+                .await?
+                .try_collect::<Vec<_>>()
+                .await
+                .map_err(status_to_arrow_error)?;
+        }
+
+        Ok(())
+    }
+
     /// Close the prepared statement, so that this PreparedStatement can not used
     /// anymore and server can free up any resources.
     pub async fn close(mut self) -> Result<(), ArrowError> {
@@ -515,6 +552,13 @@ fn status_to_arrow_error(status: tonic::Status) -> ArrowError {
     ArrowError::IpcError(format!("{status:?}"))
 }
 
+fn flight_error_to_arrow_error(err: FlightError) -> ArrowError {
+    match err {
+        FlightError::Arrow(e) => e,
+        e => ArrowError::ExternalError(Box::new(e)),
+    }
+}
+
 // A polymorphic structure to natively represent different types of data contained in `FlightData`
 pub enum ArrowFlightData {
     RecordBatch(RecordBatch),
diff --git a/arrow-flight/src/sql/server.rs b/arrow-flight/src/sql/server.rs
index 102d97105a2e..a158ed77f54d 100644
--- a/arrow-flight/src/sql/server.rs
+++ b/arrow-flight/src/sql/server.rs
@@ -19,7 +19,7 @@
 
 use std::pin::Pin;
 
-use futures::Stream;
+use futures::{stream::Peekable, Stream, StreamExt};
 use prost::Message;
 use tonic::{Request, Response, Status, Streaming};
 
@@ -366,7 +366,7 @@ pub trait FlightSqlService: Sync + Send + Sized + 'static {
     /// Implementors may override to handle additional calls to do_put()
     async fn do_put_fallback(
         &self,
-        _request: Request<Streaming<FlightData>>,
+        _request: Request<PeekableFlightDataStream>,
         message: Any,
     ) -> Result<Response<<Self as FlightService>::DoPutStream>, Status> {
         Err(Status::unimplemented(format!(
@@ -379,7 +379,7 @@ pub trait FlightSqlService: Sync + Send + Sized + 'static {
     async fn do_put_statement_update(
         &self,
         _ticket: CommandStatementUpdate,
-        _request: Request<Streaming<FlightData>>,
+        _request: Request<PeekableFlightDataStream>,
     ) -> Result<i64, Status> {
         Err(Status::unimplemented(
             "do_put_statement_update has no default implementation",
@@ -390,7 +390,7 @@ pub trait FlightSqlService: Sync + Send + Sized + 'static {
     async fn do_put_prepared_statement_query(
         &self,
         _query: CommandPreparedStatementQuery,
-        _request: Request<Streaming<FlightData>>,
+        _request: Request<PeekableFlightDataStream>,
     ) -> Result<Response<<Self as FlightService>::DoPutStream>, Status> {
         Err(Status::unimplemented(
             "do_put_prepared_statement_query has no default implementation",
@@ -401,7 +401,7 @@ pub trait FlightSqlService: Sync + Send + Sized + 'static {
     async fn do_put_prepared_statement_update(
         &self,
         _query: CommandPreparedStatementUpdate,
-        _request: Request<Streaming<FlightData>>,
+        _request: Request<PeekableFlightDataStream>,
     ) -> Result<i64, Status> {
         Err(Status::unimplemented(
             "do_put_prepared_statement_update has no default implementation",
@@ -412,7 +412,7 @@ pub trait FlightSqlService: Sync + Send + Sized + 'static {
     async fn do_put_substrait_plan(
         &self,
         _query: CommandStatementSubstraitPlan,
-        _request: Request<Streaming<FlightData>>,
+        _request: Request<PeekableFlightDataStream>,
     ) -> Result<i64, Status> {
         Err(Status::unimplemented(
             "do_put_substrait_plan has no default implementation",
@@ -688,9 +688,17 @@ where
 
     async fn do_put(
         &self,
-        mut request: Request<Streaming<FlightData>>,
+        request: Request<Streaming<FlightData>>,
     ) -> Result<Response<Self::DoPutStream>, Status> {
-        let cmd = request.get_mut().message().await?.unwrap();
+        // See issue #4658: https://github.com/apache/arrow-rs/issues/4658
+        // To dispatch to the correct `do_put` method, we cannot discard the first message,
+        // as it may contain the Arrow schema, which the `do_put` handler may need.
+        // To allow the first message to be reused by the `do_put` handler,
+        // 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 message = Any::decode(&*cmd.flight_descriptor.unwrap().cmd)
             .map_err(decode_error_to_status)?;
         match Command::try_from(message).map_err(arrow_error_to_status)? {
@@ -957,3 +965,89 @@ fn decode_error_to_status(err: prost::DecodeError) -> Status {
 fn arrow_error_to_status(err: arrow_schema::ArrowError) -> Status {
     Status::internal(format!("{err:?}"))
 }
+
+/// A wrapper around [`Streaming<FlightData>`] that allows "peeking" at the
+/// message at the front of the stream without consuming it.
+/// This is needed because sometimes the first message in the stream will contain
+/// a [`FlightDescriptor`] in addition to potentially any data, and the dispatch logic
+/// must inspect this information.
+///
+/// # Example
+///
+/// [`PeekableFlightDataStream::peek`] can be used to peek at the first message without
+/// discarding it; otherwise, `PeekableFlightDataStream` can be used as a regular stream.
+/// See the following example:
+///
+/// ```no_run
+/// use arrow_array::RecordBatch;
+/// use arrow_flight::decode::FlightRecordBatchStream;
+/// use arrow_flight::FlightDescriptor;
+/// use arrow_flight::error::FlightError;
+/// use arrow_flight::sql::server::PeekableFlightDataStream;
+/// use tonic::{Request, Status};
+/// use futures::TryStreamExt;
+///
+/// #[tokio::main]
+/// async fn main() -> Result<(), Status> {
+///     let request: Request<PeekableFlightDataStream> = todo!();
+///     let stream: PeekableFlightDataStream = request.into_inner();
+///
+///     // The first message contains the flight descriptor and the schema.
+///     // Read the flight descriptor without discarding the schema:
+///     let flight_descriptor: FlightDescriptor = stream
+///         .peek()
+///         .await
+///         .cloned()
+///         .transpose()?
+///         .and_then(|data| data.flight_descriptor)
+///         .expect("first message should contain flight descriptor");
+///
+///     // Pass the stream through a decoder
+///     let batches: Vec<RecordBatch> = FlightRecordBatchStream::new_from_flight_data(
+///         request.into_inner().map_err(|e| e.into()),
+///     )
+///     .try_collect()
+///     .await?;
+/// }
+/// ```
+pub struct PeekableFlightDataStream {
+    inner: Peekable<Streaming<FlightData>>,
+}
+
+impl PeekableFlightDataStream {
+    fn new(stream: Streaming<FlightData>) -> Self {
+        Self {
+            inner: stream.peekable(),
+        }
+    }
+
+    /// Convert this stream into a `Streaming<FlightData>`.
+    /// Any messages observed through [`Self::peek`] will be lost
+    /// after the conversion.
+    pub fn into_inner(self) -> Streaming<FlightData> {
+        self.inner.into_inner()
+    }
+
+    /// Convert this stream into a `Peekable<Streaming<FlightData>>`.
+    /// Preserves the state of the stream, so that calls to [`Self::peek`]
+    /// and [`Self::poll_next`] are the same.
+    pub fn into_peekable(self) -> Peekable<Streaming<FlightData>> {
+        self.inner
+    }
+
+    /// Peek at the head of this stream without advancing it.
+    pub async fn peek(&mut self) -> Option<&Result<FlightData, Status>> {
+        Pin::new(&mut self.inner).peek().await
+    }
+}
+
+impl Stream for PeekableFlightDataStream {
+    type Item = Result<FlightData, Status>;
+
+    fn poll_next(
+        mut self: Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<Option<Self::Item>> {
+        self.inner.poll_next_unpin(cx)
+    }
+}