Repository: influxdata/influxdb
File: influxdb3_server/src/http.rs (Rust)

Code before:
            }

            impl Future for JsonFuture {
                type Output = Option<Result<Bytes, DataFusionError>>;

                fn poll(
                    mut self: Pin<&mut Self>,
                    ctx: &mut std::task::Context<'_>,
                ) -> Poll<Self::Output> {
                    match dbg!(self.stream.poll_next_unpin(ctx)) {
                        Poll::Ready(Some(batch)) => {
                            let batch = match batch {
                                Ok(batch) => batch,
                                Err(e) => return Poll::Ready(Some(Err(e))),
                            };
                            match self.state {
                                State::FirstPoll => {
                                    let mut writer = arrow_json::ArrayWriter::new(Vec::new());
                                    writer.write(&batch).unwrap();
                                    self.state = State::Body;
                                    Poll::Ready(Some(Ok(Bytes::from(writer.into_inner()))))
                                }
                                State::Body => {
                                    let mut writer = arrow_json::WriterBuilder::new()
                                        .build::<Vec<_>, JsonMiddle>(Vec::new());
                                    writer.write(&batch).unwrap();
                                    Poll::Ready(Some(Ok(Bytes::from(writer.into_inner()))))
                                }
                                _ => unreachable!(),
                            }
                        }
                        Poll::Ready(None) => {
                            match self.state {
                                // We've never written anything so just close the
                                // stream with an empty array
                                State::FirstPoll => {
                                    self.state = State::Done;
                                    Poll::Ready(Some(Ok(Bytes::from("[]"))))
                                }
                                // We have written to this before so we should
                                // write the final bytes to close the object
                                State::Body => {
                                    self.state = State::Done;
                                    Poll::Ready(Some(Ok(Bytes::from("]"))))
                                }
                                State::Done => Poll::Ready(None),
                            }
                        }
                        Poll::Pending => Poll::Pending,
                    }
                }

Review comment (bug):
Similar to above, would it be possible to return error here instead of `unwrap`?

Code after:
                                Err(e) => return Poll::Ready(Some(Err(e))),
                            };
                            let mut writer = if self.first_poll {
                                self.first_poll = false;
                                arrow_csv::Writer::new(Vec::new())
                            } else {
                                arrow_csv::WriterBuilder::new()
                                    .with_header(false)
                                    .build(Vec::new())
                            };
                            if let Err(err) = writer.write(&batch) {
                                Poll::Ready(Some(Err(err.into())))
                            } else {
                                Poll::Ready(Some(Ok(Bytes::from(writer.into_inner()))))
                            }
                        }
                        Poll::Ready(None) => Poll::Ready(None),
                        Poll::Pending => Poll::Pending,
                    }
                }
            }

            let mut future = CsvFuture {
                first_poll: true,
                stream,
            };
            Ok(Body::wrap_stream(futures::stream::poll_fn(move |ctx| {
                future.poll_unpin(ctx)
            })))
        }
        QueryFormat::Json => {
            struct JsonFuture {
                state: State,
                stream: Pin<Box<dyn RecordBatchStream + Send>>,
            }

            enum State {
                FirstPoll,
                Body,
                Done,
            }

            #[derive(Default, Debug)]
            struct JsonMiddle;

            impl arrow_json::writer::JsonFormat for JsonMiddle {
                fn start_row<W: std::io::Write>(
                    &self,
                    writer: &mut W,
                    is_first_row: bool,
                ) -> std::result::Result<(), arrow_schema::ArrowError> {