Skip to content

GH-54649: [C++][FlightRPC][ODBC] Fetch on the current connection for arrow-flight-reuse-connection locations #54652

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
AbhishekPathania wants to merge 1 commit into
base: main
Choose a base branch
from
+27 −4
Open

GH-54649: [C++][FlightRPC][ODBC] Fetch on the current connection for arrow-flight-reuse-connection locations #54652

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions cpp/src/arrow/flight/sql/odbc/odbc_impl/flight_sql_stream_chunk_buffer.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,8 +33,8 @@ FlightStreamChunkBuffer::FlightStreamChunkBuffer(
arrow::Result<std::unique_ptr<FlightStreamReader>> result;
std::shared_ptr<FlightSqlClient> 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
Expand Down Expand Up @@ -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(
Expand Down
23 changes: 23 additions & 0 deletions cpp/src/arrow/flight/sql/odbc/odbc_impl/flight_sql_stream_chunk_buffer_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<FlightInfo>(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(&current_chunk)) {
num_chunks++;
}

ASSERT_EQ(5, num_chunks);
}
} // namespace arrow::flight::sql::odbc
Loading