Method RunExchangeCounter

cpp/src/arrow/flight/test_flight_server.cc:226–256

Read some number of record batches from the client, send a metadata message back with the count, then echo the batches back.

// Read some number of record batches from the client, send a
// metadata message back with the count, then echo the batches back.
Status TestFlightServer::RunExchangeCounter(std::unique_ptr<FlightMessageReader> reader,
                                            std::unique_ptr<FlightMessageWriter> writer) {
  std::vector<std::shared_ptr<RecordBatch>> batches;
  FlightStreamChunk chunk;
  int chunks = 0;
  while (true) {
    ARROW_ASSIGN_OR_RAISE(chunk, reader->Next());
    if (!chunk.data && !chunk.app_metadata) {
      break;
    }
    if (chunk.data) {
      batches.push_back(chunk.data);
      chunks++;
    }
  }

  // Echo back the number of record batches read.
  std::shared_ptr<Buffer> buf = Buffer::FromString(std::to_string(chunks));
  RETURN_NOT_OK(writer->WriteMetadata(buf));
  // Echo the record batches themselves.
  if (chunks > 0) {
    ARROW_ASSIGN_OR_RAISE(auto schema, reader->GetSchema());
    RETURN_NOT_OK(writer->Begin(schema));

    for (const auto& batch : batches) {
      RETURN_NOT_OK(writer->WriteRecordBatch(*batch));
    }
  }

  return Status::OK();
}

// Read int64 batches from the client, each time sending back a
// batch with a running sum of columns.