Implement application metadata in Flight

cpp/src/arrow/flight/client.cc

@@ -79,7 +79,7 @@ struct ClientRpc {
     if (auth_handler) {
       std::string token;
       RETURN_NOT_OK(auth_handler->GetToken(&token));
-      context.AddMetadata(internal::AUTH_HEADER, token);
+      context.AddMetadata(internal::kGrpcAuthHeader, token);
     }
     return Status::OK();
   }
@@ -129,15 +129,47 @@ class GrpcClientAuthReader : public ClientAuthReader {
       stream_;
 };
 
-class FlightIpcMessageReader : public ipc::MessageReader {
+// The next two classes are intertwined. To get the application
+// metadata while avoiding reimplementing RecordBatchStreamReader, we
+// create an ipc::MessageReader that is tied to the
+// MetadataRecordBatchReader. Every time an IPC message is read, it updates
+// the application metadata field of the MetadataRecordBatchReader. The
+// MetadataRecordBatchReader wraps RecordBatchStreamReader, offering an
+// additional method to get both the record batch and application
+// metadata.
+
+class GrpcIpcMessageReader;
+class GrpcStreamReader : public MetadataRecordBatchReader {
  public:
-  FlightIpcMessageReader(std::unique_ptr<ClientRpc> rpc,
-                         std::unique_ptr<grpc::ClientReader<pb::FlightData>> stream)
-      : rpc_(std::move(rpc)), stream_(std::move(stream)), stream_finished_(false) {}
+  GrpcStreamReader();
+
+  static Status Open(std::unique_ptr<ClientRpc> rpc,
+                     std::unique_ptr<grpc::ClientReader<pb::FlightData>> stream,
+                     std::unique_ptr<GrpcStreamReader>* out);
+  std::shared_ptr<Schema> schema() const override;
+  Status ReadNext(std::shared_ptr<RecordBatch>* out) override;
+  Status ReadWithMetadata(std::shared_ptr<RecordBatch>* out,
+                          std::shared_ptr<Buffer>* app_metadata) override;
+
+ private:
+  friend class GrpcIpcMessageReader;
+  std::unique_ptr<ipc::RecordBatchReader> batch_reader_;
+  std::shared_ptr<Buffer> last_app_metadata_;
+};
+
+class GrpcIpcMessageReader : public ipc::MessageReader {
+ public:
+  GrpcIpcMessageReader(GrpcStreamReader* reader, std::unique_ptr<ClientRpc> rpc,
+                       std::unique_ptr<grpc::ClientReader<pb::FlightData>> stream)
+      : flight_reader_(reader),
+        rpc_(std::move(rpc)),
+        stream_(std::move(stream)),
+        stream_finished_(false) {}
 
   Status ReadNextMessage(std::unique_ptr<ipc::Message>* out) override {
     if (stream_finished_) {
       *out = nullptr;
+      flight_reader_->last_app_metadata_ = nullptr;
       return Status::OK();
     }
     internal::FlightData data;
@@ -145,13 +177,16 @@ class FlightIpcMessageReader : public ipc::MessageReader {
       // Stream is completed
       stream_finished_ = true;
       *out = nullptr;
+      flight_reader_->last_app_metadata_ = nullptr;
       return OverrideWithServerError(Status::OK());
     }
     // Validate IPC message
     auto st = data.OpenMessage(out);
     if (!st.ok()) {
+      flight_reader_->last_app_metadata_ = nullptr;
       return OverrideWithServerError(std::move(st));
     }
+    flight_reader_->last_app_metadata_ = data.app_metadata;
     return Status::OK();
   }
 
@@ -162,23 +197,95 @@ class FlightIpcMessageReader : public ipc::MessageReader {
     return std::move(st);
   }
 
+ private:
+  GrpcStreamReader* flight_reader_;
   // The RPC context lifetime must be coupled to the ClientReader
   std::unique_ptr<ClientRpc> rpc_;
   std::unique_ptr<grpc::ClientReader<pb::FlightData>> stream_;
   bool stream_finished_;
 };
 
+GrpcStreamReader::GrpcStreamReader() {}
+
+Status GrpcStreamReader::Open(std::unique_ptr<ClientRpc> rpc,
+                              std::unique_ptr<grpc::ClientReader<pb::FlightData>> stream,
+                              std::unique_ptr<GrpcStreamReader>* out) {
+  *out = std::unique_ptr<GrpcStreamReader>(new GrpcStreamReader);
+  std::unique_ptr<GrpcIpcMessageReader> message_reader(
+      new GrpcIpcMessageReader(out->get(), std::move(rpc), std::move(stream)));
+  return ipc::RecordBatchStreamReader::Open(std::move(message_reader),
+                                            &(*out)->batch_reader_);
+}
+
+std::shared_ptr<Schema> GrpcStreamReader::schema() const {
+  return batch_reader_->schema();
+}
+
+Status GrpcStreamReader::ReadNext(std::shared_ptr<RecordBatch>* out) {
+  std::shared_ptr<Buffer> app_metadata;
+  return ReadWithMetadata(out, &app_metadata);
+}
+
+Status GrpcStreamReader::ReadWithMetadata(std::shared_ptr<RecordBatch>* out,
+                                          std::shared_ptr<Buffer>* app_metadata) {
+  RETURN_NOT_OK(batch_reader_->ReadNext(out));
+  *app_metadata = std::move(last_app_metadata_);
+  return Status::OK();
+}
+
+// Similarly, the next two classes are intertwined. 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.
+
+class DoPutPayloadWriter;
+class GrpcStreamWriter : public FlightStreamWriter {
+ public:
+  ~GrpcStreamWriter() = default;
+
+  GrpcStreamWriter() : app_metadata_(nullptr), batch_writer_(nullptr) {}
+
+  static Status Open(
+      const FlightDescriptor& descriptor, const std::shared_ptr<Schema>& schema,
+      std::unique_ptr<ClientRpc> rpc, std::unique_ptr<pb::PutResult> response,
+      std::shared_ptr<grpc::ClientReaderWriter<pb::FlightData, pb::PutResult>> writer,
+      std::unique_ptr<FlightStreamWriter>* out);
+
+  Status WriteRecordBatch(const RecordBatch& batch, bool allow_64bit = false) override {
+    return WriteWithMetadata(batch, nullptr, allow_64bit);
+  }
+  Status WriteWithMetadata(const RecordBatch& batch, std::shared_ptr<Buffer> app_metadata,
+                           bool allow_64bit = false) override {
+    app_metadata_ = app_metadata;
+    return batch_writer_->WriteRecordBatch(batch, allow_64bit);
+  }
+  void set_memory_pool(MemoryPool* pool) override {
+    batch_writer_->set_memory_pool(pool);
+  }
+  Status Close() override { return batch_writer_->Close(); }
+
+ private:
+  friend class DoPutPayloadWriter;
+  std::shared_ptr<Buffer> app_metadata_;
+  std::unique_ptr<ipc::RecordBatchWriter> batch_writer_;
+};
+
 /// A IpcPayloadWriter implementation that writes to a DoPut stream
 class DoPutPayloadWriter : public ipc::internal::IpcPayloadWriter {
  public:
-  DoPutPayloadWriter(const FlightDescriptor& descriptor, std::unique_ptr<ClientRpc> rpc,
-                     std::unique_ptr<protocol::PutResult> response,
-                     std::unique_ptr<grpc::ClientWriter<pb::FlightData>> writer)
+  DoPutPayloadWriter(
+      const FlightDescriptor& descriptor, std::unique_ptr<ClientRpc> rpc,
+      std::unique_ptr<pb::PutResult> response,
+      std::shared_ptr<grpc::ClientReaderWriter<pb::FlightData, pb::PutResult>> writer,
+      GrpcStreamWriter* stream_writer)
       : descriptor_(descriptor),
         rpc_(std::move(rpc)),
         response_(std::move(response)),
         writer_(std::move(writer)),
-        first_payload_(true) {}
+        first_payload_(true),
+        stream_writer_(stream_writer) {}
 
   ~DoPutPayloadWriter() override = default;
 
@@ -201,6 +308,8 @@ class DoPutPayloadWriter : public ipc::internal::IpcPayloadWriter {
       }
       RETURN_NOT_OK(Buffer::FromString(str_descr, &payload.descriptor));
       first_payload_ = false;
+    } else if (stream_writer_->app_metadata_) {
+      payload.app_metadata = std::move(stream_writer_->app_metadata_);
     }
 
     if (!internal::WritePayload(payload, writer_.get())) {
@@ -211,6 +320,10 @@ class DoPutPayloadWriter : public ipc::internal::IpcPayloadWriter {
 
   Status Close() override {
     bool finished_writes = writer_->WritesDone();
+    // Drain the read side to avoid hanging
+    pb::PutResult message;
+    while (writer_->Read(&message)) {
+    }
     RETURN_NOT_OK(internal::FromGrpcStatus(writer_->Finish()));
     if (!finished_writes) {
       return Status::UnknownError(
@@ -223,9 +336,47 @@ class DoPutPayloadWriter : public ipc::internal::IpcPayloadWriter {
   // TODO: there isn't a way to access this as a user.
   const FlightDescriptor descriptor_;
   std::unique_ptr<ClientRpc> rpc_;
-  std::unique_ptr<protocol::PutResult> response_;
-  std::unique_ptr<grpc::ClientWriter<pb::FlightData>> writer_;
+  std::unique_ptr<pb::PutResult> response_;
+  std::shared_ptr<grpc::ClientReaderWriter<pb::FlightData, pb::PutResult>> writer_;
   bool first_payload_;
+  GrpcStreamWriter* stream_writer_;
+};
+
+Status GrpcStreamWriter::Open(
+    const FlightDescriptor& descriptor, const std::shared_ptr<Schema>& schema,
+    std::unique_ptr<ClientRpc> rpc, std::unique_ptr<pb::PutResult> response,
+    std::shared_ptr<grpc::ClientReaderWriter<pb::FlightData, pb::PutResult>> writer,
+    std::unique_ptr<FlightStreamWriter>* out) {
+  std::unique_ptr<GrpcStreamWriter> result(new GrpcStreamWriter);
+  std::unique_ptr<ipc::internal::IpcPayloadWriter> payload_writer(new DoPutPayloadWriter(
+      descriptor, std::move(rpc), std::move(response), writer, result.get()));
+  RETURN_NOT_OK(ipc::internal::OpenRecordBatchWriter(std::move(payload_writer), schema,
+                                                     &result->batch_writer_));
+  *out = std::move(result);
+  return Status::OK();
+}
+
+FlightMetadataReader::~FlightMetadataReader() = default;
+
+class GrpcMetadataReader : public FlightMetadataReader {
+ public:
+  explicit GrpcMetadataReader(
+      std::shared_ptr<grpc::ClientReaderWriter<pb::FlightData, pb::PutResult>> reader)
+      : reader_(reader) {}
+
+  Status ReadMetadata(std::shared_ptr<Buffer>* out) override {
+    pb::PutResult message;
+    if (reader_->Read(&message)) {
+      *out = Buffer::FromString(std::move(*message.release_app_metadata()));
+    } else {
+      // Stream finished
+      *out = nullptr;
+    }
+    return Status::OK();
+  }
+
+ private:
+  std::shared_ptr<grpc::ClientReaderWriter<pb::FlightData, pb::PutResult>> reader_;
 };
 
 class FlightClient::FlightClientImpl {
@@ -367,7 +518,7 @@ class FlightClient::FlightClientImpl {
   }
 
   Status DoGet(const FlightCallOptions& options, const Ticket& ticket,
-               std::unique_ptr<RecordBatchReader>* out) {
+               std::unique_ptr<MetadataRecordBatchReader>* out) {
     pb::Ticket pb_ticket;
     internal::ToProto(ticket, &pb_ticket);
 
@@ -376,25 +527,25 @@ class FlightClient::FlightClientImpl {
     std::unique_ptr<grpc::ClientReader<pb::FlightData>> stream(
         stub_->DoGet(&rpc->context, pb_ticket));
 
-    std::unique_ptr<ipc::MessageReader> message_reader(
-        new FlightIpcMessageReader(std::move(rpc), std::move(stream)));
-    return ipc::RecordBatchStreamReader::Open(std::move(message_reader), out);
+    std::unique_ptr<GrpcStreamReader> reader;
+    RETURN_NOT_OK(GrpcStreamReader::Open(std::move(rpc), std::move(stream), &reader));
+    *out = std::move(reader);
+    return Status::OK();
   }
 
   Status DoPut(const FlightCallOptions& options, const FlightDescriptor& descriptor,
                const std::shared_ptr<Schema>& schema,
-               std::unique_ptr<ipc::RecordBatchWriter>* out) {
+               std::unique_ptr<FlightStreamWriter>* out,
+               std::unique_ptr<FlightMetadataReader>* reader) {
     std::unique_ptr<ClientRpc> rpc(new ClientRpc(options));
     RETURN_NOT_OK(rpc->SetToken(auth_handler_.get()));
-    std::unique_ptr<protocol::PutResult> response(new protocol::PutResult);
-    std::unique_ptr<grpc::ClientWriter<pb::FlightData>> writer(
-        stub_->DoPut(&rpc->context, response.get()));
+    std::unique_ptr<pb::PutResult> response(new pb::PutResult);
+    std::shared_ptr<grpc::ClientReaderWriter<pb::FlightData, pb::PutResult>> writer(
+        stub_->DoPut(&rpc->context));
 
-    std::unique_ptr<ipc::internal::IpcPayloadWriter> payload_writer(
-        new DoPutPayloadWriter(descriptor, std::move(rpc), std::move(response),
-                               std::move(writer)));
-
-    return ipc::internal::OpenRecordBatchWriter(std::move(payload_writer), schema, out);
+    *reader = std::unique_ptr<FlightMetadataReader>(new GrpcMetadataReader(writer));
+    return GrpcStreamWriter::Open(descriptor, schema, std::move(rpc), std::move(response),
+                                  writer, out);
   }
 
  private:
@@ -449,15 +600,16 @@ Status FlightClient::ListFlights(const FlightCallOptions& options,
 }
 
 Status FlightClient::DoGet(const FlightCallOptions& options, const Ticket& ticket,
-                           std::unique_ptr<RecordBatchReader>* stream) {
+                           std::unique_ptr<MetadataRecordBatchReader>* 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<ipc::RecordBatchWriter>* stream) {
-  return impl_->DoPut(options, descriptor, schema, stream);
+                           std::unique_ptr<FlightStreamWriter>* stream,
+                           std::unique_ptr<FlightMetadataReader>* reader) {
+  return impl_->DoPut(options, descriptor, schema, stream, reader);
 }
 
 }  // namespace flight

cpp/src/arrow/flight/client.h

@@ -25,6 +25,7 @@
 #include <string>
 #include <vector>
 
+#include "arrow/ipc/reader.h"
 #include "arrow/ipc/writer.h"
 #include "arrow/status.h"
 
@@ -35,7 +36,6 @@ namespace arrow {
 
 class MemoryPool;
 class RecordBatch;
-class RecordBatchReader;
 class Schema;
 
 namespace flight {
@@ -66,6 +66,24 @@ class ARROW_FLIGHT_EXPORT FlightClientOptions {
   std::string override_hostname;
 };
 
+/// \brief A RecordBatchWriter that also allows sending
+/// application-defined metadata via the Flight protocol.
+class ARROW_EXPORT FlightStreamWriter : public ipc::RecordBatchWriter {
+ public:
+  virtual Status WriteWithMetadata(const RecordBatch& batch,
+                                   std::shared_ptr<Buffer> app_metadata,
+                                   bool allow_64bit = false) = 0;
+};
+
+/// \brief A reader for application-specific metadata sent back to the
+/// client during an upload.
+class ARROW_EXPORT FlightMetadataReader {
+ public:
+  virtual ~FlightMetadataReader();
+  /// \brief Read a message from the server.
+  virtual Status ReadMetadata(std::shared_ptr<Buffer>* out) = 0;
+};
+
 /// \brief Client class for Arrow Flight RPC services (gRPC-based).
 /// API experimental for now
 class ARROW_FLIGHT_EXPORT FlightClient {
@@ -151,8 +169,8 @@ class ARROW_FLIGHT_EXPORT FlightClient {
   /// \param[out] stream the returned RecordBatchReader
   /// \return Status
   Status DoGet(const FlightCallOptions& options, const Ticket& ticket,
-               std::unique_ptr<RecordBatchReader>* stream);
-  Status DoGet(const Ticket& ticket, std::unique_ptr<RecordBatchReader>* stream) {
+               std::unique_ptr<MetadataRecordBatchReader>* stream);
+  Status DoGet(const Ticket& ticket, std::unique_ptr<MetadataRecordBatchReader>* stream) {
     return DoGet({}, ticket, stream);
   }
 
@@ -163,13 +181,16 @@ class ARROW_FLIGHT_EXPORT FlightClient {
   /// \param[in] descriptor the descriptor of the stream
   /// \param[in] schema the schema for the data to upload
   /// \param[out] stream a writer to write record batches to
+  /// \param[out] reader a reader for application metadata from the server
   /// \return Status
   Status DoPut(const FlightCallOptions& options, const FlightDescriptor& descriptor,
                const std::shared_ptr<Schema>& schema,
-               std::unique_ptr<ipc::RecordBatchWriter>* stream);
+               std::unique_ptr<FlightStreamWriter>* stream,
+               std::unique_ptr<FlightMetadataReader>* reader);
   Status DoPut(const FlightDescriptor& descriptor, const std::shared_ptr<Schema>& schema,
-               std::unique_ptr<ipc::RecordBatchWriter>* stream) {
-    return DoPut({}, descriptor, schema, stream);
+               std::unique_ptr<FlightStreamWriter>* stream,
+               std::unique_ptr<FlightMetadataReader>* reader) {
+    return DoPut({}, descriptor, schema, stream, reader);
   }
 
  private:

cpp/src/arrow/flight/flight-benchmark.cc

@@ -106,7 +106,7 @@ Status RunPerformanceTest(const std::string& hostname, const int port) {
     perf::Token token;
     token.ParseFromString(endpoint.ticket.ticket);
 
-    std::unique_ptr<RecordBatchReader> reader;
+    std::unique_ptr<MetadataRecordBatchReader> reader;
     RETURN_NOT_OK(client->DoGet(endpoint.ticket, &reader));
 
     std::shared_ptr<RecordBatch> batch;

cpp/src/arrow/flight/flight-test.cc

@@ -211,7 +211,7 @@ class TestFlightClient : public ::testing::Test {
 
     // By convention, fetch the first endpoint
     Ticket ticket = info->endpoints()[0].ticket;
-    std::unique_ptr<RecordBatchReader> stream;
+    std::unique_ptr<MetadataRecordBatchReader> stream;
     ASSERT_OK(client_->DoGet(ticket, &stream));
 
     std::shared_ptr<RecordBatch> chunk;
@@ -255,7 +255,8 @@ class TlsTestServer : public FlightServerBase {
 class DoPutTestServer : public FlightServerBase {
  public:
   Status DoPut(const ServerCallContext& context,
-               std::unique_ptr<FlightMessageReader> reader) override {
+               std::unique_ptr<FlightMessageReader> reader,
+               std::unique_ptr<FlightMetadataWriter> writer) override {
     descriptor_ = reader->descriptor();
     return reader->ReadAll(&batches_);
   }
@@ -267,6 +268,71 @@ class DoPutTestServer : public FlightServerBase {
   friend class TestDoPut;
 };
 
+class MetadataTestServer : public FlightServerBase {
+  Status DoGet(const ServerCallContext& context, const Ticket& request,
+               std::unique_ptr<FlightDataStream>* data_stream) override {
+    BatchVector batches;
+    RETURN_NOT_OK(ExampleIntBatches(&batches));
+    std::shared_ptr<RecordBatchReader> batch_reader =
+        std::make_shared<BatchIterator>(batches[0]->schema(), batches);
+
+    *data_stream = std::unique_ptr<FlightDataStream>(new NumberingStream(
+        std::unique_ptr<FlightDataStream>(new RecordBatchStream(batch_reader))));
+    return Status::OK();
+  }
+
+  Status DoPut(const ServerCallContext& context,
+               std::unique_ptr<FlightMessageReader> reader,
+               std::unique_ptr<FlightMetadataWriter> writer) override {
+    std::shared_ptr<arrow::RecordBatch> chunk;
+    std::shared_ptr<Buffer> app_metadata;
+    int counter = 0;
+    while (true) {
+      RETURN_NOT_OK(reader->ReadWithMetadata(&chunk, &app_metadata));
+      if (chunk == nullptr) break;
+      if (app_metadata == nullptr) {
+        return Status::Invalid("Expected application metadata to be provided");
+      }
+      if (std::to_string(counter) != app_metadata->ToString()) {
+        return Status::Invalid("Expected metadata value: " + std::to_string(counter) +
+                               " but got: " + app_metadata->ToString());
+      }
+      auto metadata = Buffer::FromString(std::to_string(counter));
+      RETURN_NOT_OK(writer->WriteMetadata(*metadata));
+      counter++;
+    }
+    return Status::OK();
+  }
+};
+
+template <typename T>
+class InsecureTestServer : public ::testing::Test {
+ public:
+  void SetUp() {
+    Location location;
+    ASSERT_OK(Location::ForGrpcTcp("localhost", 30000, &location));
+
+    std::unique_ptr<FlightServerBase> server(new T);
+    FlightServerOptions options(location);
+    ASSERT_OK(server->Init(options));
+
+    server_.reset(new InProcessTestServer(std::move(server), location));
+    ASSERT_OK(server_->Start());
+    ASSERT_OK(ConnectClient());
+  }
+
+  void TearDown() { server_->Stop(); }
+
+  Status ConnectClient() { return FlightClient::Connect(server_->location(), &client_); }
+
+ protected:
+  int port_;
+  std::unique_ptr<FlightClient> client_;
+  std::unique_ptr<InProcessTestServer> server_;
+};
+
+using TestMetadata = InsecureTestServer<MetadataTestServer>;
+
 class TestAuthHandler : public ::testing::Test {
  public:
   void SetUp() {
@@ -323,8 +389,9 @@ class TestDoPut : public ::testing::Test {
 
   void CheckDoPut(FlightDescriptor descr, const std::shared_ptr<Schema>& schema,
                   const BatchVector& batches) {
-    std::unique_ptr<ipc::RecordBatchWriter> stream;
-    ASSERT_OK(client_->DoPut(descr, schema, &stream));
+    std::unique_ptr<FlightStreamWriter> stream;
+    std::unique_ptr<FlightMetadataReader> reader;
+    ASSERT_OK(client_->DoPut(descr, schema, &stream, &reader));
     for (const auto& batch : batches) {
       ASSERT_OK(stream->WriteRecordBatch(*batch));
     }
@@ -485,7 +552,7 @@ TEST_F(TestFlightClient, Issue5095) {
   // Make sure the server-side error message is reflected to the
   // client
   Ticket ticket1{"ARROW-5095-fail"};
-  std::unique_ptr<RecordBatchReader> stream;
+  std::unique_ptr<MetadataRecordBatchReader> stream;
   Status status = client_->DoGet(ticket1, &stream);
   ASSERT_RAISES(IOError, status);
   ASSERT_THAT(status.message(), ::testing::HasSubstr("Server-side error"));
@@ -588,13 +655,14 @@ TEST_F(TestAuthHandler, PassAuthenticatedCalls) {
   status = client_->GetFlightInfo(FlightDescriptor{}, &info);
   ASSERT_RAISES(NotImplemented, status);
 
-  std::unique_ptr<RecordBatchReader> stream;
+  std::unique_ptr<MetadataRecordBatchReader> stream;
   status = client_->DoGet(Ticket{}, &stream);
   ASSERT_RAISES(NotImplemented, status);
 
-  std::unique_ptr<ipc::RecordBatchWriter> writer;
+  std::unique_ptr<FlightStreamWriter> writer;
+  std::unique_ptr<FlightMetadataReader> reader;
   std::shared_ptr<Schema> schema = arrow::schema({});
-  status = client_->DoPut(FlightDescriptor{}, schema, &writer);
+  status = client_->DoPut(FlightDescriptor{}, schema, &writer, &reader);
   ASSERT_OK(status);
   status = writer->Close();
   ASSERT_RAISES(NotImplemented, status);
@@ -625,15 +693,16 @@ TEST_F(TestAuthHandler, FailUnauthenticatedCalls) {
   ASSERT_RAISES(IOError, status);
   ASSERT_THAT(status.message(), ::testing::HasSubstr("Invalid token"));
 
-  std::unique_ptr<RecordBatchReader> stream;
+  std::unique_ptr<MetadataRecordBatchReader> stream;
   status = client_->DoGet(Ticket{}, &stream);
   ASSERT_RAISES(IOError, status);
   ASSERT_THAT(status.message(), ::testing::HasSubstr("Invalid token"));
 
-  std::unique_ptr<ipc::RecordBatchWriter> writer;
+  std::unique_ptr<FlightStreamWriter> writer;
+  std::unique_ptr<FlightMetadataReader> reader;
   std::shared_ptr<Schema> schema(
       (new arrow::Schema(std::vector<std::shared_ptr<Field>>())));
-  status = client_->DoPut(FlightDescriptor{}, schema, &writer);
+  status = client_->DoPut(FlightDescriptor{}, schema, &writer, &reader);
   ASSERT_OK(status);
   status = writer->Close();
   ASSERT_RAISES(IOError, status);
@@ -693,5 +762,73 @@ TEST_F(TestTls, OverrideHostname) {
   ASSERT_RAISES(IOError, client->DoAction(options, action, &results));
 }
 
+TEST_F(TestMetadata, DoGet) {
+  Ticket ticket{""};
+  std::unique_ptr<MetadataRecordBatchReader> stream;
+  ASSERT_OK(client_->DoGet(ticket, &stream));
+
+  BatchVector expected_batches;
+  ASSERT_OK(ExampleIntBatches(&expected_batches));
+
+  std::shared_ptr<RecordBatch> chunk;
+  std::shared_ptr<Buffer> metadata;
+  auto num_batches = static_cast<int>(expected_batches.size());
+  for (int i = 0; i < num_batches; ++i) {
+    ASSERT_OK(stream->ReadWithMetadata(&chunk, &metadata));
+    ASSERT_NE(nullptr, chunk);
+    ASSERT_NE(nullptr, metadata);
+    ASSERT_BATCHES_EQUAL(*expected_batches[i], *chunk);
+    ASSERT_EQ(std::to_string(i), metadata->ToString());
+  }
+  ASSERT_OK(stream->ReadNext(&chunk));
+  ASSERT_EQ(nullptr, chunk);
+}
+
+TEST_F(TestMetadata, DoPut) {
+  std::unique_ptr<FlightStreamWriter> writer;
+  std::unique_ptr<FlightMetadataReader> reader;
+  std::shared_ptr<Schema> schema = ExampleIntSchema();
+  ASSERT_OK(client_->DoPut(FlightDescriptor{}, schema, &writer, &reader));
+
+  BatchVector expected_batches;
+  ASSERT_OK(ExampleIntBatches(&expected_batches));
+
+  std::shared_ptr<RecordBatch> chunk;
+  std::shared_ptr<Buffer> metadata;
+  auto num_batches = static_cast<int>(expected_batches.size());
+  for (int i = 0; i < num_batches; ++i) {
+    ASSERT_OK(writer->WriteWithMetadata(*expected_batches[i],
+                                        Buffer::FromString(std::to_string(i))));
+  }
+  // This eventually calls grpc::ClientReaderWriter::Finish which can
+  // hang if there are unread messages. So make sure our wrapper
+  // around this doesn't hang (because it drains any unread messages)
+  ASSERT_OK(writer->Close());
+}
+
+TEST_F(TestMetadata, DoPutReadMetadata) {
+  std::unique_ptr<FlightStreamWriter> writer;
+  std::unique_ptr<FlightMetadataReader> reader;
+  std::shared_ptr<Schema> schema = ExampleIntSchema();
+  ASSERT_OK(client_->DoPut(FlightDescriptor{}, schema, &writer, &reader));
+
+  BatchVector expected_batches;
+  ASSERT_OK(ExampleIntBatches(&expected_batches));
+
+  std::shared_ptr<RecordBatch> chunk;
+  std::shared_ptr<Buffer> metadata;
+  auto num_batches = static_cast<int>(expected_batches.size());
+  for (int i = 0; i < num_batches; ++i) {
+    ASSERT_OK(writer->WriteWithMetadata(*expected_batches[i],
+                                        Buffer::FromString(std::to_string(i))));
+    ASSERT_OK(reader->ReadMetadata(&metadata));
+    ASSERT_NE(nullptr, metadata);
+    ASSERT_EQ(std::to_string(i), metadata->ToString());
+  }
+  // As opposed to DoPutDrainMetadata, now we've read the messages, so
+  // make sure this still closes as expected.
+  ASSERT_OK(writer->Close());
+}
+
 }  // namespace flight
 }  // namespace arrow

cpp/src/arrow/flight/internal.h

@@ -63,7 +63,8 @@ namespace flight {
 
 namespace internal {
 
-static const char* AUTH_HEADER = "auth-token-bin";
+/// The name of the header used to pass authentication tokens.
+static const char* kGrpcAuthHeader = "auth-token-bin";
 
 ARROW_FLIGHT_EXPORT
 Status SchemaToString(const Schema& schema, std::string* out);

cpp/src/arrow/flight/serialization-internal.cc

@@ -163,6 +163,14 @@ grpc::Status FlightDataSerialize(const FlightPayload& msg, ByteBuffer* out,
   // 1 byte for metadata tag
   header_size += 1 + WireFormatLite::LengthDelimitedSize(metadata_size);
 
+  // App metadata tag if appropriate
+  int32_t app_metadata_size = 0;
+  if (msg.app_metadata && msg.app_metadata->size() > 0) {
+    DCHECK_LT(msg.app_metadata->size(), kInt32Max);
+    app_metadata_size = static_cast<int32_t>(msg.app_metadata->size());
+    header_size += 1 + WireFormatLite::LengthDelimitedSize(app_metadata_size);
+  }
+
   for (const auto& buffer : ipc_msg.body_buffers) {
     // Buffer may be null when the row length is zero, or when all
     // entries are invalid.
@@ -214,6 +222,15 @@ grpc::Status FlightDataSerialize(const FlightPayload& msg, ByteBuffer* out,
   header_stream.WriteRawMaybeAliased(ipc_msg.metadata->data(),
                                      static_cast<int>(ipc_msg.metadata->size()));
 
+  // Write app metadata
+  if (app_metadata_size > 0) {
+    WireFormatLite::WriteTag(pb::FlightData::kAppMetadataFieldNumber,
+                             WireFormatLite::WIRETYPE_LENGTH_DELIMITED, &header_stream);
+    header_stream.WriteVarint32(app_metadata_size);
+    header_stream.WriteRawMaybeAliased(msg.app_metadata->data(),
+                                       static_cast<int>(msg.app_metadata->size()));
+  }
+
   if (has_body) {
     // Write body tag
     WireFormatLite::WriteTag(pb::FlightData::kDataBodyFieldNumber,
@@ -292,6 +309,12 @@ grpc::Status FlightDataDeserialize(ByteBuffer* buffer, FlightData* out) {
                               "Unable to read FlightData metadata");
         }
       } break;
+      case pb::FlightData::kAppMetadataFieldNumber: {
+        if (!ReadBytesZeroCopy(wrapped_buffer, &pb_stream, &out->app_metadata)) {
+          return grpc::Status(grpc::StatusCode::INTERNAL,
+                              "Unable to read FlightData application metadata");
+        }
+      } break;
       case pb::FlightData::kDataBodyFieldNumber: {
         if (!ReadBytesZeroCopy(wrapped_buffer, &pb_stream, &out->body)) {
           return grpc::Status(grpc::StatusCode::INTERNAL,
@@ -330,7 +353,7 @@ Status FlightData::OpenMessage(std::unique_ptr<ipc::Message>* message) {
 // (see customize_protobuf.h).
 
 bool WritePayload(const FlightPayload& payload,
-                  grpc::ClientWriter<pb::FlightData>* writer) {
+                  grpc::ClientReaderWriter<pb::FlightData, pb::PutResult>* writer) {
   // Pretend to be pb::FlightData and intercept in SerializationTraits
   return writer->Write(*reinterpret_cast<const pb::FlightData*>(&payload),
                        grpc::WriteOptions());
@@ -348,7 +371,8 @@ bool ReadPayload(grpc::ClientReader<pb::FlightData>* reader, FlightData* data) {
   return reader->Read(reinterpret_cast<pb::FlightData*>(data));
 }
 
-bool ReadPayload(grpc::ServerReader<pb::FlightData>* reader, FlightData* data) {
+bool ReadPayload(grpc::ServerReaderWriter<pb::PutResult, pb::FlightData>* reader,
+                 FlightData* data) {
   // Pretend to be pb::FlightData and intercept in SerializationTraits
   return reader->Read(reinterpret_cast<pb::FlightData*>(data));
 }

cpp/src/arrow/flight/serialization-internal.h

@@ -43,6 +43,9 @@ struct FlightData {
   /// Non-length-prefixed Message header as described in format/Message.fbs
   std::shared_ptr<Buffer> metadata;
 
+  /// Application-defined metadata
+  std::shared_ptr<Buffer> app_metadata;
+
   /// Message body
   std::shared_ptr<Buffer> body;
 
@@ -53,14 +56,15 @@ struct FlightData {
 /// Write Flight message on gRPC stream with zero-copy optimizations.
 /// True is returned on success, false if some error occurred (connection closed?).
 bool WritePayload(const FlightPayload& payload,
-                  grpc::ClientWriter<pb::FlightData>* writer);
+                  grpc::ClientReaderWriter<pb::FlightData, pb::PutResult>* writer);
 bool WritePayload(const FlightPayload& payload,
                   grpc::ServerWriter<pb::FlightData>* writer);
 
 /// Read Flight message from gRPC stream with zero-copy optimizations.
 /// True is returned on success, false if stream ended.
 bool ReadPayload(grpc::ClientReader<pb::FlightData>* reader, FlightData* data);
-bool ReadPayload(grpc::ServerReader<pb::FlightData>* reader, FlightData* data);
+bool ReadPayload(grpc::ServerReaderWriter<pb::PutResult, pb::FlightData>* reader,
+                 FlightData* data);
 
 }  // namespace internal
 }  // namespace flight

cpp/src/arrow/flight/server.cc

@@ -72,12 +72,15 @@ namespace {
 // A MessageReader implementation that reads from a gRPC ServerReader
 class FlightIpcMessageReader : public ipc::MessageReader {
  public:
-  explicit FlightIpcMessageReader(grpc::ServerReader<pb::FlightData>* reader)
-      : reader_(reader) {}
+  explicit FlightIpcMessageReader(
+      grpc::ServerReaderWriter<pb::PutResult, pb::FlightData>* reader,
+      std::shared_ptr<Buffer>* last_metadata)
+      : reader_(reader), app_metadata_(last_metadata) {}
 
   Status ReadNextMessage(std::unique_ptr<ipc::Message>* out) override {
     if (stream_finished_) {
       *out = nullptr;
+      *app_metadata_ = nullptr;
       return Status::OK();
     }
     internal::FlightData data;
@@ -89,6 +92,7 @@ class FlightIpcMessageReader : public ipc::MessageReader {
             "Client provided malformed message or did not provide message");
       }
       *out = nullptr;
+      *app_metadata_ = nullptr;
       return Status::OK();
     }
 
@@ -100,25 +104,29 @@ class FlightIpcMessageReader : public ipc::MessageReader {
       first_message_ = false;
     }
 
-    return data.OpenMessage(out);
+    RETURN_NOT_OK(data.OpenMessage(out));
+    *app_metadata_ = std::move(data.app_metadata);
+    return Status::OK();
   }
 
   const FlightDescriptor& descriptor() const { return descriptor_; }
 
  protected:
-  grpc::ServerReader<pb::FlightData>* reader_;
+  grpc::ServerReaderWriter<pb::PutResult, pb::FlightData>* reader_;
   bool stream_finished_ = false;
   bool first_message_ = true;
   FlightDescriptor descriptor_;
+  std::shared_ptr<Buffer>* app_metadata_;
 };
 
 class FlightMessageReaderImpl : public FlightMessageReader {
  public:
-  explicit FlightMessageReaderImpl(grpc::ServerReader<pb::FlightData>* reader)
+  explicit FlightMessageReaderImpl(
+      grpc::ServerReaderWriter<pb::PutResult, pb::FlightData>* reader)
       : reader_(reader) {}
 
   Status Init() {
-    message_reader_ = new FlightIpcMessageReader(reader_);
+    message_reader_ = new FlightIpcMessageReader(reader_, &last_metadata_);
     return ipc::RecordBatchStreamReader::Open(
         std::unique_ptr<ipc::MessageReader>(message_reader_), &batch_reader_);
   }
@@ -130,17 +138,49 @@ class FlightMessageReaderImpl : public FlightMessageReader {
   std::shared_ptr<Schema> schema() const override { return batch_reader_->schema(); }
 
   Status ReadNext(std::shared_ptr<RecordBatch>* out) override {
-    return batch_reader_->ReadNext(out);
+    return ReadWithMetadata(out, nullptr);
+  }
+
+  Status ReadWithMetadata(std::shared_ptr<RecordBatch>* out,
+                          std::shared_ptr<Buffer>* app_metadata) override {
+    if (app_metadata) {
+      *app_metadata = nullptr;
+    }
+    RETURN_NOT_OK(batch_reader_->ReadNext(out));
+    if (app_metadata) {
+      *app_metadata = std::move(last_metadata_);
+    }
+    return Status::OK();
   }
 
  private:
   std::shared_ptr<Schema> schema_;
   std::unique_ptr<ipc::DictionaryMemo> dictionary_memo_;
-  grpc::ServerReader<pb::FlightData>* reader_;
+  grpc::ServerReaderWriter<pb::PutResult, pb::FlightData>* reader_;
   FlightIpcMessageReader* message_reader_;
+  std::shared_ptr<Buffer> last_metadata_;
   std::shared_ptr<RecordBatchReader> batch_reader_;
 };
 
+class GrpcMetadataWriter : public FlightMetadataWriter {
+ public:
+  explicit GrpcMetadataWriter(
+      grpc::ServerReaderWriter<pb::PutResult, pb::FlightData>* writer)
+      : writer_(writer) {}
+
+  Status WriteMetadata(const Buffer& buffer) override {
+    pb::PutResult message{};
+    message.set_app_metadata(buffer.data(), buffer.size());
+    if (writer_->Write(message)) {
+      return Status::OK();
+    }
+    return Status::IOError("Unknown error writing metadata.");
+  }
+
+ private:
+  grpc::ServerReaderWriter<pb::PutResult, pb::FlightData>* writer_;
+};
+
 class GrpcServerAuthReader : public ServerAuthReader {
  public:
   explicit GrpcServerAuthReader(
@@ -153,7 +193,7 @@ class GrpcServerAuthReader : public ServerAuthReader {
       *token = std::move(*request.release_payload());
       return Status::OK();
     }
-    return Status::UnknownError("Could not read client handshake request.");
+    return Status::IOError("Stream is closed.");
   }
 
  private:
@@ -246,7 +286,7 @@ class FlightServiceImpl : public FlightService::Service {
     }
 
     const auto client_metadata = context->client_metadata();
-    const auto auth_header = client_metadata.find(internal::AUTH_HEADER);
+    const auto auth_header = client_metadata.find(internal::kGrpcAuthHeader);
     std::string token;
     if (auth_header == client_metadata.end()) {
       token = "";
@@ -349,16 +389,18 @@ class FlightServiceImpl : public FlightService::Service {
     return grpc::Status::OK;
   }
 
-  grpc::Status DoPut(ServerContext* context, grpc::ServerReader<pb::FlightData>* reader,
-                     pb::PutResult* response) {
+  grpc::Status DoPut(ServerContext* context,
+                     grpc::ServerReaderWriter<pb::PutResult, pb::FlightData>* reader) {
     GrpcServerCallContext flight_context;
     GRPC_RETURN_NOT_GRPC_OK(CheckAuth(context, flight_context));
 
     auto message_reader =
         std::unique_ptr<FlightMessageReaderImpl>(new FlightMessageReaderImpl(reader));
     GRPC_RETURN_NOT_OK(message_reader->Init());
-    return internal::ToGrpcStatus(
-        server_->DoPut(flight_context, std::move(message_reader)));
+    auto metadata_writer =
+        std::unique_ptr<FlightMetadataWriter>(new GrpcMetadataWriter(reader));
+    return internal::ToGrpcStatus(server_->DoPut(
+        flight_context, std::move(message_reader), std::move(metadata_writer)));
   }
 
   grpc::Status ListActions(ServerContext* context, const pb::Empty* request,
@@ -410,6 +452,8 @@ class FlightServiceImpl : public FlightService::Service {
 
 }  // namespace
 
+FlightMetadataWriter::~FlightMetadataWriter() = default;
+
 //
 // gRPC server lifecycle
 //
@@ -572,7 +616,8 @@ Status FlightServerBase::DoGet(const ServerCallContext& context, const Ticket& r
 }
 
 Status FlightServerBase::DoPut(const ServerCallContext& context,
-                               std::unique_ptr<FlightMessageReader> reader) {
+                               std::unique_ptr<FlightMessageReader> reader,
+                               std::unique_ptr<FlightMetadataWriter> writer) {
   return Status::NotImplemented("NYI");
 }
 

cpp/src/arrow/flight/server.h

@@ -81,8 +81,9 @@ class ARROW_FLIGHT_EXPORT RecordBatchStream : public FlightDataStream {
 #pragma warning(disable : 4275)
 #endif
 
-/// \brief A reader for IPC payloads uploaded by a client
-class ARROW_FLIGHT_EXPORT FlightMessageReader : public RecordBatchReader {
+/// \brief A reader for IPC payloads uploaded by a client. Also allows
+/// reading application-defined metadata via the Flight protocol.
+class ARROW_FLIGHT_EXPORT FlightMessageReader : public MetadataRecordBatchReader {
  public:
   /// \brief Get the descriptor for this upload.
   virtual const FlightDescriptor& descriptor() const = 0;
@@ -92,6 +93,15 @@ class ARROW_FLIGHT_EXPORT FlightMessageReader : public RecordBatchReader {
 #pragma warning(pop)
 #endif
 
+/// \brief A writer for application-specific metadata sent back to the
+/// client during an upload.
+class ARROW_FLIGHT_EXPORT FlightMetadataWriter {
+ public:
+  virtual ~FlightMetadataWriter();
+  /// \brief Send a message to the client.
+  virtual Status WriteMetadata(const Buffer& app_metadata) = 0;
+};
+
 /// \brief Call state/contextual data.
 class ARROW_FLIGHT_EXPORT ServerCallContext {
  public:
@@ -178,9 +188,11 @@ class ARROW_FLIGHT_EXPORT FlightServerBase {
   /// \brief Process a stream of IPC payloads sent from a client
   /// \param[in] context The call context.
   /// \param[in] reader a sequence of uploaded record batches
+  /// \param[in] writer send metadata back to the client
   /// \return Status
   virtual Status DoPut(const ServerCallContext& context,
-                       std::unique_ptr<FlightMessageReader> reader);
+                       std::unique_ptr<FlightMessageReader> reader,
+                       std::unique_ptr<FlightMetadataWriter> writer);
 
   /// \brief Execute an action, return stream of zero or more results
   /// \param[in] context The call context.

cpp/src/arrow/flight/test-integration-client.cc

@@ -44,17 +44,27 @@ DEFINE_string(host, "localhost", "Server port to connect to");
 DEFINE_int32(port, 31337, "Server port to connect to");
 DEFINE_string(path, "", "Resource path to request");
 
-/// \brief Helper to read a RecordBatchReader into a Table.
-arrow::Status ReadToTable(std::unique_ptr<arrow::RecordBatchReader>& reader,
+/// \brief Helper to read a MetadataRecordBatchReader into a Table.
+arrow::Status ReadToTable(arrow::flight::MetadataRecordBatchReader& reader,
                           std::shared_ptr<arrow::Table>* retrieved_data) {
+  // For integration testing, we expect the server numbers the
+  // batches, to test the application metadata part of the spec.
   std::vector<std::shared_ptr<arrow::RecordBatch>> retrieved_chunks;
   std::shared_ptr<arrow::RecordBatch> chunk;
+  std::shared_ptr<arrow::Buffer> metadata_chunk;
+  int counter = 0;
   while (true) {
-    RETURN_NOT_OK(reader->ReadNext(&chunk));
+    RETURN_NOT_OK(reader.ReadWithMetadata(&chunk, &metadata_chunk));
     if (chunk == nullptr) break;
     retrieved_chunks.push_back(chunk);
+    if (std::to_string(counter) != metadata_chunk->ToString()) {
+      return arrow::Status::Invalid(
+          "Expected metadata value: " + std::to_string(counter) +
+          " but got: " + metadata_chunk->ToString());
+    }
+    counter++;
   }
-  return arrow::Table::FromRecordBatches(reader->schema(), retrieved_chunks,
+  return arrow::Table::FromRecordBatches(reader.schema(), retrieved_chunks,
                                          retrieved_data);
 }
 
@@ -71,14 +81,27 @@ arrow::Status ReadToTable(std::unique_ptr<arrow::ipc::internal::json::JsonReader
                                          retrieved_data);
 }
 
-/// \brief Helper to copy a RecordBatchReader to a RecordBatchWriter.
-arrow::Status CopyReaderToWriter(std::unique_ptr<arrow::RecordBatchReader>& reader,
-                                 arrow::ipc::RecordBatchWriter& writer) {
+/// \brief Upload the contents of a RecordBatchReader to a Flight
+/// server, validating the application metadata on the side.
+arrow::Status UploadReaderToFlight(arrow::RecordBatchReader* reader,
+                                   arrow::flight::FlightStreamWriter& writer,
+                                   arrow::flight::FlightMetadataReader& metadata_reader) {
+  int counter = 0;
   while (true) {
     std::shared_ptr<arrow::RecordBatch> chunk;
     RETURN_NOT_OK(reader->ReadNext(&chunk));
     if (chunk == nullptr) break;
-    RETURN_NOT_OK(writer.WriteRecordBatch(*chunk));
+    std::shared_ptr<arrow::Buffer> metadata =
+        arrow::Buffer::FromString(std::to_string(counter));
+    RETURN_NOT_OK(writer.WriteWithMetadata(*chunk, metadata));
+    // Wait for the server to ack the result
+    std::shared_ptr<arrow::Buffer> ack_metadata;
+    RETURN_NOT_OK(metadata_reader.ReadMetadata(&ack_metadata));
+    if (!ack_metadata->Equals(*metadata)) {
+      return arrow::Status::Invalid("Expected metadata value: " + metadata->ToString() +
+                                    " but got: " + ack_metadata->ToString());
+    }
+    counter++;
   }
   return writer.Close();
 }
@@ -91,10 +114,10 @@ arrow::Status ConsumeFlightLocation(const arrow::flight::Location& location,
   std::unique_ptr<arrow::flight::FlightClient> read_client;
   RETURN_NOT_OK(arrow::flight::FlightClient::Connect(location, &read_client));
 
-  std::unique_ptr<arrow::RecordBatchReader> stream;
+  std::unique_ptr<arrow::flight::MetadataRecordBatchReader> stream;
   RETURN_NOT_OK(read_client->DoGet(ticket, &stream));
 
-  return ReadToTable(stream, retrieved_data);
+  return ReadToTable(*stream, retrieved_data);
 }
 
 int main(int argc, char** argv) {
@@ -120,11 +143,12 @@ int main(int argc, char** argv) {
   std::shared_ptr<arrow::Table> original_data;
   ABORT_NOT_OK(ReadToTable(reader, &original_data));
 
-  std::unique_ptr<arrow::ipc::RecordBatchWriter> write_stream;
-  ABORT_NOT_OK(client->DoPut(descr, reader->schema(), &write_stream));
+  std::unique_ptr<arrow::flight::FlightStreamWriter> write_stream;
+  std::unique_ptr<arrow::flight::FlightMetadataReader> metadata_reader;
+  ABORT_NOT_OK(client->DoPut(descr, reader->schema(), &write_stream, &metadata_reader));
   std::unique_ptr<arrow::RecordBatchReader> table_reader(
       new arrow::TableBatchReader(*original_data));
-  ABORT_NOT_OK(CopyReaderToWriter(table_reader, *write_stream));
+  ABORT_NOT_OK(UploadReaderToFlight(table_reader.get(), *write_stream, *metadata_reader));
 
   // 2. Get the ticket for the data.
   std::unique_ptr<arrow::flight::FlightInfo> info;

cpp/src/arrow/flight/test-integration-server.cc

@@ -79,14 +79,16 @@ class FlightIntegrationTestServer : public FlightServerBase {
     }
     auto flight = data->second;
 
-    *data_stream = std::unique_ptr<FlightDataStream>(new RecordBatchStream(
-        std::shared_ptr<RecordBatchReader>(new TableBatchReader(*flight))));
+    *data_stream = std::unique_ptr<FlightDataStream>(
+        new NumberingStream(std::unique_ptr<FlightDataStream>(new RecordBatchStream(
+            std::shared_ptr<RecordBatchReader>(new TableBatchReader(*flight))))));
 
     return Status::OK();
   }
 
   Status DoPut(const ServerCallContext& context,
-               std::unique_ptr<FlightMessageReader> reader) override {
+               std::unique_ptr<FlightMessageReader> reader,
+               std::unique_ptr<FlightMetadataWriter> writer) override {
     const FlightDescriptor& descriptor = reader->descriptor();
 
     if (descriptor.type != FlightDescriptor::DescriptorType::PATH) {
@@ -99,10 +101,12 @@ class FlightIntegrationTestServer : public FlightServerBase {
 
     std::vector<std::shared_ptr<arrow::RecordBatch>> retrieved_chunks;
     std::shared_ptr<arrow::RecordBatch> chunk;
+    std::shared_ptr<arrow::Buffer> metadata;
     while (true) {
-      RETURN_NOT_OK(reader->ReadNext(&chunk));
+      RETURN_NOT_OK(reader->ReadWithMetadata(&chunk, &metadata));
       if (chunk == nullptr) break;
       retrieved_chunks.push_back(chunk);
+      RETURN_NOT_OK(writer->WriteMetadata(*metadata));
     }
     std::shared_ptr<arrow::Table> retrieved_data;
     RETURN_NOT_OK(arrow::Table::FromRecordBatches(reader->schema(), retrieved_chunks,

cpp/src/arrow/flight/test-util.cc

@@ -260,6 +260,22 @@ Status MakeFlightInfo(const Schema& schema, const FlightDescriptor& descriptor,
   return internal::SchemaToString(schema, &out->schema);
 }
 
+NumberingStream::NumberingStream(std::unique_ptr<FlightDataStream> stream)
+    : counter_(0), stream_(std::move(stream)) {}
+
+std::shared_ptr<Schema> NumberingStream::schema() { return stream_->schema(); }
+
+Status NumberingStream::GetSchemaPayload(FlightPayload* payload) {
+  return stream_->GetSchemaPayload(payload);
+}
+
+Status NumberingStream::Next(FlightPayload* payload) {
+  RETURN_NOT_OK(stream_->Next(payload));
+  payload->app_metadata = Buffer::FromString(std::to_string(counter_));
+  counter_++;
+  return Status::OK();
+}
+
 std::shared_ptr<Schema> ExampleIntSchema() {
   auto f0 = field("f0", int32());
   auto f1 = field("f1", int32());

cpp/src/arrow/flight/test-util.h

@@ -25,6 +25,7 @@
 #include "arrow/status.h"
 
 #include "arrow/flight/client_auth.h"
+#include "arrow/flight/server.h"
 #include "arrow/flight/server_auth.h"
 #include "arrow/flight/types.h"
 #include "arrow/flight/visibility.h"
@@ -127,6 +128,23 @@ class ARROW_FLIGHT_EXPORT BatchIterator : public RecordBatchReader {
 #pragma warning(pop)
 #endif
 
+// ----------------------------------------------------------------------
+// A FlightDataStream that numbers the record batches
+/// \brief A basic implementation of FlightDataStream that will provide
+/// a sequence of FlightData messages to be written to a gRPC stream
+class ARROW_EXPORT NumberingStream : public FlightDataStream {
+ public:
+  explicit NumberingStream(std::unique_ptr<FlightDataStream> stream);
+
+  std::shared_ptr<Schema> schema() override;
+  Status GetSchemaPayload(FlightPayload* payload) override;
+  Status Next(FlightPayload* payload) override;
+
+ private:
+  int counter_;
+  std::shared_ptr<FlightDataStream> stream_;
+};
+
 // ----------------------------------------------------------------------
 // Example data for test-server and unit tests
 

cpp/src/arrow/flight/types.h

@@ -27,6 +27,7 @@
 #include <vector>
 
 #include "arrow/flight/visibility.h"
+#include "arrow/ipc/reader.h"
 #include "arrow/ipc/writer.h"
 
 namespace arrow {
@@ -205,6 +206,7 @@ struct ARROW_FLIGHT_EXPORT FlightEndpoint {
 /// This structure corresponds to FlightData in the protocol.
 struct ARROW_FLIGHT_EXPORT FlightPayload {
   std::shared_ptr<Buffer> descriptor;
+  std::shared_ptr<Buffer> app_metadata;
   ipc::internal::IpcPayload ipc_message;
 };
 
@@ -278,6 +280,16 @@ class ARROW_FLIGHT_EXPORT ResultStream {
   virtual Status Next(std::unique_ptr<Result>* info) = 0;
 };
 
+/// \brief A RecordBatchReader that also exposes application-defined
+/// metadata from the Flight protocol.
+class ARROW_EXPORT MetadataRecordBatchReader : public RecordBatchReader {
+ public:
+  virtual ~MetadataRecordBatchReader() = default;
+
+  virtual Status ReadWithMetadata(std::shared_ptr<RecordBatch>* out,
+                                  std::shared_ptr<Buffer>* app_metadata) = 0;
+};
+
 // \brief Create a FlightListing from a vector of FlightInfo objects. This can
 // be iterated once, then it is consumed
 class ARROW_FLIGHT_EXPORT SimpleFlightListing : public FlightListing {

cpp/src/arrow/ipc/reader.cc

@@ -455,7 +455,7 @@ class RecordBatchStreamReader::RecordBatchStreamReaderImpl {
     RETURN_NOT_OK(
         ReadMessageAndValidate(message_reader_.get(), /*allow_null=*/false, &message));
 
-    CHECK_MESSAGE_TYPE(message->type(), Message::SCHEMA);
+    CHECK_MESSAGE_TYPE(Message::SCHEMA, message->type());
     CHECK_HAS_NO_BODY(*message);
     if (message->header() == nullptr) {
       return Status::IOError("Header-pointer of flatbuffer-encoded Message is null.");

cpp/src/arrow/python/flight.cc

@@ -108,10 +108,12 @@ Status PyFlightServer::DoGet(const arrow::flight::ServerCallContext& context,
   });
 }
 
-Status PyFlightServer::DoPut(const arrow::flight::ServerCallContext& context,
-                             std::unique_ptr<arrow::flight::FlightMessageReader> reader) {
+Status PyFlightServer::DoPut(
+    const arrow::flight::ServerCallContext& context,
+    std::unique_ptr<arrow::flight::FlightMessageReader> reader,
+    std::unique_ptr<arrow::flight::FlightMetadataWriter> writer) {
   return SafeCallIntoPython([&] {
-    vtable_.do_put(server_.obj(), context, std::move(reader));
+    vtable_.do_put(server_.obj(), context, std::move(reader), std::move(writer));
     return CheckPyError();
   });
 }

cpp/src/arrow/python/flight.h

@@ -50,7 +50,8 @@ class ARROW_PYTHON_EXPORT PyFlightServerVtable {
                      std::unique_ptr<arrow::flight::FlightDataStream>*)>
       do_get;
   std::function<void(PyObject*, const arrow::flight::ServerCallContext&,
-                     std::unique_ptr<arrow::flight::FlightMessageReader>)>
+                     std::unique_ptr<arrow::flight::FlightMessageReader>,
+                     std::unique_ptr<arrow::flight::FlightMetadataWriter>)>
       do_put;
   std::function<void(PyObject*, const arrow::flight::ServerCallContext&,
                      const arrow::flight::Action&,
@@ -123,7 +124,8 @@ class ARROW_PYTHON_EXPORT PyFlightServer : public arrow::flight::FlightServerBas
                const arrow::flight::Ticket& request,
                std::unique_ptr<arrow::flight::FlightDataStream>* stream) override;
   Status DoPut(const arrow::flight::ServerCallContext& context,
-               std::unique_ptr<arrow::flight::FlightMessageReader> reader) override;
+               std::unique_ptr<arrow::flight::FlightMessageReader> reader,
+               std::unique_ptr<arrow::flight::FlightMetadataWriter> writer) override;
   Status DoAction(const arrow::flight::ServerCallContext& context,
                   const arrow::flight::Action& action,
                   std::unique_ptr<arrow::flight::ResultStream>* result) override;

format/Flight.proto

@@ -288,7 +288,7 @@ message FlightData {
   /*
    * Application-defined metadata.
    */
-  bytes data_app_metadata = 3;
+  bytes app_metadata = 3;
 
   /*
    * The actual batch of Arrow data. Preferably handled with minimal-copies
@@ -303,5 +303,5 @@ message FlightData {
  * The response message associated with the submission of a DoPut.
  */
 message PutResult {
-  bytes data_app_metadata = 1;
+  bytes app_metadata = 1;
 }

java/flight/src/main/java/org/apache/arrow/flight/ArrowMessage.java

@@ -74,9 +74,16 @@ class ArrowMessage implements AutoCloseable {
       (FlightData.DATA_BODY_FIELD_NUMBER << 3) | WireFormat.WIRETYPE_LENGTH_DELIMITED;
   private static final int HEADER_TAG =
       (FlightData.DATA_HEADER_FIELD_NUMBER << 3) | WireFormat.WIRETYPE_LENGTH_DELIMITED;
+  private static final int APP_METADATA_TAG =
+      (FlightData.APP_METADATA_FIELD_NUMBER << 3) | WireFormat.WIRETYPE_LENGTH_DELIMITED;
 
   private static Marshaller<FlightData> NO_BODY_MARSHALLER = ProtoUtils.marshaller(FlightData.getDefaultInstance());
 
+  /** Get the application-specific metadata in this message. */
+  public byte[] getApplicationMetadata() {
+    return appMetadata;
+  }
+
   /** Types of messages that can be sent. */
   public enum HeaderType {
     NONE,
@@ -114,6 +121,7 @@ class ArrowMessage implements AutoCloseable {
 
   private final FlightDescriptor descriptor;
   private final Message message;
+  private final byte[] appMetadata;
   private final List<ArrowBuf> bufs;
 
   public ArrowMessage(FlightDescriptor descriptor, Schema schema) {
@@ -124,9 +132,10 @@ class ArrowMessage implements AutoCloseable {
     message = Message.getRootAsMessage(serializedMessage);
     bufs = ImmutableList.of();
     this.descriptor = descriptor;
+    this.appMetadata = null;
   }
 
-  public ArrowMessage(ArrowRecordBatch batch) {
+  public ArrowMessage(ArrowRecordBatch batch, byte[] appMetadata) {
     FlatBufferBuilder builder = new FlatBufferBuilder();
     int batchOffset = batch.writeTo(builder);
     ByteBuffer serializedMessage = MessageSerializer.serializeMessage(builder, MessageHeader.RecordBatch, batchOffset,
@@ -135,11 +144,13 @@ class ArrowMessage implements AutoCloseable {
     this.message = Message.getRootAsMessage(serializedMessage);
     this.bufs = ImmutableList.copyOf(batch.getBuffers());
     this.descriptor = null;
+    this.appMetadata = appMetadata;
   }
 
-  private ArrowMessage(FlightDescriptor descriptor, Message message, ArrowBuf buf) {
+  private ArrowMessage(FlightDescriptor descriptor, Message message, byte[] appMetadata, ArrowBuf buf) {
     this.message = message;
     this.descriptor = descriptor;
+    this.appMetadata = appMetadata;
     this.bufs = buf == null ? ImmutableList.of() : ImmutableList.of(buf);
   }
 
@@ -183,6 +194,7 @@ class ArrowMessage implements AutoCloseable {
       FlightDescriptor descriptor = null;
       Message header = null;
       ArrowBuf body = null;
+      byte[] appMetadata = null;
       while (stream.available() > 0) {
         int tag = readRawVarint32(stream);
         switch (tag) {
@@ -201,6 +213,12 @@ class ArrowMessage implements AutoCloseable {
             header = Message.getRootAsMessage(ByteBuffer.wrap(bytes));
             break;
           }
+          case APP_METADATA_TAG: {
+            int size = readRawVarint32(stream);
+            appMetadata = new byte[size];
+            ByteStreams.readFully(stream, appMetadata);
+            break;
+          }
           case BODY_TAG:
             if (body != null) {
               // only read last body.
@@ -225,7 +243,7 @@ class ArrowMessage implements AutoCloseable {
         }
       }
 
-      return new ArrowMessage(descriptor, header, body);
+      return new ArrowMessage(descriptor, header, appMetadata, body);
     } catch (Exception ioe) {
       throw new RuntimeException(ioe);
     }
@@ -246,7 +264,6 @@ class ArrowMessage implements AutoCloseable {
 
       final ByteString bytes = ByteString.copyFrom(message.getByteBuffer(), message.getByteBuffer().remaining());
 
-
       if (getMessageType() == HeaderType.SCHEMA) {
 
         final FlightData.Builder builder = FlightData.newBuilder()
@@ -267,8 +284,12 @@ class ArrowMessage implements AutoCloseable {
       ByteArrayOutputStream baos = new ByteArrayOutputStream();
       CodedOutputStream cos = CodedOutputStream.newInstance(baos);
       cos.writeBytes(FlightData.DATA_HEADER_FIELD_NUMBER, bytes);
-      cos.writeTag(FlightData.DATA_BODY_FIELD_NUMBER, WireFormat.WIRETYPE_LENGTH_DELIMITED);
 
+      if (appMetadata != null && appMetadata.length > 0) {
+        cos.writeByteArray(FlightData.APP_METADATA_FIELD_NUMBER, appMetadata);
+      }
+
+      cos.writeTag(FlightData.DATA_BODY_FIELD_NUMBER, WireFormat.WIRETYPE_LENGTH_DELIMITED);
       int size = 0;
       List<ByteBuf> allBufs = new ArrayList<>();
       for (ArrowBuf b : bufs) {

java/flight/src/main/java/org/apache/arrow/flight/FlightBindingService.java

@@ -28,12 +28,11 @@ import com.google.common.collect.ImmutableSet;
 
 import io.grpc.BindableService;
 import io.grpc.MethodDescriptor;
+import io.grpc.MethodDescriptor.MethodType;
 import io.grpc.ServerMethodDefinition;
 import io.grpc.ServerServiceDefinition;
 import io.grpc.protobuf.ProtoUtils;
 import io.grpc.stub.ServerCalls;
-import io.grpc.stub.ServerCalls.ClientStreamingMethod;
-import io.grpc.stub.ServerCalls.ServerStreamingMethod;
 import io.grpc.stub.StreamObserver;
 
 /**
@@ -66,7 +65,7 @@ class FlightBindingService implements BindableService {
 
   public static MethodDescriptor<ArrowMessage, Flight.PutResult> getDoPutDescriptor(BufferAllocator allocator) {
     return MethodDescriptor.<ArrowMessage, Flight.PutResult>newBuilder()
-        .setType(io.grpc.MethodDescriptor.MethodType.CLIENT_STREAMING)
+        .setType(MethodType.BIDI_STREAMING)
         .setFullMethodName(DO_PUT)
         .setSampledToLocalTracing(false)
         .setRequestMarshaller(ArrowMessage.createMarshaller(allocator))
@@ -84,7 +83,7 @@ class FlightBindingService implements BindableService {
 
     ServerServiceDefinition.Builder serviceBuilder = ServerServiceDefinition.builder(FlightConstants.SERVICE);
     serviceBuilder.addMethod(doGetDescriptor, ServerCalls.asyncServerStreamingCall(new DoGetMethod(delegate)));
-    serviceBuilder.addMethod(doPutDescriptor, ServerCalls.asyncClientStreamingCall(new DoPutMethod(delegate)));
+    serviceBuilder.addMethod(doPutDescriptor, ServerCalls.asyncBidiStreamingCall(new DoPutMethod(delegate)));
 
     // copy over not-overridden methods.
     for (ServerMethodDefinition<?, ?> definition : baseDefinition.getMethods()) {
@@ -98,7 +97,7 @@ class FlightBindingService implements BindableService {
     return serviceBuilder.build();
   }
 
-  private class DoGetMethod implements ServerStreamingMethod<Flight.Ticket, ArrowMessage> {
+  private class DoGetMethod implements ServerCalls.ServerStreamingMethod<Flight.Ticket, ArrowMessage> {
 
     private final FlightService delegate;
 
@@ -112,7 +111,7 @@ class FlightBindingService implements BindableService {
     }
   }
 
-  private class DoPutMethod implements ClientStreamingMethod<ArrowMessage, Flight.PutResult> {
+  private class DoPutMethod implements ServerCalls.BidiStreamingMethod<ArrowMessage, PutResult> {
     private final FlightService delegate;
 
     public DoPutMethod(FlightService delegate) {

java/flight/src/main/java/org/apache/arrow/flight/FlightClient.java

@@ -17,24 +17,22 @@
 
 package org.apache.arrow.flight;
 
-import static io.grpc.stub.ClientCalls.asyncClientStreamingCall;
-import static io.grpc.stub.ClientCalls.asyncServerStreamingCall;
-
 import java.io.InputStream;
 import java.net.URISyntaxException;
 import java.util.Iterator;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
 
 import javax.net.ssl.SSLException;
 
+import org.apache.arrow.flight.FlightProducer.StreamListener;
 import org.apache.arrow.flight.auth.BasicClientAuthHandler;
 import org.apache.arrow.flight.auth.ClientAuthHandler;
 import org.apache.arrow.flight.auth.ClientAuthInterceptor;
 import org.apache.arrow.flight.auth.ClientAuthWrapper;
 import org.apache.arrow.flight.impl.Flight;
 import org.apache.arrow.flight.impl.Flight.Empty;
-import org.apache.arrow.flight.impl.Flight.PutResult;
 import org.apache.arrow.flight.impl.FlightServiceGrpc;
 import org.apache.arrow.flight.impl.FlightServiceGrpc.FlightServiceBlockingStub;
 import org.apache.arrow.flight.impl.FlightServiceGrpc.FlightServiceStub;
@@ -47,8 +45,6 @@ import com.google.common.base.Preconditions;
 import com.google.common.base.Throwables;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Iterators;
-import com.google.common.util.concurrent.ListenableFuture;
-import com.google.common.util.concurrent.SettableFuture;
 
 import io.grpc.ClientCall;
 import io.grpc.ManagedChannel;
@@ -56,6 +52,7 @@ import io.grpc.MethodDescriptor;
 import io.grpc.netty.GrpcSslContexts;
 import io.grpc.netty.NettyChannelBuilder;
 import io.grpc.stub.ClientCallStreamObserver;
+import io.grpc.stub.ClientCalls;
 import io.grpc.stub.ClientResponseObserver;
 import io.grpc.stub.StreamObserver;
 
@@ -158,18 +155,19 @@ public class FlightClient implements AutoCloseable {
    * Create or append a descriptor with another stream.
    * @param descriptor FlightDescriptor
    * @param root VectorSchemaRoot
+   * @param metadataListener A handler for metadata messages from the server.
    * @param options RPC-layer hints for this call.
    * @return ClientStreamListener
    */
-  public ClientStreamListener startPut(
-      FlightDescriptor descriptor, VectorSchemaRoot root, CallOption... options) {
+  public ClientStreamListener startPut(FlightDescriptor descriptor, VectorSchemaRoot root,
+      StreamListener<PutResult> metadataListener, CallOption... options) {
     Preconditions.checkNotNull(descriptor);
     Preconditions.checkNotNull(root);
 
-    SetStreamObserver<PutResult> resultObserver = new SetStreamObserver<>();
+    SetStreamObserver resultObserver = new SetStreamObserver(metadataListener);
     final io.grpc.CallOptions callOptions = CallOptions.wrapStub(asyncStub, options).getCallOptions();
     ClientCallStreamObserver<ArrowMessage> observer = (ClientCallStreamObserver<ArrowMessage>)
-        asyncClientStreamingCall(
+        ClientCalls.asyncBidiStreamingCall(
                 authInterceptor.interceptCall(doPutDescriptor, callOptions, channel), resultObserver);
     // send the schema to start.
     ArrowMessage message = new ArrowMessage(descriptor.toProtocol(), root.getSchema());
@@ -235,30 +233,43 @@ public class FlightClient implements AutoCloseable {
 
         };
 
-    asyncServerStreamingCall(call, ticket.toProtocol(), clientResponseObserver);
+    ClientCalls.asyncServerStreamingCall(call, ticket.toProtocol(), clientResponseObserver);
     return stream;
   }
 
-  private static class SetStreamObserver<T> implements StreamObserver<T> {
-    private final SettableFuture<T> result = SettableFuture.create();
-    private volatile T resultLocal;
+  private static class SetStreamObserver implements StreamObserver<Flight.PutResult> {
+    private final CompletableFuture<Void> result;
+    private final StreamListener<PutResult> listener;
+
+    SetStreamObserver(FlightProducer.StreamListener<PutResult> listener) {
+      this.listener = listener;
+      result = new CompletableFuture<>();
+    }
 
     @Override
-    public void onNext(T value) {
-      resultLocal = value;
+    public void onNext(Flight.PutResult value) {
+      if (listener != null) {
+        listener.onNext(PutResult.fromProtocol(value));
+      }
     }
 
     @Override
     public void onError(Throwable t) {
-      result.setException(t);
+      result.completeExceptionally(t);
+      if (listener != null) {
+        listener.onError(t);
+      }
     }
 
     @Override
     public void onCompleted() {
-      result.set(Preconditions.checkNotNull(resultLocal));
+      if (listener != null) {
+        listener.onCompleted();
+      }
+      result.complete(null);
     }
 
-    public ListenableFuture<T> getFuture() {
+    public CompletableFuture<Void> getFuture() {
       return result;
     }
   }
@@ -266,10 +277,10 @@ public class FlightClient implements AutoCloseable {
   private static class PutObserver implements ClientStreamListener {
     private final ClientCallStreamObserver<ArrowMessage> observer;
     private final VectorUnloader unloader;
-    private final ListenableFuture<PutResult> futureResult;
+    private final CompletableFuture<Void> futureResult;
 
     public PutObserver(VectorUnloader unloader, ClientCallStreamObserver<ArrowMessage> observer,
-        ListenableFuture<PutResult> futureResult) {
+        CompletableFuture<Void> futureResult) {
       this.observer = observer;
       this.unloader = unloader;
       this.futureResult = futureResult;
@@ -277,12 +288,17 @@ public class FlightClient implements AutoCloseable {
 
     @Override
     public void putNext() {
+      putNext(null);
+    }
+
+    @Override
+    public void putNext(byte[] appMetadata) {
       ArrowRecordBatch batch = unloader.getRecordBatch();
       // Check the futureResult in case server sent an exception
       while (!observer.isReady() && !futureResult.isDone()) {
         /* busy wait */
       }
-      observer.onNext(new ArrowMessage(batch));
+      observer.onNext(new ArrowMessage(batch, appMetadata));
     }
 
     @Override
@@ -296,9 +312,9 @@ public class FlightClient implements AutoCloseable {
     }
 
     @Override
-    public PutResult getResult() {
+    public void getResult() {
       try {
-        return futureResult.get();
+        futureResult.get();
       } catch (Exception ex) {
         throw Throwables.propagate(ex);
       }
@@ -312,11 +328,15 @@ public class FlightClient implements AutoCloseable {
 
     void putNext();
 
+    void putNext(byte[] appMetadata);
+
     void error(Throwable ex);
 
+    /** Indicate the stream is finished on the client side. */
     void completed();
 
-    PutResult getResult();
+    /** Wait for the stream to finish on the server side. */
+    void getResult();
 
   }
 

java/flight/src/main/java/org/apache/arrow/flight/FlightProducer.java

@@ -17,9 +17,6 @@
 
 package org.apache.arrow.flight;
 
-import java.util.concurrent.Callable;
-
-import org.apache.arrow.flight.impl.Flight.PutResult;
 import org.apache.arrow.vector.VectorSchemaRoot;
 
 /**
@@ -27,17 +24,23 @@ import org.apache.arrow.vector.VectorSchemaRoot;
  */
 public interface FlightProducer {
 
+  /**
+   * Return data for the given stream.
+   */
   void getStream(CallContext context, Ticket ticket,
       ServerStreamListener listener);
 
+  /** List available flights given some criteria. */
   void listFlights(CallContext context, Criteria criteria,
       StreamListener<FlightInfo> listener);
 
+  /** Get information on a specific flight. */
   FlightInfo getFlightInfo(CallContext context,
       FlightDescriptor descriptor);
 
-  Callable<PutResult> acceptPut(CallContext context,
-      FlightStream flightStream);
+  /** Accept uploaded data for a flight. */
+  Runnable acceptPut(CallContext context,
+      FlightStream flightStream, StreamListener<PutResult> ackStream);
 
   void doAction(CallContext context, Action action,
       StreamListener<Result> listener);
@@ -50,16 +53,35 @@ public interface FlightProducer {
    */
   interface ServerStreamListener {
 
+    /**
+     * Check if the client has cancelled their request.
+     */
     boolean isCancelled();
 
+    /**
+     * Check if the client is ready to receive further data.
+     */
     boolean isReady();
 
+    /**
+     * Begin sending data from the specified root.
+     */
     void start(VectorSchemaRoot root);
 
+    /** Send the current contents of the root to the client. */
     void putNext();
 
+    /**
+     * Send the current contents of the root to the client, along with application-defined metadata on the side.
+     */
+    void putNext(byte[] metadata);
+
+    /**
+     *  Indicate an error to the client. Do not call {@link #completed} afterwards.
+     */
     void error(Throwable ex);
 
+    /** Indicate that the stream is finished. */
     void completed();
 
   }

java/flight/src/main/java/org/apache/arrow/flight/FlightServer.java

@@ -21,6 +21,8 @@ import java.io.File;
 import java.io.FileInputStream;
 import java.io.IOException;
 import java.io.InputStream;
+import java.net.URI;
+import java.net.URISyntaxException;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.concurrent.Executor;
@@ -47,13 +49,15 @@ public class FlightServer implements AutoCloseable {
 
   private static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(FlightServer.class);
 
+  private final Location location;
   private final Server server;
 
   /** The maximum size of an individual gRPC message. This effectively disables the limit. */
   static final int MAX_GRPC_MESSAGE_SIZE = Integer.MAX_VALUE;
 
   /** Create a new instance from a gRPC server. For internal use only. */
-  private FlightServer(Server server) {
+  private FlightServer(Location location, Server server) {
+    this.location = location;
     this.server = server;
   }
 
@@ -63,10 +67,27 @@ public class FlightServer implements AutoCloseable {
     return this;
   }
 
+  /** Get the port the server is running on (if applicable). */
   public int getPort() {
     return server.getPort();
   }
 
+  /** Get the location for this server. */
+  public Location getLocation() {
+    if (location.getUri().getPort() == 0) {
+      // If the server was bound to port 0, replace the port in the location with the real port.
+      final URI uri = location.getUri();
+      try {
+        return new Location(new URI(uri.getScheme(), uri.getUserInfo(), uri.getHost(), getPort(),
+            uri.getPath(), uri.getQuery(), uri.getFragment()));
+      } catch (URISyntaxException e) {
+        // We don't expect this to happen
+        throw new RuntimeException(e);
+      }
+    }
+    return location;
+  }
+
   /** Block until the server shuts down. */
   public void awaitTermination() throws InterruptedException {
     server.awaitTermination();
@@ -211,7 +232,7 @@ public class FlightServer implements AutoCloseable {
         return null;
       });
 
-      return new FlightServer(builder.build());
+      return new FlightServer(location, builder.build());
     }
 
     /**

java/flight/src/main/java/org/apache/arrow/flight/FlightService.java

@@ -30,7 +30,6 @@ import org.apache.arrow.flight.impl.Flight.ActionType;
 import org.apache.arrow.flight.impl.Flight.Empty;
 import org.apache.arrow.flight.impl.Flight.HandshakeRequest;
 import org.apache.arrow.flight.impl.Flight.HandshakeResponse;
-import org.apache.arrow.flight.impl.Flight.PutResult;
 import org.apache.arrow.flight.impl.FlightServiceGrpc.FlightServiceImplBase;
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.vector.VectorSchemaRoot;
@@ -145,8 +144,13 @@ class FlightService extends FlightServiceImplBase {
 
     @Override
     public void putNext() {
+      putNext(null);
+    }
+
+    @Override
+    public void putNext(byte[] metadata) {
       Preconditions.checkNotNull(unloader);
-      responseObserver.onNext(new ArrowMessage(unloader.getRecordBatch()));
+      responseObserver.onNext(new ArrowMessage(unloader.getRecordBatch(), metadata));
     }
 
     @Override
@@ -161,15 +165,17 @@ class FlightService extends FlightServiceImplBase {
 
   }
 
-  public StreamObserver<ArrowMessage> doPutCustom(final StreamObserver<PutResult> responseObserverSimple) {
-    ServerCallStreamObserver<PutResult> responseObserver = (ServerCallStreamObserver<PutResult>) responseObserverSimple;
+  public StreamObserver<ArrowMessage> doPutCustom(final StreamObserver<Flight.PutResult> responseObserverSimple) {
+    ServerCallStreamObserver<Flight.PutResult> responseObserver =
+        (ServerCallStreamObserver<Flight.PutResult>) responseObserverSimple;
     responseObserver.disableAutoInboundFlowControl();
     responseObserver.request(1);
 
     FlightStream fs = new FlightStream(allocator, PENDING_REQUESTS, null, (count) -> responseObserver.request(count));
     executors.submit(() -> {
       try {
-        responseObserver.onNext(producer.acceptPut(makeContext(responseObserver), fs).call());
+        producer.acceptPut(makeContext(responseObserver), fs,
+            StreamPipe.wrap(responseObserver, PutResult::toProtocol)).run();
         responseObserver.onCompleted();
       } catch (Exception ex) {
         responseObserver.onError(ex);

java/flight/src/main/java/org/apache/arrow/flight/FlightStream.java

@@ -39,7 +39,7 @@ import io.grpc.stub.StreamObserver;
 /**
  * An adaptor between protobuf streams and flight data streams.
  */
-public class FlightStream {
+public class FlightStream implements AutoCloseable {
 
 
   private final Object DONE = new Object();
@@ -60,6 +60,7 @@ public class FlightStream {
   private volatile Throwable ex;
   private volatile FlightDescriptor descriptor;
   private volatile Schema schema;
+  private volatile byte[] applicationMetadata = null;
 
   /**
    * Constructs a new instance.
@@ -135,6 +136,7 @@ public class FlightStream {
         try (ArrowRecordBatch arb = msg.asRecordBatch()) {
           loader.load(arb);
         }
+        this.applicationMetadata = msg.getApplicationMetadata();
         return true;
       }
 
@@ -152,6 +154,16 @@ public class FlightStream {
     }
   }
 
+  /**
+   * Get the most recent metadata sent from the server. This may be cleared by calls to {@link #next()} if the server
+   * sends a message without metadata.
+   *
+   * @return the application metadata. May be null.
+   */
+  public byte[] getLatestMetadata() {
+    return applicationMetadata;
+  }
+
   private synchronized void requestOutstanding() {
     if (pending < pendingTarget) {
       requestor.request(pendingTarget - pending);

java/flight/src/main/java/org/apache/arrow/flight/NoOpFlightProducer.java

@@ -17,10 +17,6 @@
 
 package org.apache.arrow.flight;
 
-import java.util.concurrent.Callable;
-
-import org.apache.arrow.flight.impl.Flight.PutResult;
-
 /**
  * A {@link FlightProducer} that throws on all operations.
  */
@@ -45,8 +41,8 @@ public class NoOpFlightProducer implements FlightProducer {
   }
 
   @Override
-  public Callable<PutResult> acceptPut(CallContext context,
-      FlightStream flightStream) {
+  public Runnable acceptPut(CallContext context,
+      FlightStream flightStream, StreamListener<PutResult> ackStream) {
     throw new UnsupportedOperationException("NYI");
   }
 

java/flight/src/main/java/org/apache/arrow/flight/PutResult.java

@@ -0,0 +1,62 @@
+/*
+ * 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.
+ */
+
+package org.apache.arrow.flight;
+
+import java.nio.ByteBuffer;
+
+import org.apache.arrow.flight.impl.Flight;
+
+import com.google.protobuf.ByteString;
+
+/**
+ * A message from the server during a DoPut operation.
+ */
+public class PutResult {
+
+  private ByteBuffer applicationMetadata;
+
+  private PutResult(ByteBuffer metadata) {
+    applicationMetadata = metadata;
+  }
+
+  /** Create a PutResult with application-specific metadata. */
+  public static PutResult metadata(byte[] metadata) {
+    return new PutResult(ByteBuffer.wrap(metadata));
+  }
+
+  /** Create an empty PutResult. */
+  public static PutResult empty() {
+    return new PutResult(null);
+  }
+
+  /** Get the metadata in this message. May be null. */
+  public ByteBuffer getApplicationMetadata() {
+    return applicationMetadata;
+  }
+
+  Flight.PutResult toProtocol() {
+    if (applicationMetadata == null) {
+      return Flight.PutResult.getDefaultInstance();
+    }
+    return Flight.PutResult.newBuilder().setAppMetadata(ByteString.copyFrom(applicationMetadata)).build();
+  }
+
+  static PutResult fromProtocol(Flight.PutResult message) {
+    return new PutResult(message.getAppMetadata().asReadOnlyByteBuffer());
+  }
+}

java/flight/src/main/java/org/apache/arrow/flight/example/InMemoryStore.java

@@ -17,7 +17,6 @@
 
 package org.apache.arrow.flight.example;
 
-import java.util.concurrent.Callable;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 
@@ -29,10 +28,10 @@ import org.apache.arrow.flight.FlightInfo;
 import org.apache.arrow.flight.FlightProducer;
 import org.apache.arrow.flight.FlightStream;
 import org.apache.arrow.flight.Location;
+import org.apache.arrow.flight.PutResult;
 import org.apache.arrow.flight.Result;
 import org.apache.arrow.flight.Ticket;
 import org.apache.arrow.flight.example.Stream.StreamCreator;
-import org.apache.arrow.flight.impl.Flight.PutResult;
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.util.AutoCloseables;
 import org.apache.arrow.vector.VectorSchemaRoot;
@@ -116,8 +115,8 @@ public class InMemoryStore implements FlightProducer, AutoCloseable {
   }
 
   @Override
-  public Callable<PutResult> acceptPut(CallContext context,
-      final FlightStream flightStream) {
+  public Runnable acceptPut(CallContext context,
+      final FlightStream flightStream, final StreamListener<PutResult> ackStream) {
     return () -> {
       StreamCreator creator = null;
       boolean success = false;
@@ -130,11 +129,11 @@ public class InMemoryStore implements FlightProducer, AutoCloseable {
 
         VectorUnloader unloader = new VectorUnloader(root);
         while (flightStream.next()) {
+          ackStream.onNext(PutResult.metadata(flightStream.getLatestMetadata()));
           creator.add(unloader.getRecordBatch());
         }
         creator.complete();
         success = true;
-        return PutResult.getDefaultInstance();
       } finally {
         if (!success) {
           creator.drop();

java/flight/src/main/java/org/apache/arrow/flight/example/Stream.java

@@ -17,6 +17,7 @@
 
 package org.apache.arrow.flight.example;
 
+import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
 import java.util.Iterator;
 import java.util.List;
@@ -84,9 +85,11 @@ public class Stream implements AutoCloseable, Iterable<ArrowRecordBatch> {
     try (VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) {
       listener.start(root);
       final VectorLoader loader = new VectorLoader(root);
+      int counter = 0;
       for (ArrowRecordBatch batch : batches) {
         loader.load(batch);
-        listener.putNext();
+        listener.putNext(Integer.toString(counter).getBytes(StandardCharsets.UTF_8));
+        counter++;
       }
       listener.completed();
     } catch (Exception ex) {

java/flight/src/main/java/org/apache/arrow/flight/example/integration/IntegrationTestClient.java

@@ -19,6 +19,7 @@ package org.apache.arrow.flight.example.integration;
 
 import java.io.File;
 import java.io.IOException;
+import java.nio.charset.StandardCharsets;
 import java.util.Collections;
 import java.util.List;
 
@@ -26,8 +27,10 @@ import org.apache.arrow.flight.FlightClient;
 import org.apache.arrow.flight.FlightDescriptor;
 import org.apache.arrow.flight.FlightEndpoint;
 import org.apache.arrow.flight.FlightInfo;
+import org.apache.arrow.flight.FlightProducer.StreamListener;
 import org.apache.arrow.flight.FlightStream;
 import org.apache.arrow.flight.Location;
+import org.apache.arrow.flight.PutResult;
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.memory.RootAllocator;
 import org.apache.arrow.vector.VectorLoader;
@@ -93,11 +96,33 @@ class IntegrationTestClient {
       jsonRoot = VectorSchemaRoot.create(root.getSchema(), allocator);
       VectorUnloader unloader = new VectorUnloader(root);
       VectorLoader jsonLoader = new VectorLoader(jsonRoot);
-      FlightClient.ClientStreamListener stream = client.startPut(descriptor, root);
+      FlightClient.ClientStreamListener stream = client.startPut(descriptor, root, new StreamListener<PutResult>() {
+        int counter = 0;
+
+        @Override
+        public void onNext(PutResult val) {
+          final String metadata = StandardCharsets.UTF_8.decode(val.getApplicationMetadata()).toString();
+          if (!Integer.toString(counter).equals(metadata)) {
+            throw new RuntimeException(
+                String.format("Invalid ACK from server. Expected '%d' but got '%s'.", counter, metadata));
+          }
+          counter++;
+        }
+
+        @Override
+        public void onError(Throwable t) {
+        }
+
+        @Override
+        public void onCompleted() {
+        }
+      });
+      int counter = 0;
       while (reader.read(root)) {
-        stream.putNext();
+        stream.putNext(Integer.toString(counter).getBytes(StandardCharsets.UTF_8));
         jsonLoader.load(unloader.getRecordBatch());
         root.clear();
+        counter++;
       }
       stream.completed();
       // Need to call this, or exceptions from the server get swallowed

java/flight/src/test/java/org/apache/arrow/flight/FlightTestUtil.java

@@ -43,14 +43,15 @@ public class FlightTestUtil {
    * Returns a a FlightServer (actually anything that is startable)
    * that has been started bound to a random port.
    */
-  public static <T> T getStartedServer(Function<Integer, T> newServerFromPort) throws IOException {
+  public static <T> T getStartedServer(Function<Location, T> newServerFromLocation) throws IOException {
     IOException lastThrown = null;
     T server = null;
     for (int x = 0; x < 3; x++) {
       final int port = 49152 + RANDOM.nextInt(5000);
+      final Location location = Location.forGrpcInsecure(LOCALHOST, port);
       lastThrown = null;
       try {
-        server = newServerFromPort.apply(port);
+        server = newServerFromLocation.apply(location);
         try {
           server.getClass().getMethod("start").invoke(server);
         } catch (NoSuchMethodException | IllegalAccessException e) {

java/flight/src/test/java/org/apache/arrow/flight/TestApplicationMetadata.java

@@ -0,0 +1,167 @@
+/*
+ * 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.
+ */
+
+package org.apache.arrow.flight;
+
+import java.util.Arrays;
+import java.util.Collections;
+
+import org.apache.arrow.flight.FlightProducer.StreamListener;
+import org.apache.arrow.flight.auth.ServerAuthHandler;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.IntVector;
+import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.types.pojo.ArrowType;
+import org.apache.arrow.vector.types.pojo.Field;
+import org.apache.arrow.vector.types.pojo.Schema;
+
+import org.junit.Assert;
+import org.junit.Test;
+
+/**
+ * Tests for application-specific metadata support in Flight.
+ */
+public class TestApplicationMetadata {
+
+  /**
+   * Ensure that a client can read the metadata sent from the server.
+   */
+  @Test
+  public void retrieveMetadata() {
+    try (final BufferAllocator a = new RootAllocator(Long.MAX_VALUE);
+        final FlightServer s =
+            FlightTestUtil.getStartedServer(
+                (location) -> FlightServer.builder(a, location, new MetadataFlightProducer(a)).build());
+        final FlightClient client = FlightClient.builder(a, s.getLocation()).build();
+        final FlightStream stream = client.getStream(new Ticket(new byte[0]))) {
+      byte i = 0;
+      while (stream.next()) {
+        final IntVector vector = (IntVector) stream.getRoot().getVector("a");
+        Assert.assertEquals(1, vector.getValueCount());
+        Assert.assertEquals(10, vector.get(0));
+        Assert.assertArrayEquals(new byte[]{i}, stream.getLatestMetadata());
+        i++;
+      }
+    } catch (Exception e) {
+      throw new RuntimeException(e);
+    }
+  }
+
+  /**
+   * Ensure that a client can send metadata to the server.
+   */
+  @Test
+  public void uploadMetadata() {
+    final FlightDescriptor descriptor = FlightDescriptor.path("test");
+    final Schema schema = new Schema(Collections.singletonList(Field.nullable("a", new ArrowType.Int(32, true))));
+    try (final BufferAllocator a = new RootAllocator(Long.MAX_VALUE);
+        VectorSchemaRoot root = VectorSchemaRoot.create(schema, a);
+        final FlightServer s =
+            FlightTestUtil.getStartedServer(
+                (location) -> FlightServer.builder(a, location, new MetadataFlightProducer(a)).build());
+        final FlightClient client = FlightClient.builder(a, s.getLocation()).build()) {
+
+      final StreamListener<PutResult> listener = new StreamListener<PutResult>() {
+        int counter = 0;
+
+        @Override
+        public void onNext(PutResult val) {
+          Assert.assertNotNull(val);
+          Assert.assertEquals(counter, val.getApplicationMetadata().get(0));
+          counter++;
+        }
+
+        @Override
+        public void onError(Throwable t) {
+          Assert.fail(t.toString());
+        }
+
+        @Override
+        public void onCompleted() {
+          Assert.assertEquals(10, counter);
+        }
+      };
+      final FlightClient.ClientStreamListener writer = client.startPut(descriptor, root, listener);
+
+      root.allocateNew();
+      for (byte i = 0; i < 10; i++) {
+        final IntVector vector = (IntVector) root.getVector("a");
+        vector.set(0, 10);
+        vector.setValueCount(1);
+        root.setRowCount(1);
+        writer.putNext(new byte[]{i});
+      }
+      writer.completed();
+      // Must attempt to retrieve the result to get any server-side errors.
+      writer.getResult();
+    } catch (Exception e) {
+      throw new RuntimeException(e);
+    }
+  }
+
+  /**
+   * A FlightProducer that always produces a fixed data stream with metadata on the side.
+   */
+  private static class MetadataFlightProducer extends NoOpFlightProducer {
+
+    private final BufferAllocator allocator;
+
+    public MetadataFlightProducer(BufferAllocator allocator) {
+      this.allocator = allocator;
+    }
+
+    @Override
+    public void getStream(CallContext context, Ticket ticket, ServerStreamListener listener) {
+      final Schema schema = new Schema(Collections.singletonList(Field.nullable("a", new ArrowType.Int(32, true))));
+      try (VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) {
+        root.allocateNew();
+        listener.start(root);
+        for (byte i = 0; i < 10; i++) {
+          final IntVector vector = (IntVector) root.getVector("a");
+          vector.set(0, 10);
+          vector.setValueCount(1);
+          root.setRowCount(1);
+          listener.putNext(new byte[]{i});
+        }
+        listener.completed();
+      }
+    }
+
+    @Override
+    public Runnable acceptPut(CallContext context, FlightStream stream, StreamListener<PutResult> ackStream) {
+      return () -> {
+        try (FlightStream flightStream = stream) {
+          byte current = 0;
+          while (flightStream.next()) {
+            byte[] metadata = flightStream.getLatestMetadata();
+            if (!Arrays.equals(new byte[]{current}, metadata)) {
+              throw new IllegalArgumentException("Metadata does not match expected value.");
+            }
+            ackStream.onNext(PutResult.metadata(new byte[]{current}));
+            current++;
+          }
+          if (current != 10) {
+            throw new IllegalArgumentException("Wrong number of messages sent.");
+          }
+        } catch (Exception e) {
+          throw new RuntimeException(e);
+        }
+      };
+    }
+  }
+}

java/flight/src/test/java/org/apache/arrow/flight/TestBackPressure.java

@@ -46,7 +46,7 @@ public class TestBackPressure {
     try (
         final BufferAllocator a = new RootAllocator(Long.MAX_VALUE);
         final PerformanceTestServer server = FlightTestUtil.getStartedServer(
-            (port) -> (new PerformanceTestServer(a, Location.forGrpcInsecure(FlightTestUtil.LOCALHOST, port))));
+            (location) -> (new PerformanceTestServer(a, location)));
         final FlightClient client = FlightClient.builder(a, server.getLocation()).build()
     ) {
       FlightStream fs1 = client.getStream(client.getInfo(
@@ -121,13 +121,12 @@ public class TestBackPressure {
       try (
           BufferAllocator serverAllocator = allocator.newChildAllocator("server", 0, Long.MAX_VALUE);
           FlightServer server =
-              FlightTestUtil.getStartedServer(
-                  (port) -> FlightServer.builder(serverAllocator, Location.forGrpcInsecure("localhost", port), producer)
-                      .build());
+              FlightTestUtil.getStartedServer((location) -> FlightServer.builder(serverAllocator, location, producer)
+                  .build());
           BufferAllocator clientAllocator = allocator.newChildAllocator("client", 0, Long.MAX_VALUE);
           FlightClient client =
               FlightClient
-                  .builder(clientAllocator, Location.forGrpcInsecure(FlightTestUtil.LOCALHOST, server.getPort()))
+                  .builder(clientAllocator, server.getLocation())
                   .build()
       ) {
         FlightStream stream = client.getStream(new Ticket(new byte[1]));

java/flight/src/test/java/org/apache/arrow/flight/TestBasicOperation.java

@@ -19,14 +19,13 @@ package org.apache.arrow.flight;
 
 import java.net.URISyntaxException;
 import java.util.Iterator;
-import java.util.concurrent.Callable;
 import java.util.function.BiConsumer;
 import java.util.function.Consumer;
 
 import org.apache.arrow.flight.FlightClient.ClientStreamListener;
+import org.apache.arrow.flight.FlightProducer.StreamListener;
 import org.apache.arrow.flight.impl.Flight;
 import org.apache.arrow.flight.impl.Flight.FlightDescriptor.DescriptorType;
-import org.apache.arrow.flight.impl.Flight.PutResult;
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.memory.RootAllocator;
 import org.apache.arrow.vector.IntVector;
@@ -98,7 +97,19 @@ public class TestBasicOperation {
       IntVector iv = new IntVector("c1", a);
 
       VectorSchemaRoot root = VectorSchemaRoot.of(iv);
-      ClientStreamListener listener = c.startPut(FlightDescriptor.path("hello"), root);
+      ClientStreamListener listener = c.startPut(FlightDescriptor.path("hello"), root, new StreamListener<PutResult>() {
+        @Override
+        public void onNext(PutResult val) {
+        }
+
+        @Override
+        public void onError(Throwable t) {
+        }
+
+        @Override
+        public void onCompleted() {
+        }
+      });
 
       //batch 1
       root.allocateNew();
@@ -155,12 +166,11 @@ public class TestBasicOperation {
         Producer producer = new Producer(a);
         FlightServer s =
             FlightTestUtil.getStartedServer(
-                (port) -> FlightServer.builder(a, Location.forGrpcInsecure("localhost", port), producer).build()
+                (location) -> FlightServer.builder(a, location, producer).build()
             )) {
 
       try (
-          FlightClient c = FlightClient.builder(a, Location.forGrpcInsecure(FlightTestUtil.LOCALHOST, s.getPort()))
-              .build()
+          FlightClient c = FlightClient.builder(a, s.getLocation()).build()
       ) {
         try (BufferAllocator testAllocator = a.newChildAllocator("testcase", 0, Long.MAX_VALUE)) {
           consumer.accept(c, testAllocator);
@@ -199,14 +209,12 @@ public class TestBasicOperation {
     }
 
     @Override
-    public Callable<PutResult> acceptPut(CallContext context,
-        FlightStream flightStream) {
+    public Runnable acceptPut(CallContext context, FlightStream flightStream, StreamListener<PutResult> ackStream) {
       return () -> {
         try (VectorSchemaRoot root = flightStream.getRoot()) {
           while (flightStream.next()) {
 
           }
-          return PutResult.getDefaultInstance();
         }
       };
     }

java/flight/src/test/java/org/apache/arrow/flight/TestCallOptions.java

@@ -69,11 +69,8 @@ public class TestCallOptions {
         BufferAllocator a = new RootAllocator(Long.MAX_VALUE);
         Producer producer = new Producer(a);
         FlightServer s =
-            FlightTestUtil.getStartedServer(
-                (port) -> FlightServer.builder(a, Location.forGrpcInsecure(FlightTestUtil.LOCALHOST, port), producer)
-                    .build());
-        FlightClient client = FlightClient.builder(a, Location.forGrpcInsecure(FlightTestUtil.LOCALHOST, s.getPort()))
-            .build()) {
+            FlightTestUtil.getStartedServer((location) -> FlightServer.builder(a, location, producer).build());
+        FlightClient client = FlightClient.builder(a, s.getLocation()).build()) {
       testFn.accept(client);
     } catch (InterruptedException | IOException e) {
       throw new RuntimeException(e);

java/flight/src/test/java/org/apache/arrow/flight/TestLargeMessage.java

@@ -19,10 +19,9 @@ package org.apache.arrow.flight;
 
 import java.util.Arrays;
 import java.util.List;
-import java.util.concurrent.Callable;
 import java.util.stream.Stream;
 
-import org.apache.arrow.flight.impl.Flight;
+import org.apache.arrow.flight.FlightProducer.StreamListener;
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.memory.RootAllocator;
 import org.apache.arrow.vector.IntVector;
@@ -43,11 +42,9 @@ public class TestLargeMessage {
     try (final BufferAllocator a = new RootAllocator(Long.MAX_VALUE);
          final Producer producer = new Producer(a);
          final FlightServer s =
-             FlightTestUtil.getStartedServer(
-                 (port) -> FlightServer.builder(a, Location.forGrpcInsecure("localhost", port), producer).build())) {
+             FlightTestUtil.getStartedServer((location) -> FlightServer.builder(a, location, producer).build())) {
 
-      try (FlightClient client = FlightClient
-          .builder(a, Location.forGrpcInsecure(FlightTestUtil.LOCALHOST, s.getPort())).build()) {
+      try (FlightClient client = FlightClient.builder(a, s.getLocation()).build()) {
         FlightStream stream = client.getStream(new Ticket(new byte[]{}));
         try (VectorSchemaRoot root = stream.getRoot()) {
           while (stream.next()) {
@@ -74,18 +71,32 @@ public class TestLargeMessage {
     try (final BufferAllocator a = new RootAllocator(Long.MAX_VALUE);
          final Producer producer = new Producer(a);
          final FlightServer s =
-             FlightTestUtil.getStartedServer(
-                 (port) -> FlightServer.builder(a, Location.forGrpcInsecure("localhost", port), producer).build()
+             FlightTestUtil.getStartedServer((location) -> FlightServer.builder(a, location, producer).build()
              )) {
 
-      try (FlightClient client = FlightClient
-          .builder(a, Location.forGrpcInsecure(FlightTestUtil.LOCALHOST, s.getPort())).build();
+      try (FlightClient client = FlightClient.builder(a, s.getLocation()).build();
            BufferAllocator testAllocator = a.newChildAllocator("testcase", 0, Long.MAX_VALUE);
            VectorSchemaRoot root = generateData(testAllocator)) {
-        final FlightClient.ClientStreamListener listener = client.startPut(FlightDescriptor.path("hello"), root);
+        final FlightClient.ClientStreamListener listener = client.startPut(FlightDescriptor.path("hello"), root,
+            new StreamListener<PutResult>() {
+              @Override
+              public void onNext(PutResult val) {
+
+              }
+
+              @Override
+              public void onError(Throwable t) {
+
+              }
+
+              @Override
+              public void onCompleted() {
+
+              }
+            });
         listener.putNext();
         listener.completed();
-        Assert.assertEquals(listener.getResult(), Flight.PutResult.getDefaultInstance());
+        listener.getResult();
       }
     }
   }
@@ -141,14 +152,12 @@ public class TestLargeMessage {
     }
 
     @Override
-    public Callable<Flight.PutResult> acceptPut(CallContext context,
-        FlightStream flightStream) {
+    public Runnable acceptPut(CallContext context, FlightStream flightStream, StreamListener<PutResult> ackStream) {
       return () -> {
         try (VectorSchemaRoot root = flightStream.getRoot()) {
           while (flightStream.next()) {
             ;
           }
-          return Flight.PutResult.getDefaultInstance();
         }
       };
     }

java/flight/src/test/java/org/apache/arrow/flight/auth/TestAuth.java

@@ -119,9 +119,9 @@ public class TestAuth {
       }
     };
 
-    server = FlightTestUtil.getStartedServer((port) -> FlightServer.builder(
+    server = FlightTestUtil.getStartedServer((location) -> FlightServer.builder(
         allocator,
-        Location.forGrpcInsecure("localhost", port),
+        location,
         new NoOpFlightProducer() {
           @Override
           public void listFlights(CallContext context, Criteria criteria,
@@ -150,8 +150,7 @@ public class TestAuth {
             listener.completed();
           }
         }).authHandler(new BasicServerAuthHandler(validator)).build());
-    client = FlightClient.builder(allocator, Location.forGrpcInsecure(FlightTestUtil.LOCALHOST, server.getPort()))
-        .build();
+    client = FlightClient.builder(allocator, server.getLocation()).build();
   }
 
   @After

java/flight/src/test/java/org/apache/arrow/flight/example/TestExampleServer.java

@@ -23,9 +23,11 @@ import org.apache.arrow.flight.FlightClient;
 import org.apache.arrow.flight.FlightClient.ClientStreamListener;
 import org.apache.arrow.flight.FlightDescriptor;
 import org.apache.arrow.flight.FlightInfo;
+import org.apache.arrow.flight.FlightProducer.StreamListener;
 import org.apache.arrow.flight.FlightStream;
 import org.apache.arrow.flight.FlightTestUtil;
 import org.apache.arrow.flight.Location;
+import org.apache.arrow.flight.PutResult;
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.memory.RootAllocator;
 import org.apache.arrow.util.AutoCloseables;
@@ -75,7 +77,20 @@ public class TestExampleServer {
     IntVector iv = new IntVector("c1", a);
 
     VectorSchemaRoot root = VectorSchemaRoot.of(iv);
-    ClientStreamListener listener = client.startPut(FlightDescriptor.path("hello"), root);
+    ClientStreamListener listener = client.startPut(FlightDescriptor.path("hello"), root,
+        new StreamListener<PutResult>() {
+          @Override
+          public void onNext(PutResult val) {
+          }
+
+          @Override
+          public void onError(Throwable t) {
+          }
+
+          @Override
+          public void onCompleted() {
+          }
+        });
 
     //batch 1
     root.allocateNew();

java/flight/src/test/java/org/apache/arrow/flight/perf/PerformanceTestServer.java

@@ -33,6 +33,7 @@ import org.apache.arrow.flight.FlightProducer;
 import org.apache.arrow.flight.FlightServer;
 import org.apache.arrow.flight.FlightStream;
 import org.apache.arrow.flight.Location;
+import org.apache.arrow.flight.NoOpFlightProducer;
 import org.apache.arrow.flight.Result;
 import org.apache.arrow.flight.Ticket;
 import org.apache.arrow.flight.impl.Flight.PutResult;
@@ -79,7 +80,7 @@ public class PerformanceTestServer implements AutoCloseable {
     AutoCloseables.close(flightServer, allocator);
   }
 
-  private final class PerfProducer implements FlightProducer {
+  private final class PerfProducer extends NoOpFlightProducer {
 
     @Override
     public void getStream(CallContext context, Ticket ticket,
@@ -145,11 +146,6 @@ public class PerformanceTestServer implements AutoCloseable {
       }
     }
 
-    @Override
-    public void listFlights(CallContext context, Criteria criteria,
-        StreamListener<FlightInfo> listener) {
-    }
-
     @Override
     public FlightInfo getFlightInfo(CallContext context,
         FlightDescriptor descriptor) {
@@ -181,24 +177,6 @@ public class PerformanceTestServer implements AutoCloseable {
         throw new RuntimeException(e);
       }
     }
-
-    @Override
-    public Callable<PutResult> acceptPut(CallContext context,
-        FlightStream flightStream) {
-      return null;
-    }
-
-    @Override
-    public void doAction(CallContext context, Action action,
-        StreamListener<Result> listener) {
-      listener.onCompleted();
-    }
-
-    @Override
-    public void listActions(CallContext context,
-        StreamListener<ActionType> listener) {
-    }
-
   }
 }
 

java/flight/src/test/java/org/apache/arrow/flight/perf/TestPerf.java

@@ -81,8 +81,7 @@ public class TestPerf {
       try (
           final BufferAllocator a = new RootAllocator(Long.MAX_VALUE);
           final PerformanceTestServer server =
-              FlightTestUtil.getStartedServer((port) -> new PerformanceTestServer(a,
-                  Location.forGrpcInsecure(FlightTestUtil.LOCALHOST, port)));
+              FlightTestUtil.getStartedServer((location) -> new PerformanceTestServer(a, location));
           final FlightClient client = FlightClient.builder(a, server.getLocation()).build();
       ) {
         final FlightInfo info = client.getInfo(getPerfFlightDescriptor(50_000_000L, 4095, 2));

python/pyarrow/_flight.pyx

@@ -44,6 +44,15 @@ cdef class FlightCallOptions:
         CFlightCallOptions options
 
     def __init__(self, timeout=None):
+        """Create call options.
+
+        Parameters
+        ----------
+        timeout : float or None
+            A timeout for the call, in seconds. None means that the
+            timeout defaults to an implementation-specific value.
+
+        """
         if timeout is not None:
             self.options.timeout = CTimeoutDuration(timeout)
 
@@ -70,14 +79,24 @@ cdef class Action:
         CAction action
 
     def __init__(self, action_type, buf):
+        """Create an action from a type and a buffer.
+
+        Parameters
+        ----------
+        action_type : bytes or str
+        buf : Buffer or bytes-like object
+        """
         self.action.type = tobytes(action_type)
         self.action.body = pyarrow_unwrap_buffer(as_buffer(buf))
 
     @property
     def type(self):
+        """The action type."""
         return frombytes(self.action.type)
 
+    @property
     def body(self):
+        """The action body (arguments for the action)."""
         return pyarrow_wrap_buffer(self.action.body)
 
     @staticmethod
@@ -92,10 +111,16 @@ _ActionType = collections.namedtuple('_ActionType', ['type', 'description'])
 
 
 class ActionType(_ActionType):
-    """A type of action executable on a Flight service."""
+    """A type of action that is executable on a Flight service."""
 
     def make_action(self, buf):
-        """Create an Action with this type."""
+        """Create an Action with this type.
+
+        Parameters
+        ----------
+        buf : obj
+            An Arrow buffer or Python bytes or bytes-like object.
+        """
         return Action(self.type, buf)
 
 
@@ -105,6 +130,12 @@ cdef class Result:
         unique_ptr[CResult] result
 
     def __init__(self, buf):
+        """Create a new result.
+
+        Parameters
+        ----------
+        buf : Buffer or bytes-like object
+        """
         self.result.reset(new CResult())
         self.result.get().body = pyarrow_unwrap_buffer(as_buffer(buf))
 
@@ -115,6 +146,23 @@ cdef class Result:
 
 
 class DescriptorType(enum.Enum):
+    """
+    The type of a FlightDescriptor.
+
+    Attributes
+    ----------
+
+    UNKNOWN
+        An unknown descriptor type.
+
+    PATH
+        A Flight stream represented by a path.
+
+    CMD
+        A Flight stream represented by an application-defined command.
+
+    """
+
     UNKNOWN = 0
     PATH = 1
     CMD = 2
@@ -151,6 +199,7 @@ cdef class FlightDescriptor:
 
     @property
     def descriptor_type(self):
+        """Get the type of this descriptor."""
         if self.descriptor.type == CDescriptorTypeUnknown:
             return DescriptorType.UNKNOWN
         elif self.descriptor.type == CDescriptorTypePath:
@@ -309,6 +358,7 @@ cdef class FlightEndpoint:
 
     @property
     def ticket(self):
+        """Get the ticket in this endpoint."""
         return Ticket(self.endpoint.ticket.ticket)
 
     @property
@@ -400,12 +450,110 @@ cdef class FlightInfo:
         return result
 
 
-cdef class FlightRecordBatchReader(_CRecordBatchReader, _ReadPandasOption):
+_AppMetadataRecordBatch = collections.namedtuple(
+    '_BatchWithMetadata',
+    ['record_batch', 'app_metadata'],
+)
+
+
+class AppMetadataRecordBatch(_AppMetadataRecordBatch):
+    """A RecordBatch with application metadata on the side."""
+
+
+cdef class MetadataRecordBatchReader(_CRecordBatchReader, _ReadPandasOption):
+    """A RecordBatchReader that also allows reading application metadata."""
+
     cdef dict __dict__
 
+    cdef readonly:
+        Schema schema
 
-cdef class FlightRecordBatchWriter(_CRecordBatchWriter):
-    pass
+    def read_with_metadata(self):
+        """Read the next RecordBatch along with any metadata.
+
+        Returns
+        -------
+        batch : RecordBatch
+            The next RecordBatch in the stream.
+        metadata : Buffer or None
+            Application-specific metadata for the batch as defined by
+            Flight.
+        """
+        cdef:
+            shared_ptr[CRecordBatch] batch
+            shared_ptr[CBuffer] app_metadata
+
+        with nogil:
+            check_status((<CMetadataRecordBatchReader*> self.reader.get())
+                         .ReadWithMetadata(&batch, &app_metadata))
+
+        if batch.get() == NULL:
+            raise StopIteration
+
+        if app_metadata != NULL:
+            metadata = pyarrow_wrap_buffer(app_metadata)
+        else:
+            metadata = None
+
+        return AppMetadataRecordBatch(pyarrow_wrap_batch(batch), metadata)
+
+
+cdef class FlightStreamWriter(_CRecordBatchWriter):
+    """A RecordBatchWriter that also allows writing application metadata."""
+
+    def write_with_metadata(self, RecordBatch batch, buf):
+        """Write a RecordBatch along with Flight metadata.
+
+        Parameters
+        ----------
+        batch : RecordBatch
+            The next RecordBatch in the stream.
+        buf : Buffer
+            Application-specific metadata for the batch as defined by
+            Flight.
+        """
+        cdef shared_ptr[CBuffer] c_buf = pyarrow_unwrap_buffer(as_buffer(buf))
+        with nogil:
+            check_status(
+                (<CFlightStreamWriter*> self.writer.get())
+                .WriteWithMetadata(deref(batch.batch),
+                                   c_buf,
+                                   1))
+
+
+cdef class FlightMetadataReader:
+    """A reader for Flight metadata messages sent during a DoPut."""
+
+    cdef:
+        unique_ptr[CFlightMetadataReader] reader
+
+    def read(self):
+        """Read the next metadata message."""
+        cdef shared_ptr[CBuffer] buf
+        with nogil:
+            check_status(self.reader.get().ReadMetadata(&buf))
+        if buf == NULL:
+            return None
+        return pyarrow_wrap_buffer(buf)
+
+
+cdef class FlightMetadataWriter:
+    """A sender for Flight metadata messages during a DoPut."""
+
+    cdef:
+        unique_ptr[CFlightMetadataWriter] writer
+
+    def write(self, message):
+        """Write the next metadata message.
+
+        Parameters
+        ----------
+        message : Buffer
+        """
+        cdef shared_ptr[CBuffer] buf = \
+            pyarrow_unwrap_buffer(as_buffer(message))
+        with nogil:
+            check_status(self.writer.get().WriteMetadata(deref(buf)))
 
 
 cdef class FlightClient:
@@ -451,7 +599,15 @@ cdef class FlightClient:
         return result
 
     def authenticate(self, auth_handler, options: FlightCallOptions = None):
-        """Authenticate to the server."""
+        """Authenticate to the server.
+
+        Parameters
+        ----------
+        auth_handler : ClientAuthHandler
+            The authentication mechanism to use.
+        options : FlightCallOptions
+            Options for this call.
+        """
         cdef:
             unique_ptr[CClientAuthHandler] handler
             CFlightCallOptions* c_options = FlightCallOptions.unwrap(options)
@@ -539,34 +695,53 @@ cdef class FlightClient:
         return result
 
     def do_get(self, ticket: Ticket, options: FlightCallOptions = None):
-        """Request the data for a flight."""
+        """Request the data for a flight.
+
+        Returns
+        -------
+        reader : MetadataRecordBatchReader
+        """
         cdef:
-            unique_ptr[CRecordBatchReader] reader
+            unique_ptr[CMetadataRecordBatchReader] reader
             CFlightCallOptions* c_options = FlightCallOptions.unwrap(options)
 
         with nogil:
-            check_status(self.client.get().DoGet(
-                deref(c_options), ticket.ticket, &reader))
-        result = FlightRecordBatchReader()
+            check_status(
+                self.client.get().DoGet(
+                    deref(c_options), ticket.ticket, &reader))
+        result = MetadataRecordBatchReader()
         result.reader.reset(reader.release())
+        result.schema = pyarrow_wrap_schema(result.reader.get().schema())
         return result
 
     def do_put(self, descriptor: FlightDescriptor, schema: Schema,
                options: FlightCallOptions = None):
-        """Upload data to a flight."""
+        """Upload data to a flight.
+
+        Returns
+        -------
+        writer : FlightStreamWriter
+        reader : FlightMetadataReader
+        """
         cdef:
             shared_ptr[CSchema] c_schema = pyarrow_unwrap_schema(schema)
-            unique_ptr[CRecordBatchWriter] writer
+            unique_ptr[CFlightStreamWriter] writer
+            unique_ptr[CFlightMetadataReader] metadata_reader
             CFlightCallOptions* c_options = FlightCallOptions.unwrap(options)
             CFlightDescriptor c_descriptor = \
                 FlightDescriptor.unwrap(descriptor)
+            FlightMetadataReader reader = FlightMetadataReader()
 
         with nogil:
             check_status(self.client.get().DoPut(
-                deref(c_options), c_descriptor, c_schema, &writer))
-        result = FlightRecordBatchWriter()
+                deref(c_options),
+                c_descriptor,
+                c_schema,
+                &writer,
+                &reader.reader))
+        result = FlightStreamWriter()
         result.writer.reset(writer.release())
-        return result
+        return result, reader
 
 
 cdef class FlightDataStream:
@@ -809,11 +984,22 @@ cdef void _data_stream_next(void* self, CFlightPayload* payload) except *:
         payload.ipc_message.metadata.reset(<CBuffer*> nullptr)
         return
 
+    if isinstance(result, (list, tuple)):
+        result, metadata = result
+    else:
+        result, metadata = result, None
+
     if isinstance(result, (Table, _CRecordBatchReader)):
+        if metadata:
+            raise ValueError("Can only return metadata alongside a "
+                             "RecordBatch.")
         result = RecordBatchStream(result)
 
     stream_schema = pyarrow_wrap_schema(stream.schema)
     if isinstance(result, FlightDataStream):
+        if metadata:
+            raise ValueError("Can only return metadata alongside a "
+                             "RecordBatch.")
         data_stream = unique_ptr[CFlightDataStream](
             (<FlightDataStream> result).to_stream())
         substream_schema = pyarrow_wrap_schema(data_stream.get().schema())
@@ -838,6 +1024,8 @@ cdef void _data_stream_next(void* self, CFlightPayload* payload) except *:
             deref(batch.batch),
             c_default_memory_pool(),
             &payload.ipc_message))
+        if metadata:
+            payload.app_metadata = pyarrow_unwrap_buffer(as_buffer(metadata))
     else:
         raise TypeError("GeneratorStream must be initialized with "
                         "an iterator of FlightDataStream, Table, "
@@ -880,17 +1068,22 @@ cdef void _get_flight_info(void* self, const CServerCallContext& context,
 
 
 cdef void _do_put(void* self, const CServerCallContext& context,
-                  unique_ptr[CFlightMessageReader] reader) except *:
+                  unique_ptr[CFlightMessageReader] reader,
+                  unique_ptr[CFlightMetadataWriter] writer) except *:
     """Callback for implementing Flight servers in Python."""
     cdef:
-        FlightRecordBatchReader py_reader = FlightRecordBatchReader()
+        MetadataRecordBatchReader py_reader = MetadataRecordBatchReader()
+        FlightMetadataWriter py_writer = FlightMetadataWriter()
         FlightDescriptor descriptor = \
             FlightDescriptor.__new__(FlightDescriptor)
 
     descriptor.descriptor = reader.get().descriptor()
     py_reader.reader.reset(reader.release())
+    py_reader.schema = pyarrow_wrap_schema(
+        py_reader.reader.get().schema())
+    py_writer.writer.reset(writer.release())
     (<object> self).do_put(ServerCallContext.wrap(context), descriptor,
-                           py_reader)
+                           py_reader, py_writer)
 
 
 cdef void _do_get(void* self, const CServerCallContext& context,
@@ -990,10 +1183,25 @@ cdef void _get_token(void* self, c_string* token) except *:
 
 
 cdef class ServerAuthHandler:
-    """Authentication middleware for a server."""
+    """Authentication middleware for a server.
+
+    To implement an authentication mechanism, subclass this class and
+    override its methods.
+
+    """
 
     def authenticate(self, outgoing, incoming):
-        """Conduct the handshake with the client."""
+        """Conduct the handshake with the client.
+
+        May raise an error if the client cannot authenticate.
+
+        Parameters
+        ----------
+        outgoing : ServerAuthSender
+            A channel to send messages to the client.
+        incoming : ServerAuthReader
+            A channel to read messages from the client.
+        """
         raise NotImplementedError
 
     def is_valid(self, token):
@@ -1003,6 +1211,11 @@ cdef class ServerAuthHandler:
         name the peer) or raise an exception (if the token is
         invalid).
 
+        Parameters
+        ----------
+        token : bytes
+            The authentication token from the client.
+
         """
         raise NotImplementedError
 
@@ -1017,7 +1230,15 @@ cdef class ClientAuthHandler:
     """Authentication plugin for a client."""
 
     def authenticate(self, outgoing, incoming):
-        """Conduct the handshake with the server."""
+        """Conduct the handshake with the server.
+
+        Parameters
+        ----------
+        outgoing : ClientAuthSender
+            A channel to send messages to the server.
+        incoming : ClientAuthReader
+            A channel to read messages from the server.
+        """
         raise NotImplementedError
 
     def get_token(self):
@@ -1032,12 +1253,26 @@ cdef class ClientAuthHandler:
 
 
 cdef class FlightServerBase:
-    """A Flight service definition."""
+    """A Flight service definition.
+
+    Override methods to define your Flight service.
+
+    """
 
     cdef:
         unique_ptr[PyFlightServer] server
 
     def run(self, location, auth_handler=None, tls_certificates=None):
+        """Start this server.
+
+        Parameters
+        ----------
+        location : Location
+        auth_handler : ServerAuthHandler
+            An authentication mechanism to use. May be None.
+        tls_certificates : list
+            A list of (certificate, key) pairs.
+        """
         cdef:
             PyFlightServerVtable vtable = PyFlightServerVtable()
             PyFlightServer* c_server
@@ -1078,7 +1313,8 @@ cdef class FlightServerBase:
     def get_flight_info(self, context, descriptor):
         raise NotImplementedError
 
-    def do_put(self, context, descriptor, reader):
+    def do_put(self, context, descriptor, reader,
+               writer: FlightMetadataWriter):
         raise NotImplementedError
 
     def do_get(self, context, ticket):

python/pyarrow/includes/libarrow_flight.pxd

@@ -112,22 +112,40 @@ cdef extern from "arrow/flight/api.h" namespace "arrow" nogil:
     cdef cppclass CSimpleFlightListing" arrow::flight::SimpleFlightListing":
         CSimpleFlightListing(vector[CFlightInfo]&& info)
 
-    cdef cppclass CFlightMessageReader \
-            " arrow::flight::FlightMessageReader"(CRecordBatchReader):
-        CFlightDescriptor& descriptor()
-
     cdef cppclass CFlightPayload" arrow::flight::FlightPayload":
         shared_ptr[CBuffer] descriptor
+        shared_ptr[CBuffer] app_metadata
         CIpcPayload ipc_message
 
     cdef cppclass CFlightDataStream" arrow::flight::FlightDataStream":
         shared_ptr[CSchema] schema()
         CStatus Next(CFlightPayload*)
 
+    cdef cppclass CMetadataRecordBatchReader \
+            " arrow::flight::MetadataRecordBatchReader"(CRecordBatchReader):
+        CStatus ReadWithMetadata(shared_ptr[CRecordBatch]* out,
+                                 shared_ptr[CBuffer]* app_metadata)
+
+    cdef cppclass CFlightMessageReader \
+            " arrow::flight::FlightMessageReader"(CMetadataRecordBatchReader):
+        CFlightDescriptor& descriptor()
+
+    cdef cppclass CFlightStreamWriter \
+            " arrow::flight::FlightStreamWriter"(CRecordBatchWriter):
+        CStatus WriteWithMetadata(const CRecordBatch& batch,
+                                  shared_ptr[CBuffer] app_metadata,
+                                  c_bool allow_64bit)
+
     cdef cppclass CRecordBatchStream \
             " arrow::flight::RecordBatchStream"(CFlightDataStream):
         CRecordBatchStream(shared_ptr[CRecordBatchReader]& reader)
 
+    cdef cppclass CFlightMetadataReader" arrow::flight::FlightMetadataReader":
+        CStatus ReadMetadata(shared_ptr[CBuffer]* out)
+
+    cdef cppclass CFlightMetadataWriter" arrow::flight::FlightMetadataWriter":
+        CStatus WriteMetadata(const CBuffer& message)
+
     cdef cppclass CServerAuthReader" arrow::flight::ServerAuthReader":
         CStatus Read(c_string* token)
 
@@ -193,11 +211,12 @@ cdef extern from "arrow/flight/api.h" namespace "arrow" nogil:
                               unique_ptr[CFlightInfo]* info)
 
         CStatus DoGet(CFlightCallOptions& options, CTicket& ticket,
-                      unique_ptr[CRecordBatchReader]* stream)
+                      unique_ptr[CMetadataRecordBatchReader]* stream)
         CStatus DoPut(CFlightCallOptions& options,
                       CFlightDescriptor& descriptor,
                       shared_ptr[CSchema]& schema,
-                      unique_ptr[CRecordBatchWriter]* stream)
+                      unique_ptr[CFlightStreamWriter]* stream,
+                      unique_ptr[CFlightMetadataReader]* reader)
 
 
 # Callbacks for implementing Flight servers
@@ -209,7 +228,8 @@ ctypedef void cb_get_flight_info(object, const CServerCallContext&,
                                  const CFlightDescriptor&,
                                  unique_ptr[CFlightInfo]*)
 ctypedef void cb_do_put(object, const CServerCallContext&,
-                        unique_ptr[CFlightMessageReader])
+                        unique_ptr[CFlightMessageReader],
+                        unique_ptr[CFlightMetadataWriter])
 ctypedef void cb_do_get(object, const CServerCallContext&,
                         const CTicket&,
                         unique_ptr[CFlightDataStream]*)

python/pyarrow/tests/test_flight.py

@@ -20,6 +20,7 @@ import base64
 import contextlib
 import os
 import socket
+import struct
 import tempfile
 import threading
 import time
@@ -114,17 +115,57 @@ class ConstantFlightServer(flight.FlightServerBase):
         return flight.RecordBatchStream(table)
 
 
+class MetadataFlightServer(flight.FlightServerBase):
+    """A Flight server that numbers incoming/outgoing data."""
+
+    def do_get(self, context, ticket):
+        data = [
+            pa.array([-10, -5, 0, 5, 10])
+        ]
+        table = pa.Table.from_arrays(data, names=['a'])
+        return flight.GeneratorStream(
+            table.schema,
+            self.number_batches(table))
+
+    def do_put(self, context, descriptor, reader, writer):
+        counter = 0
+        expected_data = [-10, -5, 0, 5, 10]
+        while True:
+            try:
+                batch, buf = reader.read_with_metadata()
+                assert batch.equals(pa.RecordBatch.from_arrays(
+                    [pa.array([expected_data[counter]])],
+                    ['a']
+                ))
+                assert buf is not None
+                client_counter, = struct.unpack('<i', buf.to_pybytes())
+                assert counter == client_counter
+                writer.write(struct.pack('<i', counter))
+                counter += 1
+            except StopIteration:
+                return
+
+    @staticmethod
+    def number_batches(table):
+        for idx, batch in enumerate(table.to_batches()):
+            buf = struct.pack('<i', idx)
+            yield batch, buf
+
+
 class EchoFlightServer(flight.FlightServerBase):
     """A Flight server that returns the last data uploaded."""
 
-    def __init__(self):
+    def __init__(self, expected_schema=None):
         super(EchoFlightServer, self).__init__()
         self.last_message = None
+        self.expected_schema = expected_schema
 
     def do_get(self, context, ticket):
         return flight.RecordBatchStream(self.last_message)
 
-    def do_put(self, context, descriptor, reader):
+    def do_put(self, context, descriptor, reader, writer):
+        if self.expected_schema:
+            assert self.expected_schema == reader.schema
         self.last_message = reader.read_all()
 
 
@@ -392,15 +433,23 @@ def test_flight_get_info():
                     reason="Unix sockets can't be tested on Windows")
 def test_flight_domain_socket():
     """Try a simple do_get call over a Unix domain socket."""
-    table = simple_ints_table()
-
     with tempfile.NamedTemporaryFile() as sock:
         sock.close()
         location = flight.Location.for_grpc_unix(sock.name)
         with flight_server(ConstantFlightServer,
                            location=location) as server_location:
             client = flight.FlightClient.connect(server_location)
-            data = client.do_get(flight.Ticket(b'ints')).read_all()
+
+            reader = client.do_get(flight.Ticket(b'ints'))
+            table = simple_ints_table()
+            assert reader.schema.equals(table.schema)
+            data = reader.read_all()
+            assert data.equals(table)
+
+            reader = client.do_get(flight.Ticket(b'dicts'))
+            table = simple_dicts_table()
+            assert reader.schema.equals(table.schema)
+            data = reader.read_all()
             assert data.equals(table)
 
 
@@ -415,10 +464,11 @@ def test_flight_large_message():
         pa.array(range(0, 10 * 1024 * 1024))
     ], names=['a'])
 
-    with flight_server(EchoFlightServer) as server_location:
+    with flight_server(EchoFlightServer,
+                       expected_schema=data.schema) as server_location:
         client = flight.FlightClient.connect(server_location)
-        writer = client.do_put(flight.FlightDescriptor.for_path('test'),
-                               data.schema)
+        writer, _ = client.do_put(flight.FlightDescriptor.for_path('test'),
+                                  data.schema)
         # Write a single giant chunk
         writer.write_table(data, 10 * 1024 * 1024)
         writer.close()
@@ -434,8 +484,8 @@ def test_flight_generator_stream():
 
     with flight_server(EchoStreamFlightServer) as server_location:
         client = flight.FlightClient.connect(server_location)
-        writer = client.do_put(flight.FlightDescriptor.for_path('test'),
-                               data.schema)
+        writer, _ = client.do_put(flight.FlightDescriptor.for_path('test'),
+                                  data.schema)
         writer.write_table(data)
         writer.close()
         result = client.do_get(flight.Ticket(b'')).read_all()
@@ -585,3 +635,50 @@ def test_tls_override_hostname():
             override_hostname="fakehostname")
         with pytest.raises(pa.ArrowIOError):
             client.do_get(flight.Ticket(b'ints'))
+
+
+def test_flight_do_get_metadata():
+    """Try a simple do_get call with metadata."""
+    data = [
+        pa.array([-10, -5, 0, 5, 10])
+    ]
+    table = pa.Table.from_arrays(data, names=['a'])
+
+    batches = []
+    with flight_server(MetadataFlightServer) as server_location:
+        client = flight.FlightClient.connect(server_location)
+        reader = client.do_get(flight.Ticket(b''))
+        idx = 0
+        while True:
+            try:
+                batch, metadata = reader.read_with_metadata()
+                batches.append(batch)
+                server_idx, = struct.unpack('<i', metadata.to_pybytes())
+                assert idx == server_idx
+                idx += 1
+            except StopIteration:
+                break
+        data = pa.Table.from_batches(batches)
+        assert data.equals(table)
+
+
+def test_flight_do_put_metadata():
+    """Try a simple do_put call with metadata."""
+    data = [
+        pa.array([-10, -5, 0, 5, 10])
+    ]
+    table = pa.Table.from_arrays(data, names=['a'])
+
+    with flight_server(MetadataFlightServer) as server_location:
+        client = flight.FlightClient.connect(server_location)
+        writer, metadata_reader = client.do_put(
+            flight.FlightDescriptor.for_path(''),
+            table.schema)
+        with writer:
+            for idx, batch in enumerate(table.to_batches(chunksize=1)):
+                metadata = struct.pack('<i', idx)
+                writer.write_with_metadata(batch, metadata)
+                buf = metadata_reader.read()
+                assert buf is not None
+                server_idx, = struct.unpack('<i', buf.to_pybytes())
+                assert idx == server_idx