// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

#include "arrow/flight/client.h"

// Platform-specific defines
#include "arrow/flight/platform.h"

#include <map>
#include <memory>
#include <mutex>
#include <sstream>
#include <string>
#include <unordered_map>
#include <utility>

#ifdef GRPCPP_PP_INCLUDE
#include <grpcpp/grpcpp.h>
#if defined(GRPC_NAMESPACE_FOR_TLS_CREDENTIALS_OPTIONS)
#include <grpcpp/security/tls_credentials_options.h>
#endif
#else
#include <grpc++/grpc++.h>
#endif

#include <grpc/grpc_security_constants.h>

#include "arrow/buffer.h"
#include "arrow/ipc/reader.h"
#include "arrow/ipc/writer.h"
#include "arrow/record_batch.h"
#include "arrow/result.h"
#include "arrow/status.h"
#include "arrow/table.h"
#include "arrow/type.h"
#include "arrow/util/logging.h"
#include "arrow/util/uri.h"

#include "arrow/flight/client_auth.h"
#include "arrow/flight/client_header_internal.h"
#include "arrow/flight/client_middleware.h"
#include "arrow/flight/internal.h"
#include "arrow/flight/middleware.h"
#include "arrow/flight/middleware_internal.h"
#include "arrow/flight/serialization_internal.h"
#include "arrow/flight/types.h"

namespace arrow {

namespace flight {

namespace pb = arrow::flight::protocol;

const char* kWriteSizeDetailTypeId = "flight::FlightWriteSizeStatusDetail";

FlightCallOptions::FlightCallOptions()
    : timeout(-1),
      read_options(ipc::IpcReadOptions::Defaults()),
      write_options(ipc::IpcWriteOptions::Defaults()) {}

const char* FlightWriteSizeStatusDetail::type_id() const {
  return kWriteSizeDetailTypeId;
}

std::string FlightWriteSizeStatusDetail::ToString() const {
  std::stringstream ss;
  ss << "IPC payload size (" << actual_ << " bytes) exceeded soft limit (" << limit_
     << " bytes)";
  return ss.str();
}

std::shared_ptr<FlightWriteSizeStatusDetail> FlightWriteSizeStatusDetail::UnwrapStatus(
    const arrow::Status& status) {
  if (!status.detail() || status.detail()->type_id() != kWriteSizeDetailTypeId) {
    return nullptr;
  }
  return std::dynamic_pointer_cast<FlightWriteSizeStatusDetail>(status.detail());
}

FlightClientOptions FlightClientOptions::Defaults() { return FlightClientOptions(); }

Status FlightStreamReader::ReadAll(std::shared_ptr<Table>* table,
                                   const StopToken& stop_token) {
  std::vector<std::shared_ptr<RecordBatch>> batches;
  RETURN_NOT_OK(ReadAll(&batches, stop_token));
  ARROW_ASSIGN_OR_RAISE(auto schema, GetSchema());
  return Table::FromRecordBatches(schema, std::move(batches)).Value(table);
}

struct ClientRpc {
  grpc::ClientContext context;

  explicit ClientRpc(const FlightCallOptions& options) {
    if (options.timeout.count() >= 0) {
      std::chrono::system_clock::time_point deadline =
          std::chrono::time_point_cast<std::chrono::system_clock::time_point::duration>(
              std::chrono::system_clock::now() + options.timeout);
      context.set_deadline(deadline);
    }
    for (auto header : options.headers) {
      context.AddMetadata(header.first, header.second);
    }
  }

  /// \brief Add an auth token via an auth handler
  Status SetToken(ClientAuthHandler* auth_handler) {
    if (auth_handler) {
      std::string token;
      RETURN_NOT_OK(auth_handler->GetToken(&token));
      context.AddMetadata(internal::kGrpcAuthHeader, token);
    }
    return Status::OK();
  }
};

/// Helper that manages Finish() of a gRPC stream.
///
/// When we encounter an error (e.g. could not decode an IPC message),
/// we want to provide both the client-side error context and any
/// available server-side context. This helper helps wrap up that
/// logic.
///
/// This class protects the stream with a flag (so that Finish is
/// idempotent), and drains the read side (so that Finish won't hang).
///
/// The template lets us abstract between DoGet/DoExchange and DoPut,
/// which respectively read internal::FlightData and pb::PutResult.
template <typename Stream, typename ReadT>
class FinishableStream {
 public:
  FinishableStream(std::shared_ptr<ClientRpc> rpc, std::shared_ptr<Stream> stream)
      : rpc_(rpc), stream_(stream), finished_(false), server_status_() {}
  virtual ~FinishableStream() = default;

  /// \brief Get the underlying stream.
  std::shared_ptr<Stream> stream() const { return stream_; }

  /// \brief Finish the call, adding server context to the given status.
  virtual Status Finish(Status st) {
    if (finished_) {
      return MergeStatus(std::move(st));
    }

    // Drain the read side, as otherwise gRPC Finish() will hang. We
    // only call Finish() when the client closes the writer or the
    // reader finishes, so it's OK to assume the client no longer
    // wants to read and drain the read side. (If the client wants to
    // indicate that it is done writing, but not done reading, it
    // should use DoneWriting.
    ReadT message;
    while (internal::ReadPayload(stream_.get(), &message)) {
      // Drain the read side to avoid gRPC hanging in Finish()
    }

    server_status_ = internal::FromGrpcStatus(stream_->Finish(), &rpc_->context);
    finished_ = true;

    return MergeStatus(std::move(st));
  }

 private:
  Status MergeStatus(Status&& st) {
    if (server_status_.ok()) {
      return std::move(st);
    }
    return Status::FromDetailAndArgs(
        server_status_.code(), server_status_.detail(), server_status_.message(),
        ". Client context: ", st.ToString(),
        ". gRPC client debug context: ", rpc_->context.debug_error_string());
  }

  std::shared_ptr<ClientRpc> rpc_;
  std::shared_ptr<Stream> stream_;
  bool finished_;
  Status server_status_;
};

/// Helper that manages \a Finish() of a read-write gRPC stream.
///
/// This also calls \a WritesDone() and protects itself with a mutex
/// to enable sharing between the reader and writer.
template <typename Stream, typename ReadT>
class FinishableWritableStream : public FinishableStream<Stream, ReadT> {
 public:
  FinishableWritableStream(std::shared_ptr<ClientRpc> rpc,
                           std::shared_ptr<std::mutex> read_mutex,
                           std::shared_ptr<Stream> stream)
      : FinishableStream<Stream, ReadT>(rpc, stream),
        finish_mutex_(),
        read_mutex_(read_mutex),
        done_writing_(false) {}
  virtual ~FinishableWritableStream() = default;

  /// \brief Indicate to gRPC that the write half of the stream is done.
  Status DoneWriting() {
    // This is only used by the writer side of a stream, so it need
    // not be protected with a lock.
    if (done_writing_) {
      return Status::OK();
    }
    done_writing_ = true;
    if (!this->stream()->WritesDone()) {
      // Error happened, try to close the stream to get more detailed info
      return Finish(MakeFlightError(FlightStatusCode::Internal,
                                    "Could not flush pending record batches"));
    }
    return Status::OK();
  }

  Status Finish(Status st) override {
    // This may be used concurrently by reader/writer side of a
    // stream, so it needs to be protected.
    std::lock_guard<std::mutex> guard(finish_mutex_);

    // Now that we're shared between a reader and writer, we need to
    // protect ourselves from being called while there's an
    // outstanding read.
    std::unique_lock<std::mutex> read_guard(*read_mutex_, std::try_to_lock);
    if (!read_guard.owns_lock()) {
      return MakeFlightError(
          FlightStatusCode::Internal,
          "Cannot close stream with pending read operation. Client context: " +
              st.ToString());
    }

    // Try to flush pending writes. Don't use our WritesDone() to
    // avoid recursion.
    bool finished_writes = done_writing_ || this->stream()->WritesDone();
    done_writing_ = true;

    st = FinishableStream<Stream, ReadT>::Finish(std::move(st));

    if (!finished_writes) {
      return Status::FromDetailAndArgs(
          st.code(), st.detail(), st.message(),
          ". Additionally, could not finish writing record batches before closing");
    }
    return st;
  }

 private:
  std::mutex finish_mutex_;
  std::shared_ptr<std::mutex> read_mutex_;
  bool done_writing_;
};

class GrpcAddCallHeaders : public AddCallHeaders {
 public:
  explicit GrpcAddCallHeaders(std::multimap<grpc::string, grpc::string>* metadata)
      : metadata_(metadata) {}
  ~GrpcAddCallHeaders() override = default;

  void AddHeader(const std::string& key, const std::string& value) override {
    metadata_->insert(std::make_pair(key, value));
  }

 private:
  std::multimap<grpc::string, grpc::string>* metadata_;
};

class GrpcClientInterceptorAdapter : public grpc::experimental::Interceptor {
 public:
  explicit GrpcClientInterceptorAdapter(
      std::vector<std::unique_ptr<ClientMiddleware>> middleware)
      : middleware_(std::move(middleware)), received_headers_(false) {}

  void Intercept(grpc::experimental::InterceptorBatchMethods* methods) {
    using InterceptionHookPoints = grpc::experimental::InterceptionHookPoints;
    if (methods->QueryInterceptionHookPoint(
            InterceptionHookPoints::PRE_SEND_INITIAL_METADATA)) {
      GrpcAddCallHeaders add_headers(methods->GetSendInitialMetadata());
      for (const auto& middleware : middleware_) {
        middleware->SendingHeaders(&add_headers);
      }
    }

    if (methods->QueryInterceptionHookPoint(
            InterceptionHookPoints::POST_RECV_INITIAL_METADATA)) {
      if (!methods->GetRecvInitialMetadata()->empty()) {
        ReceivedHeaders(*methods->GetRecvInitialMetadata());
      }
    }

    if (methods->QueryInterceptionHookPoint(InterceptionHookPoints::POST_RECV_STATUS)) {
      DCHECK_NE(nullptr, methods->GetRecvStatus());
      DCHECK_NE(nullptr, methods->GetRecvTrailingMetadata());
      ReceivedHeaders(*methods->GetRecvTrailingMetadata());
      const Status status = internal::FromGrpcStatus(*methods->GetRecvStatus());
      for (const auto& middleware : middleware_) {
        middleware->CallCompleted(status);
      }
    }

    methods->Proceed();
  }

 private:
  void ReceivedHeaders(
      const std::multimap<grpc::string_ref, grpc::string_ref>& metadata) {
    if (received_headers_) {
      return;
    }
    received_headers_ = true;
    CallHeaders headers;
    for (const auto& entry : metadata) {
      headers.insert({util::string_view(entry.first.data(), entry.first.length()),
                      util::string_view(entry.second.data(), entry.second.length())});
    }
    for (const auto& middleware : middleware_) {
      middleware->ReceivedHeaders(headers);
    }
  }

  std::vector<std::unique_ptr<ClientMiddleware>> middleware_;
  // When communicating with a gRPC-Java server, the server may not
  // send back headers if the call fails right away. Instead, the
  // headers will be consolidated into the trailers. We don't want to
  // call the client middleware callback twice, so instead track
  // whether we saw headers - if not, then we need to check trailers.
  bool received_headers_;
};

class GrpcClientInterceptorAdapterFactory
    : public grpc::experimental::ClientInterceptorFactoryInterface {
 public:
  GrpcClientInterceptorAdapterFactory(
      std::vector<std::shared_ptr<ClientMiddlewareFactory>> middleware)
      : middleware_(middleware) {}

  grpc::experimental::Interceptor* CreateClientInterceptor(
      grpc::experimental::ClientRpcInfo* info) override {
    std::vector<std::unique_ptr<ClientMiddleware>> middleware;

    FlightMethod flight_method = FlightMethod::Invalid;
    util::string_view method(info->method());
    if (method.ends_with("/Handshake")) {
      flight_method = FlightMethod::Handshake;
    } else if (method.ends_with("/ListFlights")) {
      flight_method = FlightMethod::ListFlights;
    } else if (method.ends_with("/GetFlightInfo")) {
      flight_method = FlightMethod::GetFlightInfo;
    } else if (method.ends_with("/GetSchema")) {
      flight_method = FlightMethod::GetSchema;
    } else if (method.ends_with("/DoGet")) {
      flight_method = FlightMethod::DoGet;
    } else if (method.ends_with("/DoPut")) {
      flight_method = FlightMethod::DoPut;
    } else if (method.ends_with("/DoExchange")) {
      flight_method = FlightMethod::DoExchange;
    } else if (method.ends_with("/DoAction")) {
      flight_method = FlightMethod::DoAction;
    } else if (method.ends_with("/ListActions")) {
      flight_method = FlightMethod::ListActions;
    } else {
      DCHECK(false) << "Unknown Flight method: " << info->method();
    }

    const CallInfo flight_info{flight_method};
    for (auto& factory : middleware_) {
      std::unique_ptr<ClientMiddleware> instance;
      factory->StartCall(flight_info, &instance);
      if (instance) {
        middleware.push_back(std::move(instance));
      }
    }
    return new GrpcClientInterceptorAdapter(std::move(middleware));
  }

 private:
  std::vector<std::shared_ptr<ClientMiddlewareFactory>> middleware_;
};

class GrpcClientAuthSender : public ClientAuthSender {
 public:
  explicit GrpcClientAuthSender(
      std::shared_ptr<
          grpc::ClientReaderWriter<pb::HandshakeRequest, pb::HandshakeResponse>>
          stream)
      : stream_(stream) {}

  Status Write(const std::string& token) override {
    pb::HandshakeRequest response;
    response.set_payload(token);
    if (stream_->Write(response)) {
      return Status::OK();
    }
    return internal::FromGrpcStatus(stream_->Finish());
  }

 private:
  std::shared_ptr<grpc::ClientReaderWriter<pb::HandshakeRequest, pb::HandshakeResponse>>
      stream_;
};

class GrpcClientAuthReader : public ClientAuthReader {
 public:
  explicit GrpcClientAuthReader(
      std::shared_ptr<
          grpc::ClientReaderWriter<pb::HandshakeRequest, pb::HandshakeResponse>>
          stream)
      : stream_(stream) {}

  Status Read(std::string* token) override {
    pb::HandshakeResponse request;
    if (stream_->Read(&request)) {
      *token = std::move(*request.mutable_payload());
      return Status::OK();
    }
    return internal::FromGrpcStatus(stream_->Finish());
  }

 private:
  std::shared_ptr<grpc::ClientReaderWriter<pb::HandshakeRequest, pb::HandshakeResponse>>
      stream_;
};

// An ipc::MessageReader that adapts any readable gRPC stream
// returning FlightData.
template <typename Reader>
class GrpcIpcMessageReader : public ipc::MessageReader {
 public:
  GrpcIpcMessageReader(
      std::shared_ptr<ClientRpc> rpc, std::shared_ptr<std::mutex> read_mutex,
      std::shared_ptr<FinishableStream<Reader, internal::FlightData>> stream,
      std::shared_ptr<internal::PeekableFlightDataReader<std::shared_ptr<Reader>>>
          peekable_reader,
      std::shared_ptr<Buffer>* app_metadata)
      : rpc_(rpc),
        read_mutex_(read_mutex),
        stream_(std::move(stream)),
        peekable_reader_(peekable_reader),
        app_metadata_(app_metadata),
        stream_finished_(false) {}

  ::arrow::Result<std::unique_ptr<ipc::Message>> ReadNextMessage() override {
    if (stream_finished_) {
      return nullptr;
    }
    internal::FlightData* data;
    {
      auto guard = read_mutex_ ? std::unique_lock<std::mutex>(*read_mutex_)
                               : std::unique_lock<std::mutex>();
      peekable_reader_->Next(&data);
    }
    if (!data) {
      stream_finished_ = true;
      return stream_->Finish(Status::OK());
    }
    // Validate IPC message
    auto result = data->OpenMessage();
    if (!result.ok()) {
      return stream_->Finish(std::move(result).status());
    }
    *app_metadata_ = std::move(data->app_metadata);
    return result;
  }

 private:
  // The RPC context lifetime must be coupled to the ClientReader
  std::shared_ptr<ClientRpc> rpc_;
  // Guard reads with a mutex to prevent concurrent reads if the write
  // side calls Finish(). Nullable as DoGet doesn't need this.
  std::shared_ptr<std::mutex> read_mutex_;
  std::shared_ptr<FinishableStream<Reader, internal::FlightData>> stream_;
  std::shared_ptr<internal::PeekableFlightDataReader<std::shared_ptr<Reader>>>
      peekable_reader_;
  // A reference to GrpcStreamReader.app_metadata_. That class
  // can't access the app metadata because when it Peek()s the stream,
  // it may be looking at a dictionary batch, not the record
  // batch. Updating it here ensures the reader is always updated with
  // the last metadata message read.
  std::shared_ptr<Buffer>* app_metadata_;
  bool stream_finished_;
};

/// The implementation of the public-facing API for reading from a
/// FlightData stream
template <typename Reader>
class GrpcStreamReader : public FlightStreamReader {
 public:
  GrpcStreamReader(std::shared_ptr<ClientRpc> rpc, std::shared_ptr<std::mutex> read_mutex,
                   const ipc::IpcReadOptions& options, StopToken stop_token,
                   std::shared_ptr<FinishableStream<Reader, internal::FlightData>> stream)
      : rpc_(rpc),
        read_mutex_(read_mutex),
        options_(options),
        stop_token_(std::move(stop_token)),
        stream_(stream),
        peekable_reader_(new internal::PeekableFlightDataReader<std::shared_ptr<Reader>>(
            stream->stream())),
        app_metadata_(nullptr) {}

  Status EnsureDataStarted() {
    if (!batch_reader_) {
      bool skipped_to_data = false;
      {
        auto guard = TakeGuard();
        skipped_to_data = peekable_reader_->SkipToData();
      }
      // peek() until we find the first data message; discard metadata
      if (!skipped_to_data) {
        return OverrideWithServerError(MakeFlightError(
            FlightStatusCode::Internal, "Server never sent a data message"));
      }

      auto message_reader =
          std::unique_ptr<ipc::MessageReader>(new GrpcIpcMessageReader<Reader>(
              rpc_, read_mutex_, stream_, peekable_reader_, &app_metadata_));
      auto result =
          ipc::RecordBatchStreamReader::Open(std::move(message_reader), options_);
      RETURN_NOT_OK(OverrideWithServerError(std::move(result).Value(&batch_reader_)));
    }
    return Status::OK();
  }
  arrow::Result<std::shared_ptr<Schema>> GetSchema() override {
    RETURN_NOT_OK(EnsureDataStarted());
    return batch_reader_->schema();
  }
  Status Next(FlightStreamChunk* out) override {
    internal::FlightData* data;
    {
      auto guard = TakeGuard();
      peekable_reader_->Peek(&data);
    }
    if (!data) {
      out->app_metadata = nullptr;
      out->data = nullptr;
      return stream_->Finish(Status::OK());
    }

    if (!data->metadata) {
      // Metadata-only (data->metadata is the IPC header)
      out->app_metadata = data->app_metadata;
      out->data = nullptr;
      {
        auto guard = TakeGuard();
        peekable_reader_->Next(&data);
      }
      return Status::OK();
    }

    if (!batch_reader_) {
      RETURN_NOT_OK(EnsureDataStarted());
      // Re-peek here since EnsureDataStarted() advances the stream
      return Next(out);
    }
    RETURN_NOT_OK(batch_reader_->ReadNext(&out->data));
    out->app_metadata = std::move(app_metadata_);
    return Status::OK();
  }
  Status ReadAll(std::vector<std::shared_ptr<RecordBatch>>* batches) override {
    return ReadAll(batches, stop_token_);
  }
  Status ReadAll(std::vector<std::shared_ptr<RecordBatch>>* batches,
                 const StopToken& stop_token) override {
    FlightStreamChunk chunk;

    while (true) {
      if (stop_token.IsStopRequested()) {
        Cancel();
        return stop_token.Poll();
      }
      RETURN_NOT_OK(Next(&chunk));
      if (!chunk.data) break;
      batches->emplace_back(std::move(chunk.data));
    }
    return Status::OK();
  }
  Status ReadAll(std::shared_ptr<Table>* table) override {
    return ReadAll(table, stop_token_);
  }
  using FlightStreamReader::ReadAll;
  void Cancel() override { rpc_->context.TryCancel(); }

 private:
  std::unique_lock<std::mutex> TakeGuard() {
    return read_mutex_ ? std::unique_lock<std::mutex>(*read_mutex_)
                       : std::unique_lock<std::mutex>();
  }

  Status OverrideWithServerError(Status&& st) {
    if (st.ok()) {
      return std::move(st);
    }
    return stream_->Finish(std::move(st));
  }

  friend class GrpcIpcMessageReader<Reader>;
  std::shared_ptr<ClientRpc> rpc_;
  // Guard reads with a lock to prevent Finish()/Close() from being
  // called on the writer while the reader has a pending
  // read. Nullable, as DoGet() doesn't need this.
  std::shared_ptr<std::mutex> read_mutex_;
  ipc::IpcReadOptions options_;
  StopToken stop_token_;
  std::shared_ptr<FinishableStream<Reader, internal::FlightData>> stream_;
  std::shared_ptr<internal::PeekableFlightDataReader<std::shared_ptr<Reader>>>
      peekable_reader_;
  std::shared_ptr<ipc::RecordBatchReader> batch_reader_;
  std::shared_ptr<Buffer> app_metadata_;
};

// The next two classes implement writing to a FlightData stream.
// Similarly to the read side, we want to reuse the implementation of
// RecordBatchWriter. As a result, these two classes are intertwined
// in order to pass application metadata "through" RecordBatchWriter.
// In order to get application-specific metadata to the
// IpcPayloadWriter, DoPutPayloadWriter takes a pointer to
// GrpcStreamWriter. GrpcStreamWriter updates a metadata field on
// write; DoPutPayloadWriter reads that metadata field to determine
// what to write.

template <typename ProtoReadT, typename FlightReadT>
class DoPutPayloadWriter;

template <typename ProtoReadT, typename FlightReadT>
class GrpcStreamWriter : public FlightStreamWriter {
 public:
  ~GrpcStreamWriter() override = default;

  using GrpcStream = grpc::ClientReaderWriter<pb::FlightData, ProtoReadT>;

  explicit GrpcStreamWriter(
      const FlightDescriptor& descriptor, std::shared_ptr<ClientRpc> rpc,
      int64_t write_size_limit_bytes, const ipc::IpcWriteOptions& options,
      std::shared_ptr<FinishableWritableStream<GrpcStream, FlightReadT>> writer)
      : app_metadata_(nullptr),
        batch_writer_(nullptr),
        writer_(std::move(writer)),
        rpc_(std::move(rpc)),
        write_size_limit_bytes_(write_size_limit_bytes),
        options_(options),
        descriptor_(descriptor),
        writer_closed_(false) {}

  static Status Open(
      const FlightDescriptor& descriptor, std::shared_ptr<Schema> schema,
      const ipc::IpcWriteOptions& options, std::shared_ptr<ClientRpc> rpc,
      int64_t write_size_limit_bytes,
      std::shared_ptr<FinishableWritableStream<GrpcStream, FlightReadT>> writer,
      std::unique_ptr<FlightStreamWriter>* out);

  Status CheckStarted() {
    if (!batch_writer_) {
      return Status::Invalid("Writer not initialized. Call Begin() with a schema.");
    }
    return Status::OK();
  }

  Status Begin(const std::shared_ptr<Schema>& schema,
               const ipc::IpcWriteOptions& options) override {
    if (batch_writer_) {
      return Status::Invalid("This writer has already been started.");
    }
    std::unique_ptr<ipc::internal::IpcPayloadWriter> payload_writer(
        new DoPutPayloadWriter<ProtoReadT, FlightReadT>(
            descriptor_, std::move(rpc_), write_size_limit_bytes_, writer_, this));
    // XXX: this does not actually write the message to the stream.
    // See Close().
    ARROW_ASSIGN_OR_RAISE(batch_writer_, ipc::internal::OpenRecordBatchWriter(
                                             std::move(payload_writer), schema, options));
    return Status::OK();
  }

  Status Begin(const std::shared_ptr<Schema>& schema) override {
    return Begin(schema, options_);
  }

  Status WriteRecordBatch(const RecordBatch& batch) override {
    RETURN_NOT_OK(CheckStarted());
    return WriteWithMetadata(batch, nullptr);
  }

  Status WriteMetadata(std::shared_ptr<Buffer> app_metadata) override {
    FlightPayload payload{};
    payload.app_metadata = app_metadata;
    auto status = internal::WritePayload(payload, writer_->stream().get());
    if (status.IsIOError()) {
      return writer_->Finish(MakeFlightError(FlightStatusCode::Internal,
                                             "Could not write metadata to stream"));
    }
    return status;
  }

  Status WriteWithMetadata(const RecordBatch& batch,
                           std::shared_ptr<Buffer> app_metadata) override {
    RETURN_NOT_OK(CheckStarted());
    app_metadata_ = app_metadata;
    return batch_writer_->WriteRecordBatch(batch);
  }

  Status DoneWriting() override {
    // Do not CheckStarted - DoneWriting applies to data and metadata
    if (batch_writer_) {
      // Close the writer if we have one; this will force it to flush any
      // remaining data, before we close the write side of the stream.
      writer_closed_ = true;
      Status st = batch_writer_->Close();
      if (!st.ok()) {
        return writer_->Finish(std::move(st));
      }
    }
    return writer_->DoneWriting();
  }

  Status Close() override {
    // Do not CheckStarted - Close applies to data and metadata
    if (batch_writer_ && !writer_closed_) {
      // This is important! Close() calls
      // IpcPayloadWriter::CheckStarted() which will force the initial
      // schema message to be written to the stream. This is required
      // to unstick the server, else the client and the server end up
      // waiting for each other. This happens if the client never
      // wrote anything before calling Close().
      writer_closed_ = true;
      return writer_->Finish(batch_writer_->Close());
    }
    return writer_->Finish(Status::OK());
  }

  ipc::WriteStats stats() const override {
    ARROW_CHECK_NE(batch_writer_, nullptr);
    return batch_writer_->stats();
  }

 private:
  friend class DoPutPayloadWriter<ProtoReadT, FlightReadT>;
  std::shared_ptr<Buffer> app_metadata_;
  std::unique_ptr<ipc::RecordBatchWriter> batch_writer_;
  std::shared_ptr<FinishableWritableStream<GrpcStream, FlightReadT>> writer_;

  // Fields used to lazy-initialize the IpcPayloadWriter. They're
  // invalid once Begin() is called.
  std::shared_ptr<ClientRpc> rpc_;
  int64_t write_size_limit_bytes_;
  ipc::IpcWriteOptions options_;
  FlightDescriptor descriptor_;
  bool writer_closed_;
};

/// A IpcPayloadWriter implementation that writes to a gRPC stream of
/// FlightData messages.
template <typename ProtoReadT, typename FlightReadT>
class DoPutPayloadWriter : public ipc::internal::IpcPayloadWriter {
 public:
  using GrpcStream = grpc::ClientReaderWriter<pb::FlightData, ProtoReadT>;

  DoPutPayloadWriter(
      const FlightDescriptor& descriptor, std::shared_ptr<ClientRpc> rpc,
      int64_t write_size_limit_bytes,
      std::shared_ptr<FinishableWritableStream<GrpcStream, FlightReadT>> writer,
      GrpcStreamWriter<ProtoReadT, FlightReadT>* stream_writer)
      : descriptor_(descriptor),
        rpc_(rpc),
        write_size_limit_bytes_(write_size_limit_bytes),
        writer_(std::move(writer)),
        first_payload_(true),
        stream_writer_(stream_writer) {}

  ~DoPutPayloadWriter() override = default;

  Status Start() override { return Status::OK(); }

  Status WritePayload(const ipc::IpcPayload& ipc_payload) override {
    FlightPayload payload;
    payload.ipc_message = ipc_payload;

    if (first_payload_) {
      // First Flight message needs to encore the Flight descriptor
      if (ipc_payload.type != ipc::MessageType::SCHEMA) {
        return Status::Invalid("First IPC message should be schema");
      }
      // Write the descriptor to begin with
      RETURN_NOT_OK(internal::ToPayload(descriptor_, &payload.descriptor));
      first_payload_ = false;
    } else if (ipc_payload.type == ipc::MessageType::RECORD_BATCH &&
               stream_writer_->app_metadata_) {
      payload.app_metadata = std::move(stream_writer_->app_metadata_);
    }

    if (write_size_limit_bytes_ > 0) {
      // Check if the total size is greater than the user-configured
      // soft-limit.
      int64_t size = ipc_payload.body_length + ipc_payload.metadata->size();
      if (payload.descriptor) {
        size += payload.descriptor->size();
      }
      if (payload.app_metadata) {
        size += payload.app_metadata->size();
      }
      if (size > write_size_limit_bytes_) {
        return arrow::Status(
            arrow::StatusCode::Invalid, "IPC payload size exceeded soft limit",
            std::make_shared<FlightWriteSizeStatusDetail>(write_size_limit_bytes_, size));
      }
    }

    auto status = internal::WritePayload(payload, writer_->stream().get());
    if (status.IsIOError()) {
      return writer_->Finish(MakeFlightError(FlightStatusCode::Internal,
                                             "Could not write record batch to stream"));
    }
    return status;
  }

  Status Close() override {
    // Closing is handled one layer up in GrpcStreamWriter::Close
    return Status::OK();
  }

 protected:
  const FlightDescriptor descriptor_;
  std::shared_ptr<ClientRpc> rpc_;
  int64_t write_size_limit_bytes_;
  std::shared_ptr<FinishableWritableStream<GrpcStream, FlightReadT>> writer_;
  bool first_payload_;
  GrpcStreamWriter<ProtoReadT, FlightReadT>* stream_writer_;
};

template <typename ProtoReadT, typename FlightReadT>
Status GrpcStreamWriter<ProtoReadT, FlightReadT>::Open(
    const FlightDescriptor& descriptor,
    std::shared_ptr<Schema> schema,  // this schema is nullable
    const ipc::IpcWriteOptions& options, std::shared_ptr<ClientRpc> rpc,
    int64_t write_size_limit_bytes,
    std::shared_ptr<FinishableWritableStream<GrpcStream, FlightReadT>> writer,
    std::unique_ptr<FlightStreamWriter>* out) {
  std::unique_ptr<GrpcStreamWriter<ProtoReadT, FlightReadT>> instance(
      new GrpcStreamWriter<ProtoReadT, FlightReadT>(
          descriptor, std::move(rpc), write_size_limit_bytes, options, writer));
  if (schema) {
    // The schema was provided (DoPut). Eagerly write the schema and
    // descriptor together as the first message.
    RETURN_NOT_OK(instance->Begin(schema, options));
  } else {
    // The schema was not provided (DoExchange). Eagerly write just
    // the descriptor as the first message. Note that if the client
    // calls Begin() to send data, we'll send a redundant descriptor.
    FlightPayload payload{};
    RETURN_NOT_OK(internal::ToPayload(descriptor, &payload.descriptor));
    auto status = internal::WritePayload(payload, instance->writer_->stream().get());
    if (status.IsIOError()) {
      return writer->Finish(MakeFlightError(FlightStatusCode::Internal,
                                            "Could not write descriptor to stream"));
    }
    RETURN_NOT_OK(status);
  }
  *out = std::move(instance);
  return Status::OK();
}

FlightMetadataReader::~FlightMetadataReader() = default;

class GrpcMetadataReader : public FlightMetadataReader {
 public:
  explicit GrpcMetadataReader(
      std::shared_ptr<grpc::ClientReaderWriter<pb::FlightData, pb::PutResult>> reader,
      std::shared_ptr<std::mutex> read_mutex)
      : reader_(reader), read_mutex_(read_mutex) {}

  Status ReadMetadata(std::shared_ptr<Buffer>* out) override {
    std::lock_guard<std::mutex> guard(*read_mutex_);
    pb::PutResult message;
    if (reader_->Read(&message)) {
      *out = Buffer::FromString(std::move(*message.mutable_app_metadata()));
    } else {
      // Stream finished
      *out = nullptr;
    }
    return Status::OK();
  }

 private:
  std::shared_ptr<grpc::ClientReaderWriter<pb::FlightData, pb::PutResult>> reader_;
  std::shared_ptr<std::mutex> read_mutex_;
};

namespace {
// Dummy self-signed certificate to be used because TlsCredentials
// requires root CA certs, even if you are skipping server
// verification.
#if defined(GRPC_NAMESPACE_FOR_TLS_CREDENTIALS_OPTIONS)
constexpr char kDummyRootCert[] =
    "-----BEGIN CERTIFICATE-----\n"
    "MIICwzCCAaugAwIBAgIJAM12DOkcaqrhMA0GCSqGSIb3DQEBBQUAMBQxEjAQBgNV\n"
    "BAMTCWxvY2FsaG9zdDAeFw0yMDEwMDcwODIyNDFaFw0zMDEwMDUwODIyNDFaMBQx\n"
    "EjAQBgNVBAMTCWxvY2FsaG9zdDCCASIwDQYJKoZIhvcNAQEBBQADggEPADCCAQoC\n"
    "ggEBALjJ8KPEpF0P4GjMPrJhjIBHUL0AX9E4oWdgJRCSFkPKKEWzQabTQBikMOhI\n"
    "W4VvBMaHEBuECE5OEyrzDRiAO354I4F4JbBfxMOY8NIW0uWD6THWm2KkCzoZRIPW\n"
    "yZL6dN+mK6cEH+YvbNuy5ZQGNjGG43tyiXdOCAc4AI9POeTtjdMpbbpR2VY4Ad/E\n"
    "oTEiS3gNnN7WIAdgMhCJxjzvPwKszV3f7pwuTHzFMsuHLKr6JeaVUYfbi4DxxC8Z\n"
    "k6PF6dLlLf3ngTSLBJyaXP1BhKMvz0TaMK3F0y2OGwHM9J8np2zWjTlNVEzffQZx\n"
    "SWMOQManlJGs60xYx9KCPJMZZsMCAwEAAaMYMBYwFAYDVR0RBA0wC4IJbG9jYWxo\n"
    "b3N0MA0GCSqGSIb3DQEBBQUAA4IBAQC0LrmbcNKgO+D50d/wOc+vhi9K04EZh8bg\n"
    "WYAK1kLOT4eShbzqWGV/1EggY4muQ6ypSELCLuSsg88kVtFQIeRilA6bHFqQSj6t\n"
    "sqgh2cWsMwyllCtmX6Maf3CLb2ZdoJlqUwdiBdrbIbuyeAZj3QweCtLKGSQzGDyI\n"
    "KH7G8nC5d0IoRPiCMB6RnMMKsrhviuCdWbAFHop7Ff36JaOJ8iRa2sSf2OXE8j/5\n"
    "obCXCUvYHf4Zw27JcM2AnnQI9VJLnYxis83TysC5s2Z7t0OYNS9kFmtXQbUNlmpS\n"
    "doQ/Eu47vWX7S0TXeGziGtbAOKxbHE0BGGPDOAB/jGW/JVbeTiXY\n"
    "-----END CERTIFICATE-----\n";
#endif
}  // namespace
class FlightClient::FlightClientImpl {
 public:
  Status Connect(const Location& location, const FlightClientOptions& options) {
    const std::string& scheme = location.scheme();

    std::stringstream grpc_uri;
    std::shared_ptr<grpc::ChannelCredentials> creds;
    if (scheme == kSchemeGrpc || scheme == kSchemeGrpcTcp || scheme == kSchemeGrpcTls) {
      grpc_uri << arrow::internal::UriEncodeHost(location.uri_->host()) << ':'
               << location.uri_->port_text();

      if (scheme == kSchemeGrpcTls) {
        if (options.disable_server_verification) {
#if defined(GRPC_NAMESPACE_FOR_TLS_CREDENTIALS_OPTIONS)
          namespace ge = GRPC_NAMESPACE_FOR_TLS_CREDENTIALS_OPTIONS;

          // A callback to supply to TlsCredentialsOptions that accepts any server
          // arguments.
          struct NoOpTlsAuthorizationCheck
              : public ge::TlsServerAuthorizationCheckInterface {
            int Schedule(ge::TlsServerAuthorizationCheckArg* arg) override {
              arg->set_success(1);
              arg->set_status(GRPC_STATUS_OK);
              return 0;
            }
          };
          auto server_authorization_check = std::make_shared<NoOpTlsAuthorizationCheck>();
          noop_auth_check_ = std::make_shared<ge::TlsServerAuthorizationCheckConfig>(
              server_authorization_check);
#if defined(GRPC_USE_TLS_CHANNEL_CREDENTIALS_OPTIONS)
          auto certificate_provider =
              std::make_shared<grpc::experimental::StaticDataCertificateProvider>(
                  kDummyRootCert);
#if defined(GRPC_USE_TLS_CHANNEL_CREDENTIALS_OPTIONS_ROOT_CERTS)
          grpc::experimental::TlsChannelCredentialsOptions tls_options(
              certificate_provider);
#else
          // While gRPC >= 1.36 does not require a root cert (it has a default)
          // in practice the path it hardcodes is broken. See grpc/grpc#21655.
          grpc::experimental::TlsChannelCredentialsOptions tls_options;
          tls_options.set_certificate_provider(certificate_provider);
#endif
          tls_options.watch_root_certs();
          tls_options.set_root_cert_name("dummy");
          tls_options.set_server_verification_option(
              grpc_tls_server_verification_option::GRPC_TLS_SKIP_ALL_SERVER_VERIFICATION);
          tls_options.set_server_authorization_check_config(noop_auth_check_);
#elif defined(GRPC_NAMESPACE_FOR_TLS_CREDENTIALS_OPTIONS)
          auto materials_config = std::make_shared<ge::TlsKeyMaterialsConfig>();
          materials_config->set_pem_root_certs(kDummyRootCert);
          ge::TlsCredentialsOptions tls_options(
              GRPC_SSL_DONT_REQUEST_CLIENT_CERTIFICATE,
              GRPC_TLS_SKIP_ALL_SERVER_VERIFICATION, materials_config,
              std::shared_ptr<ge::TlsCredentialReloadConfig>(), noop_auth_check_);
#endif
          creds = ge::TlsCredentials(tls_options);
#else
          return Status::NotImplemented(
              "Using encryption with server verification disabled is unsupported. "
              "Please use a release of Arrow Flight built with gRPC 1.27 or higher.");
#endif
        } else {
          grpc::SslCredentialsOptions ssl_options;
          if (!options.tls_root_certs.empty()) {
            ssl_options.pem_root_certs = options.tls_root_certs;
          }
          if (!options.cert_chain.empty()) {
            ssl_options.pem_cert_chain = options.cert_chain;
          }
          if (!options.private_key.empty()) {
            ssl_options.pem_private_key = options.private_key;
          }
          creds = grpc::SslCredentials(ssl_options);
        }
      } else {
        creds = grpc::InsecureChannelCredentials();
      }
    } else if (scheme == kSchemeGrpcUnix) {
      grpc_uri << "unix://" << location.uri_->path();
      creds = grpc::InsecureChannelCredentials();
    } else {
      return Status::NotImplemented("Flight scheme " + scheme + " is not supported.");
    }

    grpc::ChannelArguments args;
    // We can't set the same config value twice, so for values where
    // we want to set defaults, keep them in a map and update them;
    // then update them all at once
    std::unordered_map<std::string, int> default_args;
    // Try to reconnect quickly at first, in case the server is still starting up
    default_args[GRPC_ARG_INITIAL_RECONNECT_BACKOFF_MS] = 100;
    // Receive messages of any size
    default_args[GRPC_ARG_MAX_RECEIVE_MESSAGE_LENGTH] = -1;
    // Setting this arg enables each client to open it's own TCP connection to server,
    // not sharing one single connection, which becomes bottleneck under high load.
    default_args[GRPC_ARG_USE_LOCAL_SUBCHANNEL_POOL] = 1;

    if (options.override_hostname != "") {
      args.SetSslTargetNameOverride(options.override_hostname);
    }

    // Allow setting generic gRPC options.
    for (const auto& arg : options.generic_options) {
      if (util::holds_alternative<int>(arg.second)) {
        default_args[arg.first] = util::get<int>(arg.second);
      } else if (util::holds_alternative<std::string>(arg.second)) {
        args.SetString(arg.first, util::get<std::string>(arg.second));
      }
      // Otherwise unimplemented
    }
    for (const auto& pair : default_args) {
      args.SetInt(pair.first, pair.second);
    }

    std::vector<std::unique_ptr<grpc::experimental::ClientInterceptorFactoryInterface>>
        interceptors;
    interceptors.emplace_back(
        new GrpcClientInterceptorAdapterFactory(std::move(options.middleware)));

    stub_ = pb::FlightService::NewStub(
        grpc::experimental::CreateCustomChannelWithInterceptors(
            grpc_uri.str(), creds, args, std::move(interceptors)));

    write_size_limit_bytes_ = options.write_size_limit_bytes;
    return Status::OK();
  }

  Status Authenticate(const FlightCallOptions& options,
                      std::unique_ptr<ClientAuthHandler> auth_handler) {
    auth_handler_ = std::move(auth_handler);
    ClientRpc rpc(options);
    std::shared_ptr<grpc::ClientReaderWriter<pb::HandshakeRequest, pb::HandshakeResponse>>
        stream = stub_->Handshake(&rpc.context);
    GrpcClientAuthSender outgoing{stream};
    GrpcClientAuthReader incoming{stream};
    RETURN_NOT_OK(auth_handler_->Authenticate(&outgoing, &incoming));
    // Explicitly close our side of the connection
    bool finished_writes = stream->WritesDone();
    RETURN_NOT_OK(internal::FromGrpcStatus(stream->Finish(), &rpc.context));
    if (!finished_writes) {
      return MakeFlightError(FlightStatusCode::Internal,
                             "Could not finish writing before closing");
    }
    return Status::OK();
  }

  arrow::Result<std::pair<std::string, std::string>> AuthenticateBasicToken(
      const FlightCallOptions& options, const std::string& username,
      const std::string& password) {
    // Add basic auth headers to outgoing headers.
    ClientRpc rpc(options);
    internal::AddBasicAuthHeaders(&rpc.context, username, password);

    std::shared_ptr<grpc::ClientReaderWriter<pb::HandshakeRequest, pb::HandshakeResponse>>
        stream = stub_->Handshake(&rpc.context);
    GrpcClientAuthSender outgoing{stream};
    GrpcClientAuthReader incoming{stream};

    // Explicitly close our side of the connection.
    bool finished_writes = stream->WritesDone();
    RETURN_NOT_OK(internal::FromGrpcStatus(stream->Finish(), &rpc.context));
    if (!finished_writes) {
      return MakeFlightError(FlightStatusCode::Internal,
                             "Could not finish writing before closing");
    }

    // Grab bearer token from incoming headers.
    return internal::GetBearerTokenHeader(rpc.context);
  }

  Status ListFlights(const FlightCallOptions& options, const Criteria& criteria,
                     std::unique_ptr<FlightListing>* listing) {
    pb::Criteria pb_criteria;
    RETURN_NOT_OK(internal::ToProto(criteria, &pb_criteria));

    ClientRpc rpc(options);
    RETURN_NOT_OK(rpc.SetToken(auth_handler_.get()));
    std::unique_ptr<grpc::ClientReader<pb::FlightInfo>> stream(
        stub_->ListFlights(&rpc.context, pb_criteria));

    std::vector<FlightInfo> flights;

    pb::FlightInfo pb_info;
    while (!options.stop_token.IsStopRequested() && stream->Read(&pb_info)) {
      FlightInfo::Data info_data;
      RETURN_NOT_OK(internal::FromProto(pb_info, &info_data));
      flights.emplace_back(std::move(info_data));
    }
    if (options.stop_token.IsStopRequested()) rpc.context.TryCancel();
    RETURN_NOT_OK(options.stop_token.Poll());
    listing->reset(new SimpleFlightListing(std::move(flights)));
    return internal::FromGrpcStatus(stream->Finish(), &rpc.context);
  }

  Status DoAction(const FlightCallOptions& options, const Action& action,
                  std::unique_ptr<ResultStream>* results) {
    pb::Action pb_action;
    RETURN_NOT_OK(internal::ToProto(action, &pb_action));

    ClientRpc rpc(options);
    RETURN_NOT_OK(rpc.SetToken(auth_handler_.get()));
    std::unique_ptr<grpc::ClientReader<pb::Result>> stream(
        stub_->DoAction(&rpc.context, pb_action));

    pb::Result pb_result;

    std::vector<Result> materialized_results;
    while (!options.stop_token.IsStopRequested() && stream->Read(&pb_result)) {
      Result result;
      RETURN_NOT_OK(internal::FromProto(pb_result, &result));
      materialized_results.emplace_back(std::move(result));
    }
    if (options.stop_token.IsStopRequested()) rpc.context.TryCancel();
    RETURN_NOT_OK(options.stop_token.Poll());

    *results = std::unique_ptr<ResultStream>(
        new SimpleResultStream(std::move(materialized_results)));
    return internal::FromGrpcStatus(stream->Finish(), &rpc.context);
  }

  Status ListActions(const FlightCallOptions& options, std::vector<ActionType>* types) {
    pb::Empty empty;

    ClientRpc rpc(options);
    RETURN_NOT_OK(rpc.SetToken(auth_handler_.get()));
    std::unique_ptr<grpc::ClientReader<pb::ActionType>> stream(
        stub_->ListActions(&rpc.context, empty));

    pb::ActionType pb_type;
    ActionType type;
    while (!options.stop_token.IsStopRequested() && stream->Read(&pb_type)) {
      RETURN_NOT_OK(internal::FromProto(pb_type, &type));
      types->emplace_back(std::move(type));
    }
    if (options.stop_token.IsStopRequested()) rpc.context.TryCancel();
    RETURN_NOT_OK(options.stop_token.Poll());
    return internal::FromGrpcStatus(stream->Finish(), &rpc.context);
  }

  Status GetFlightInfo(const FlightCallOptions& options,
                       const FlightDescriptor& descriptor,
                       std::unique_ptr<FlightInfo>* info) {
    pb::FlightDescriptor pb_descriptor;
    pb::FlightInfo pb_response;

    RETURN_NOT_OK(internal::ToProto(descriptor, &pb_descriptor));

    ClientRpc rpc(options);
    RETURN_NOT_OK(rpc.SetToken(auth_handler_.get()));
    Status s = internal::FromGrpcStatus(
        stub_->GetFlightInfo(&rpc.context, pb_descriptor, &pb_response), &rpc.context);
    RETURN_NOT_OK(s);

    FlightInfo::Data info_data;
    RETURN_NOT_OK(internal::FromProto(pb_response, &info_data));
    info->reset(new FlightInfo(std::move(info_data)));
    return Status::OK();
  }

  Status GetSchema(const FlightCallOptions& options, const FlightDescriptor& descriptor,
                   std::unique_ptr<SchemaResult>* schema_result) {
    pb::FlightDescriptor pb_descriptor;
    pb::SchemaResult pb_response;

    RETURN_NOT_OK(internal::ToProto(descriptor, &pb_descriptor));

    ClientRpc rpc(options);
    RETURN_NOT_OK(rpc.SetToken(auth_handler_.get()));
    Status s = internal::FromGrpcStatus(
        stub_->GetSchema(&rpc.context, pb_descriptor, &pb_response), &rpc.context);
    RETURN_NOT_OK(s);

    std::string str;
    RETURN_NOT_OK(internal::FromProto(pb_response, &str));
    schema_result->reset(new SchemaResult(str));
    return Status::OK();
  }

  Status DoGet(const FlightCallOptions& options, const Ticket& ticket,
               std::unique_ptr<FlightStreamReader>* out) {
    using StreamReader = GrpcStreamReader<grpc::ClientReader<pb::FlightData>>;
    pb::Ticket pb_ticket;
    internal::ToProto(ticket, &pb_ticket);

    auto rpc = std::make_shared<ClientRpc>(options);
    RETURN_NOT_OK(rpc->SetToken(auth_handler_.get()));
    std::shared_ptr<grpc::ClientReader<pb::FlightData>> stream =
        stub_->DoGet(&rpc->context, pb_ticket);
    auto finishable_stream = std::make_shared<
        FinishableStream<grpc::ClientReader<pb::FlightData>, internal::FlightData>>(
        rpc, stream);
    *out = std::unique_ptr<StreamReader>(new StreamReader(
        rpc, nullptr, options.read_options, options.stop_token, finishable_stream));
    // Eagerly read the schema
    return static_cast<StreamReader*>(out->get())->EnsureDataStarted();
  }

  Status DoPut(const FlightCallOptions& options, const FlightDescriptor& descriptor,
               const std::shared_ptr<Schema>& schema,
               std::unique_ptr<FlightStreamWriter>* out,
               std::unique_ptr<FlightMetadataReader>* reader) {
    using GrpcStream = grpc::ClientReaderWriter<pb::FlightData, pb::PutResult>;
    using StreamWriter = GrpcStreamWriter<pb::PutResult, pb::PutResult>;

    auto rpc = std::make_shared<ClientRpc>(options);
    RETURN_NOT_OK(rpc->SetToken(auth_handler_.get()));
    std::shared_ptr<GrpcStream> stream = stub_->DoPut(&rpc->context);
    // The writer drains the reader on close to avoid hanging inside
    // gRPC. Concurrent reads are unsafe, so a mutex protects this operation.
    std::shared_ptr<std::mutex> read_mutex = std::make_shared<std::mutex>();
    auto finishable_stream =
        std::make_shared<FinishableWritableStream<GrpcStream, pb::PutResult>>(
            rpc, read_mutex, stream);
    *reader =
        std::unique_ptr<FlightMetadataReader>(new GrpcMetadataReader(stream, read_mutex));
    return StreamWriter::Open(descriptor, schema, options.write_options, rpc,
                              write_size_limit_bytes_, finishable_stream, out);
  }

  Status DoExchange(const FlightCallOptions& options, const FlightDescriptor& descriptor,
                    std::unique_ptr<FlightStreamWriter>* writer,
                    std::unique_ptr<FlightStreamReader>* reader) {
    using GrpcStream = grpc::ClientReaderWriter<pb::FlightData, pb::FlightData>;
    using StreamReader = GrpcStreamReader<GrpcStream>;
    using StreamWriter = GrpcStreamWriter<pb::FlightData, internal::FlightData>;

    auto rpc = std::make_shared<ClientRpc>(options);
    RETURN_NOT_OK(rpc->SetToken(auth_handler_.get()));
    std::shared_ptr<grpc::ClientReaderWriter<pb::FlightData, pb::FlightData>> stream =
        stub_->DoExchange(&rpc->context);
    // The writer drains the reader on close to avoid hanging inside
    // gRPC. Concurrent reads are unsafe, so a mutex protects this operation.
    std::shared_ptr<std::mutex> read_mutex = std::make_shared<std::mutex>();
    auto finishable_stream =
        std::make_shared<FinishableWritableStream<GrpcStream, internal::FlightData>>(
            rpc, read_mutex, stream);
    *reader = std::unique_ptr<StreamReader>(new StreamReader(
        rpc, read_mutex, options.read_options, options.stop_token, finishable_stream));
    // Do not eagerly read the schema. There may be metadata messages
    // before any data is sent, or data may not be sent at all.
    return StreamWriter::Open(descriptor, nullptr, options.write_options, rpc,
                              write_size_limit_bytes_, finishable_stream, writer);
  }

 private:
  std::unique_ptr<pb::FlightService::Stub> stub_;
  std::shared_ptr<ClientAuthHandler> auth_handler_;
#if defined(GRPC_NAMESPACE_FOR_TLS_CREDENTIALS_OPTIONS)
  // Scope the TlsServerAuthorizationCheckConfig to be at the class instance level, since
  // it gets created during Connect() and needs to persist to DoAction() calls. gRPC does
  // not correctly increase the reference count of this object:
  // https://github.com/grpc/grpc/issues/22287
  std::shared_ptr<
      GRPC_NAMESPACE_FOR_TLS_CREDENTIALS_OPTIONS::TlsServerAuthorizationCheckConfig>
      noop_auth_check_;
#endif
  int64_t write_size_limit_bytes_;
};

FlightClient::FlightClient() { impl_.reset(new FlightClientImpl); }

FlightClient::~FlightClient() {}

Status FlightClient::Connect(const Location& location,
                             std::unique_ptr<FlightClient>* client) {
  return Connect(location, FlightClientOptions::Defaults(), client);
}

Status FlightClient::Connect(const Location& location, const FlightClientOptions& options,
                             std::unique_ptr<FlightClient>* client) {
  client->reset(new FlightClient);
  return (*client)->impl_->Connect(location, options);
}

Status FlightClient::Authenticate(const FlightCallOptions& options,
                                  std::unique_ptr<ClientAuthHandler> auth_handler) {
  return impl_->Authenticate(options, std::move(auth_handler));
}

arrow::Result<std::pair<std::string, std::string>> FlightClient::AuthenticateBasicToken(
    const FlightCallOptions& options, const std::string& username,
    const std::string& password) {
  return impl_->AuthenticateBasicToken(options, username, password);
}

Status FlightClient::DoAction(const FlightCallOptions& options, const Action& action,
                              std::unique_ptr<ResultStream>* results) {
  return impl_->DoAction(options, action, results);
}

Status FlightClient::ListActions(const FlightCallOptions& options,
                                 std::vector<ActionType>* actions) {
  return impl_->ListActions(options, actions);
}

Status FlightClient::GetFlightInfo(const FlightCallOptions& options,
                                   const FlightDescriptor& descriptor,
                                   std::unique_ptr<FlightInfo>* info) {
  return impl_->GetFlightInfo(options, descriptor, info);
}

Status FlightClient::GetSchema(const FlightCallOptions& options,
                               const FlightDescriptor& descriptor,
                               std::unique_ptr<SchemaResult>* schema_result) {
  return impl_->GetSchema(options, descriptor, schema_result);
}

Status FlightClient::ListFlights(std::unique_ptr<FlightListing>* listing) {
  return ListFlights({}, {}, listing);
}

Status FlightClient::ListFlights(const FlightCallOptions& options,
                                 const Criteria& criteria,
                                 std::unique_ptr<FlightListing>* listing) {
  return impl_->ListFlights(options, criteria, listing);
}

Status FlightClient::DoGet(const FlightCallOptions& options, const Ticket& ticket,
                           std::unique_ptr<FlightStreamReader>* stream) {
  return impl_->DoGet(options, ticket, stream);
}

Status FlightClient::DoPut(const FlightCallOptions& options,
                           const FlightDescriptor& descriptor,
                           const std::shared_ptr<Schema>& schema,
                           std::unique_ptr<FlightStreamWriter>* stream,
                           std::unique_ptr<FlightMetadataReader>* reader) {
  return impl_->DoPut(options, descriptor, schema, stream, reader);
}

Status FlightClient::DoExchange(const FlightCallOptions& options,
                                const FlightDescriptor& descriptor,
                                std::unique_ptr<FlightStreamWriter>* writer,
                                std::unique_ptr<FlightStreamReader>* reader) {
  return impl_->DoExchange(options, descriptor, writer, reader);
}

}  // namespace flight
}  // namespace arrow
