Method TestSizeLimit

cpp/src/arrow/flight/test_definitions.cc:826–864

}

void DoPutTest::TestSizeLimit() {
  const int64_t size_limit = 4096;
  ASSERT_OK_AND_ASSIGN(auto location,
                       Location::ForScheme(transport(), "127.0.0.1", server_->port()));
  auto client_options = FlightClientOptions::Defaults();
  client_options.write_size_limit_bytes = size_limit;
  ASSERT_OK_AND_ASSIGN(auto client, FlightClient::Connect(location, client_options));

  auto descr = FlightDescriptor::Command("simple");
  // Batch is too large to fit in one message
  auto schema = arrow::schema({field("f1", arrow::int64())});
  auto batch = arrow::ConstantArrayGenerator::Zeroes(768, schema);
  auto batch1 = batch->Slice(0, 384);
  auto batch2 = batch->Slice(384);

  ASSERT_OK_AND_ASSIGN(auto do_put_result, client->DoPut(descr, schema));
  std::unique_ptr<FlightStreamWriter> writer = std::move(do_put_result.writer);
  std::unique_ptr<FlightMetadataReader> reader = std::move(do_put_result.reader);

  // Large batch will exceed the limit
  const auto status = writer->WriteRecordBatch(*batch);
  EXPECT_RAISES_WITH_MESSAGE_THAT(Invalid, ::testing::HasSubstr("exceeded soft limit"),
                                  status);
  auto detail = FlightWriteSizeStatusDetail::UnwrapStatus(status);
  ASSERT_NE(nullptr, detail);
  ASSERT_EQ(size_limit, detail->limit());
  ASSERT_GT(detail->actual(), size_limit);

  // But we can retry with smaller batches
  ASSERT_OK(writer->WriteRecordBatch(*batch1));
  ASSERT_OK(writer->WriteWithMetadata(*batch2, Buffer::FromString("1")));

  // Write a metadata-only message
  ASSERT_OK(writer->WriteMetadata(Buffer::FromString(kExpectedMetadata)));

  ASSERT_OK(writer->DoneWriting());
  ASSERT_OK(writer->Close());
  CheckBatches(descr, {batch1, batch2});
}
void DoPutTest::TestUndrained() {
  // Ensure if the application doesn't drain all messages, that the
  // server/transport does