diff --git a/cpp/src/arrow/flight/sql/odbc/odbc_impl/flight_sql_stream_chunk_buffer.cc b/cpp/src/arrow/flight/sql/odbc/odbc_impl/flight_sql_stream_chunk_buffer.cc index be788d9d08a..45c5f3a9126 100644 --- a/cpp/src/arrow/flight/sql/odbc/odbc_impl/flight_sql_stream_chunk_buffer.cc +++ b/cpp/src/arrow/flight/sql/odbc/odbc_impl/flight_sql_stream_chunk_buffer.cc @@ -33,8 +33,8 @@ FlightStreamChunkBuffer::FlightStreamChunkBuffer( arrow::Result> result; std::shared_ptr temp_flight_sql_client; auto endpoint_locations = endpoint.locations; - if (endpoint_locations.empty()) { - // list of Locations needs to be empty to proceed + if (endpoint_locations.empty() || + endpoint_locations[0].scheme() == Location::ReuseConnection().scheme()) { result = flight_sql_client.DoGet(call_options, ticket); } else { // If it is non-empty, the driver should create a FlightSqlClient to connect to one @@ -67,8 +67,8 @@ FlightStreamChunkBuffer::FlightStreamChunkBuffer( bool is_not_empty = result.ok() && (result.ValueOrDie().data != nullptr); // If result is valid, save the temp Flight SQL Client for future stream reader - // call. temp_flight_sql_client is intentionally null if the list of endpoint - // Locations is empty. + // call. temp_flight_sql_client is intentionally null if the endpoint is read + // from the original connection. // After all data is fetched from reader, the temp client is closed. if (is_not_ok || is_not_empty) { return std::make_optional( diff --git a/cpp/src/arrow/flight/sql/odbc/odbc_impl/flight_sql_stream_chunk_buffer_test.cc b/cpp/src/arrow/flight/sql/odbc/odbc_impl/flight_sql_stream_chunk_buffer_test.cc index cbe5cd8f7e5..79f20cd8c4a 100644 --- a/cpp/src/arrow/flight/sql/odbc/odbc_impl/flight_sql_stream_chunk_buffer_test.cc +++ b/cpp/src/arrow/flight/sql/odbc/odbc_impl/flight_sql_stream_chunk_buffer_test.cc @@ -128,4 +128,27 @@ TEST_F(FlightStreamChunkBufferTest, TestMultipleEndpointsInt) { // In total 10 batches should be returned. ASSERT_EQ(10, num_chunks); } + +TEST_F(FlightStreamChunkBufferTest, TestReuseConnectionLocation) { + ASSERT_OK_AND_ASSIGN(auto client, + arrow::flight::FlightClient::Connect(server_location1)); + FlightSqlClient server1_client(std::move(client)); + + FlightEndpoint endpoint( + {Ticket{"ticket-ints-1"}, {Location::ReuseConnection()}, std::nullopt, {}}); + FlightDescriptor descr{FlightDescriptor::PATH, "", {"examples", "ints"}}; + auto info_ptr = std::make_shared(arrow::flight::MakeFlightInfo( + *arrow::flight::ExampleIntSchema(), descr, {endpoint}, 1000, 100000, false, "")); + + FlightStreamChunkBuffer chunk_buffer(server1_client, FlightClientOptions::Defaults(), + FlightCallOptions(), info_ptr); + + FlightStreamChunk current_chunk; + int num_chunks = 0; + while (chunk_buffer.GetNext(¤t_chunk)) { + num_chunks++; + } + + ASSERT_EQ(5, num_chunks); +} } // namespace arrow::flight::sql::odbc