Fix flight sql do put handling, add bind parameter support to FlightSQL cli client #4797

suremarc · 2023-09-07T21:04:59Z

Which issue does this PR close?

Closes #4658 and closes #3598

Rationale for this change

DoPut requests with nonempty flight streams cannot be handled properly by the Rust FlightSQL server implementation in its current state.

What changes are included in this PR?

We change all DoPut methods on the FlightSqlService trait to accept a Peekable<Streaming<FlightData>> instead of a regular Streaming<FlightData>. We also finished the parameter binding functionality in the client, in order to test prepared statements properly.

Are there any user-facing changes?

Yes, there is unfortunately an API change for the FlightSqlService trait. I am open to alternatives, as it is probably possible to do evil things with Peekable, but I do not think it is possible to fix this without a breaking change.

suremarc · 2023-09-07T21:06:25Z

arrow-flight/src/sql/client.rs

+fn flight_error_to_arrow_error(err: FlightError) -> ArrowError {
+    match err {
+        FlightError::Arrow(e) => e,
+        FlightError::NotYetImplemented(s) => ArrowError::NotYetImplemented(s),
+        FlightError::Tonic(status) => status_to_arrow_error(status),
+        FlightError::ProtocolError(e) => ArrowError::IpcError(e),
+        FlightError::DecodeError(s) => ArrowError::IpcError(s),
+        FlightError::ExternalError(e) => ArrowError::ExternalError(e),
+    }
+}


I am not sure why we return ArrowError from the Flight client (instead of FlightError), but I am trying to keep this PR scoped, so I just decided to stay consistent

I agree this should probably be FlightError

tustvold

This looks good to me, however, I'm not very well versed in FlightSQL and so have asked some potentially stupid questions

tustvold · 2023-09-08T12:18:06Z

arrow-flight/tests/flight_sql_client_cli.rs

+const PREPARED_QUERY: &str = "SELECT * FROM table WHERE field = $1";
+const PREPARED_STATEMENT_HANDLE: &str = "prepared_statement_handle";
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 1)]


Do we need multi thread here?

I don't really think so. I just followed what the existing test did, but I don't really know see why it should. I will try and see if it works without it

Seems to work fine without it, so I just removed it

tustvold · 2023-09-08T12:18:34Z

arrow-flight/src/sql/server.rs

+        // 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(futures::StreamExt::peekable);


So if I am following correctly, the issue is do_put accepts a FlightData stream, but the first request in the stream will contain a FlightDescriptor in addition to potentially any data. I continue to be utterly baffled by the design of Flight 😅

Yes, that is my understanding as well. Prior to this change, decoding a flight stream inside one of the do_put methods would give you an error like Received RecordBatch prior to schema

arrow-flight/src/sql/client.rs

tustvold · 2023-09-08T12:22:02Z

arrow-flight/src/sql/client.rs

+fn flight_error_to_arrow_error(err: FlightError) -> ArrowError {
+    match err {
+        FlightError::Arrow(e) => e,
+        FlightError::NotYetImplemented(s) => ArrowError::NotYetImplemented(s),
+        FlightError::Tonic(status) => status_to_arrow_error(status),
+        FlightError::ProtocolError(e) => ArrowError::IpcError(e),
+        FlightError::DecodeError(s) => ArrowError::IpcError(s),
+        FlightError::ExternalError(e) => ArrowError::ExternalError(e),
+    }
+}


I agree this should probably be FlightError

tustvold · 2023-09-08T12:28:56Z

arrow-flight/src/sql/client.rs

+                .await
+                .map_err(flight_error_to_arrow_error)?;
+
+            self.flight_sql_client


This appears consistent with the FlightSQL specification, it uses do_put to bind the parameter arguments. What isn't clear to me is if the result should be being used in some way.

This would seem to imply some sort of server-side state which I had perhaps expected FlightSQL to not rely on

Yeah, I think we are in agreement about it implying server-side state. FWIW FlightSQL also supports transactions which I think (maybe wrongly) would also require state. There was also some discussion happening about adding new RPC's for managing session state at some point (like a close RPC or something)

This seems like a fundamental flaw in FlightSQL tbh, gRPC is not a connection-oriented protocol and so the lifetime of any server state is non-deterministic... I believe @alamb plans to start a discussion to see if we can't fix this

I filed apache/arrow#37720 and will circulate this around

tustvold · 2023-09-08T12:29:31Z

arrow-flight/src/sql/server.rs

@@ -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<Peekable<Streaming<FlightData>>>,


Another option might be to pass the first ticket request as a separate argument. I don't feel strongly either way

It's a tough decision for me. I prefer using Peekable as it can be used as if it were the original stream, but I hate the fact that we have to leak its usage. We could pass the first FlightData as a separate argument, but it would require the user to chain it with the stream, if they wanted to use any APIs expecting a stream of FlightData. So I think I would stick with Peekable in the absence of any preference from others.

This server is the scaffolding to help people build flightsql servers -- they can always use the raw FlightService if they prefer. Thus I think the change in API is less critical and given the requirements it seems inevitable.

I think the only thing we should try and improve in this PR is improving the documentation to explain why peekable is used somehow (to lower the cognative burden on people trying to use this).

One potential option to document this is rather than using Peekable<Streaming<...>> dirctly, would make our own wrapper, PeekableStreaming or something. While this would require duplicate a bunch of the peekable API, there would be a natural place to document what it was for and how to use it which I think would lower the barrier to usage.

For example:

/// A wrapper around `Streaming` that allows inspection of the first message. /// This is needed because sometimes the first request in the stream will contain /// a [`FlightDescriptor`] in addition to potentially any data and the dispatch logic /// must inspect this information. /// /// # example: /// <show an example here of calling `into_inner()` to get the original data back struct PeekableStreaming { inner: Peekable<Streaming<FlightData>> } impl PeekableStreaming { /// return the inner stream pub fn into_inner(self) -> Streaming<FlightData> { self.inner.into_inner() } ... }

We could also potentially use something like BoxStream<FlightData> but that would lose the gRPC specific stuff like status codes and trailers exposed by Streaming as well as being an API change as well.

Thus I think this design is the best of several less than ideal solutions. To proceed perhaps we can add some documentation on the do_*_fallback methods that mentions the stream comes from peekable

I added a new type, PeekableFlightDataStream, which exposes into_inner and peek, similarly to Peekable. I think this is a good enough subset of functionality for FlightSQL use cases, and if users need access to more of the lower-level functionality of Peekable, they can call PeekableFlightDataStream::into_peekable.

Co-authored-by: Raphael Taylor-Davies <1781103+tustvold@users.noreply.github.com>

suremarc · 2023-09-08T15:43:40Z

arrow-flight/src/sql/client.rs

-        let cmd = CommandPreparedStatementQuery {
+        self.write_bind_params().await?;
+
+        let cmd = CommandPreparedStatementUpdate {


I also forgot to mention, I think this was a bug in the existing implementation. ExecuteUpdate should be performed with a CommandPreparedStatementUpdate command, not a CommandPreparedStatementQuery.

I agree - update should use CommandPreparedStatementUpdate

https://github.com/apache/arrow/blob/15a8ac3ce4e3ac31f9f361770ad4a38c69102aa1/format/FlightSql.proto#L1769-L1780

alamb

Thank you @suremarc and @tustvold -- I think this PR is pretty close to ready. I left some comments: I think we need to document / explain the peekable use better (I left some suggestions) and get CI passing but then we should be ready to go

alamb · 2023-09-14T09:29:40Z

arrow-flight/src/sql/client.rs

-        let cmd = CommandPreparedStatementQuery {
+        self.write_bind_params().await?;
+
+        let cmd = CommandPreparedStatementUpdate {


I agree - update should use CommandPreparedStatementUpdate

https://github.com/apache/arrow/blob/15a8ac3ce4e3ac31f9f361770ad4a38c69102aa1/format/FlightSql.proto#L1769-L1780

alamb · 2023-09-14T09:31:15Z

arrow-flight/tests/flight_sql_client_cli.rs

+            .arg(addr.port().to_string())
+            .arg("prepared-statement-query")
+            .arg(PREPARED_QUERY)
+            .args(["-p", "$1=string"])


💯 for the tests

alamb · 2023-09-14T09:48:15Z

arrow-flight/src/sql/server.rs

@@ -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<Peekable<Streaming<FlightData>>>,


This server is the scaffolding to help people build flightsql servers -- they can always use the raw FlightService if they prefer. Thus I think the change in API is less critical and given the requirements it seems inevitable.

I think the only thing we should try and improve in this PR is improving the documentation to explain why peekable is used somehow (to lower the cognative burden on people trying to use this).

One potential option to document this is rather than using Peekable<Streaming<...>> dirctly, would make our own wrapper, PeekableStreaming or something. While this would require duplicate a bunch of the peekable API, there would be a natural place to document what it was for and how to use it which I think would lower the barrier to usage.

For example:

/// A wrapper around `Streaming` that allows inspection of the first message. /// This is needed because sometimes the first request in the stream will contain /// a [`FlightDescriptor`] in addition to potentially any data and the dispatch logic /// must inspect this information. /// /// # example: /// <show an example here of calling `into_inner()` to get the original data back struct PeekableStreaming { inner: Peekable<Streaming<FlightData>> } impl PeekableStreaming { /// return the inner stream pub fn into_inner(self) -> Streaming<FlightData> { self.inner.into_inner() } ... }

We could also potentially use something like BoxStream<FlightData> but that would lose the gRPC specific stuff like status codes and trailers exposed by Streaming as well as being an API change as well.

Thus I think this design is the best of several less than ideal solutions. To proceed perhaps we can add some documentation on the do_*_fallback methods that mentions the stream comes from peekable

alamb · 2023-09-14T10:40:16Z

arrow-flight/src/sql/client.rs

+                .await
+                .map_err(flight_error_to_arrow_error)?;
+
+            self.flight_sql_client


I filed apache/arrow#37720 and will circulate this around

alamb

I think this is really nicely done -- thank you @suremarc

alamb · 2023-09-18T22:01:33Z

arrow-flight/src/sql/server.rs

+
+/// 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


suremarc added 4 commits September 7, 2023 12:36

change Streaming<FlightData> to Peekable<Streaming<FlightData>>

b5f6c83

add explanatory comment

e52373d

working test

352555f

trigger pre-commit hooks?

c82dd1c

github-actions bot added arrow Changes to the arrow crate arrow-flight Changes to the arrow-flight crate labels Sep 7, 2023

suremarc commented Sep 7, 2023

View reviewed changes

tustvold reviewed Sep 8, 2023

View reviewed changes

suremarc and others added 2 commits September 8, 2023 09:28

Update arrow-flight/src/sql/client.rs

7f06bc5

Co-authored-by: Raphael Taylor-Davies <1781103+tustvold@users.noreply.github.com>

remove unnecessary multi-thread annotation

80855ec

suremarc requested a review from tustvold September 8, 2023 14:54

suremarc commented Sep 8, 2023

View reviewed changes

alamb changed the title ~~Fix flight sql do put handling~~ Fix flight sql do put handling, add bind parameter support to FlightSQL cli client Sep 14, 2023

alamb mentioned this pull request Sep 14, 2023

[FlightSQL] Stateless prepared statement with parameter support apache/arrow#37720

Closed

alamb reviewed Sep 14, 2023

View reviewed changes

rework api

6cfc5c4

suremarc requested a review from alamb September 15, 2023 20:17

alamb approved these changes Sep 18, 2023

View reviewed changes

alamb merged commit 47e8a8d into apache:master Sep 18, 2023
12 checks passed

alamb added the api-change Changes to the arrow API label Sep 18, 2023

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Fix flight sql do put handling, add bind parameter support to FlightSQL cli client #4797

Fix flight sql do put handling, add bind parameter support to FlightSQL cli client #4797

suremarc commented Sep 7, 2023 •

edited by alamb

Loading

suremarc Sep 7, 2023

tustvold Sep 8, 2023

tustvold left a comment

tustvold Sep 8, 2023

suremarc Sep 8, 2023

suremarc Sep 8, 2023

tustvold Sep 8, 2023

suremarc Sep 8, 2023

tustvold Sep 8, 2023

tustvold Sep 8, 2023

suremarc Sep 8, 2023 •

edited

Loading

tustvold Sep 8, 2023

alamb Sep 14, 2023

tustvold Sep 8, 2023

suremarc Sep 8, 2023

alamb Sep 14, 2023

suremarc Sep 15, 2023

suremarc Sep 8, 2023

alamb Sep 14, 2023

alamb left a comment

alamb Sep 14, 2023

alamb Sep 14, 2023

alamb Sep 14, 2023

alamb Sep 14, 2023

alamb left a comment

alamb Sep 18, 2023

Fix flight sql do put handling, add bind parameter support to FlightSQL cli client #4797

Fix flight sql do put handling, add bind parameter support to FlightSQL cli client #4797

Conversation

suremarc commented Sep 7, 2023 • edited by alamb Loading

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are there any user-facing changes?

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tustvold left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

suremarc Sep 8, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

alamb left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

alamb left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

suremarc commented Sep 7, 2023 •

edited by alamb

Loading

suremarc Sep 8, 2023 •

edited

Loading