(
&self,
request: Request<Streaming<FlightData>>,
)
| 709 | } |
| 710 | |
| 711 | async fn do_put( |
| 712 | &self, |
| 713 | request: Request<Streaming<FlightData>>, |
| 714 | ) -> Result<Response<Self::DoPutStream>, Status> { |
| 715 | // See issue #4658: https://github.com/apache/arrow-rs/issues/4658 |
| 716 | // To dispatch to the correct `do_put` method, we cannot discard the first message, |
| 717 | // as it may contain the Arrow schema, which the `do_put` handler may need. |
| 718 | // To allow the first message to be reused by the `do_put` handler, |
| 719 | // we wrap this stream in a `Peekable` one, which allows us to peek at |
| 720 | // the first message without discarding it. |
| 721 | let mut request = request.map(PeekableFlightDataStream::new); |
| 722 | let mut stream = Pin::new(request.get_mut()); |
| 723 | |
| 724 | let peeked_item = stream.peek().await.cloned(); |
| 725 | let Some(cmd) = peeked_item else { |
| 726 | return self |
| 727 | .do_put_error_callback(request, DoPutError::MissingCommand) |
| 728 | .await; |
| 729 | }; |
| 730 | |
| 731 | let Some(flight_descriptor) = cmd?.flight_descriptor else { |
| 732 | return self |
| 733 | .do_put_error_callback(request, DoPutError::MissingFlightDescriptor) |
| 734 | .await; |
| 735 | }; |
| 736 | let message = Any::decode(flight_descriptor.cmd).map_err(decode_error_to_status)?; |
| 737 | match Command::try_from(message).map_err(arrow_error_to_status)? { |
| 738 | Command::CommandStatementUpdate(command) => { |
| 739 | let record_count = self.do_put_statement_update(command, request).await?; |
| 740 | let result = DoPutUpdateResult { record_count }; |
| 741 | let output = futures::stream::iter(vec![Ok(PutResult { |
| 742 | app_metadata: result.encode_to_vec().into(), |
| 743 | })]); |
| 744 | Ok(Response::new(Box::pin(output))) |
| 745 | } |
| 746 | Command::CommandStatementIngest(command) => { |
| 747 | let record_count = self.do_put_statement_ingest(command, request).await?; |
| 748 | let result = DoPutUpdateResult { record_count }; |
| 749 | let output = futures::stream::iter(vec![Ok(PutResult { |
| 750 | app_metadata: result.encode_to_vec().into(), |
| 751 | })]); |
| 752 | Ok(Response::new(Box::pin(output))) |
| 753 | } |
| 754 | Command::CommandPreparedStatementQuery(command) => { |
| 755 | let result = self |
| 756 | .do_put_prepared_statement_query(command, request) |
| 757 | .await?; |
| 758 | let output = futures::stream::iter(vec![Ok(PutResult { |
| 759 | app_metadata: result.encode_to_vec().into(), |
| 760 | })]); |
| 761 | Ok(Response::new(Box::pin(output))) |
| 762 | } |
| 763 | Command::CommandStatementSubstraitPlan(command) => { |
| 764 | let record_count = self.do_put_substrait_plan(command, request).await?; |
| 765 | let result = DoPutUpdateResult { record_count }; |
| 766 | let output = futures::stream::iter(vec![Ok(PutResult { |
| 767 | app_metadata: result.encode_to_vec().into(), |
| 768 | })]); |
nothing calls this directly
no test coverage detected