diff --git a/cpp/src/arrow/flight/CMakeLists.txt b/cpp/src/arrow/flight/CMakeLists.txt index 663e7e241019..68d51da17dcd 100644 --- a/cpp/src/arrow/flight/CMakeLists.txt +++ b/cpp/src/arrow/flight/CMakeLists.txt @@ -291,7 +291,9 @@ if(ARROW_TESTING) STATIC_INSTALL_INTERFACE_LIBS ${ARROW_FLIGHT_TESTING_STATIC_INSTALL_INTERFACE_LIBS} PRIVATE_INCLUDES - "${Protobuf_INCLUDE_DIRS}") + "${Protobuf_INCLUDE_DIRS}" + SHARED_PRIVATE_LINK_LIBS + GTest::gmock) foreach(LIB_TARGET ${ARROW_FLIGHT_TESTING_LIBRARIES}) target_compile_definitions(${LIB_TARGET} PRIVATE ARROW_FLIGHT_EXPORTING) diff --git a/cpp/src/arrow/flight/sql/odbc/flight_sql/CMakeLists.txt b/cpp/src/arrow/flight/sql/odbc/flight_sql/CMakeLists.txt index 895eff8ed9a2..bd876804279c 100644 --- a/cpp/src/arrow/flight/sql/odbc/flight_sql/CMakeLists.txt +++ b/cpp/src/arrow/flight/sql/odbc/flight_sql/CMakeLists.txt @@ -136,9 +136,11 @@ add_arrow_test(arrow_odbc_spi_impl_test accessors/time_array_accessor_test.cc accessors/timestamp_array_accessor_test.cc flight_sql_connection_test.cc + flight_sql_stream_chunk_buffer_test.cc parse_table_types_test.cc json_converter_test.cc record_batch_transformer_test.cc utils_test.cc EXTRA_LINK_LIBS - arrow_odbc_spi_impl) + arrow_odbc_spi_impl + arrow_flight_testing_shared) diff --git a/cpp/src/arrow/flight/sql/odbc/flight_sql/flight_sql_auth_method.cc b/cpp/src/arrow/flight/sql/odbc/flight_sql/flight_sql_auth_method.cc index 6c265673837a..b2d57e5df857 100644 --- a/cpp/src/arrow/flight/sql/odbc/flight_sql/flight_sql_auth_method.cc +++ b/cpp/src/arrow/flight/sql/odbc/flight_sql/flight_sql_auth_method.cc @@ -58,9 +58,9 @@ class NoOpClientAuthHandler : public arrow::flight::ClientAuthHandler { arrow::Status Authenticate(arrow::flight::ClientAuthSender* outgoing, arrow::flight::ClientAuthReader* incoming) override { - // Write a blank string. The server should ignore this and just accept any Handshake + // Return OK Status. The server should ignore this and just accept any Handshake // request. - return outgoing->Write(std::string()); + return arrow::Status::OK(); } arrow::Status GetToken(std::string* token) override { diff --git a/cpp/src/arrow/flight/sql/odbc/flight_sql/flight_sql_connection.cc b/cpp/src/arrow/flight/sql/odbc/flight_sql/flight_sql_connection.cc index 6d5d95865ba6..c87c394fc314 100644 --- a/cpp/src/arrow/flight/sql/odbc/flight_sql/flight_sql_connection.cc +++ b/cpp/src/arrow/flight/sql/odbc/flight_sql/flight_sql_connection.cc @@ -112,6 +112,8 @@ inline std::string GetCerts() { return ""; } #endif const std::set BUILT_IN_PROPERTIES = { + FlightSqlConnection::DRIVER, + FlightSqlConnection::DSN, FlightSqlConnection::HOST, FlightSqlConnection::PORT, FlightSqlConnection::USER, @@ -165,15 +167,15 @@ void FlightSqlConnection::Connect(const ConnPropertyMap& properties, auto flight_ssl_configs = LoadFlightSslConfigs(properties); Location location = BuildLocation(properties, missing_attr, flight_ssl_configs); - FlightClientOptions client_options = + client_options_ = BuildFlightClientOptions(properties, missing_attr, flight_ssl_configs); const std::shared_ptr& cookie_factory = arrow::flight::GetCookieFactory(); - client_options.middleware.push_back(cookie_factory); + client_options_.middleware.push_back(cookie_factory); std::unique_ptr flight_client; - ThrowIfNotOK(FlightClient::Connect(location, client_options).Value(&flight_client)); + ThrowIfNotOK(FlightClient::Connect(location, client_options_).Value(&flight_client)); PopulateMetadataSettings(properties); PopulateCallOptions(properties); @@ -377,7 +379,7 @@ void FlightSqlConnection::Close() { std::shared_ptr FlightSqlConnection::CreateStatement() { return std::shared_ptr(new FlightSqlStatement( - diagnostics_, *sql_client_, call_options_, metadata_settings_)); + diagnostics_, *sql_client_, client_options_, call_options_, metadata_settings_)); } bool FlightSqlConnection::SetAttribute(Connection::AttributeId attribute, @@ -423,7 +425,7 @@ FlightSqlConnection::FlightSqlConnection(OdbcVersion odbc_version, const std::string& driver_version) : diagnostics_("Apache Arrow", "Flight SQL", odbc_version), odbc_version_(odbc_version), - info_(call_options_, sql_client_, driver_version), + info_(client_options_, call_options_, sql_client_, driver_version), closed_(true) { attribute_[CONNECTION_DEAD] = static_cast(SQL_TRUE); attribute_[LOGIN_TIMEOUT] = static_cast(0); diff --git a/cpp/src/arrow/flight/sql/odbc/flight_sql/flight_sql_result_set.cc b/cpp/src/arrow/flight/sql/odbc/flight_sql/flight_sql_result_set.cc index 6744537ae5be..258c810996ac 100644 --- a/cpp/src/arrow/flight/sql/odbc/flight_sql/flight_sql_result_set.cc +++ b/cpp/src/arrow/flight/sql/odbc/flight_sql/flight_sql_result_set.cc @@ -44,13 +44,14 @@ using odbcabstraction::DriverException; FlightSqlResultSet::FlightSqlResultSet( FlightSqlClient& flight_sql_client, + const arrow::flight::FlightClientOptions& client_options, const arrow::flight::FlightCallOptions& call_options, const std::shared_ptr& flight_info, const std::shared_ptr& transformer, odbcabstraction::Diagnostics& diagnostics, const odbcabstraction::MetadataSettings& metadata_settings) : metadata_settings_(metadata_settings), - chunk_buffer_(flight_sql_client, call_options, flight_info, + chunk_buffer_(flight_sql_client, client_options, call_options, flight_info, metadata_settings_.chunk_buffer_capacity_), transformer_(transformer), metadata_(transformer diff --git a/cpp/src/arrow/flight/sql/odbc/flight_sql/flight_sql_result_set.h b/cpp/src/arrow/flight/sql/odbc/flight_sql/flight_sql_result_set.h index c82a2d835435..5a03b16f0660 100644 --- a/cpp/src/arrow/flight/sql/odbc/flight_sql/flight_sql_result_set.h +++ b/cpp/src/arrow/flight/sql/odbc/flight_sql/flight_sql_result_set.h @@ -63,6 +63,7 @@ class FlightSqlResultSet : public ResultSet { ~FlightSqlResultSet() override; FlightSqlResultSet(FlightSqlClient& flight_sql_client, + const arrow::flight::FlightClientOptions& client_options, const arrow::flight::FlightCallOptions& call_options, const std::shared_ptr& flight_info, const std::shared_ptr& transformer, diff --git a/cpp/src/arrow/flight/sql/odbc/flight_sql/flight_sql_statement.cc b/cpp/src/arrow/flight/sql/odbc/flight_sql/flight_sql_statement.cc index 4f94ef701688..efe333d836a4 100644 --- a/cpp/src/arrow/flight/sql/odbc/flight_sql/flight_sql_statement.cc +++ b/cpp/src/arrow/flight/sql/odbc/flight_sql/flight_sql_statement.cc @@ -64,11 +64,12 @@ void ClosePreparedStatementIfAny( FlightSqlStatement::FlightSqlStatement( const odbcabstraction::Diagnostics& diagnostics, FlightSqlClient& sql_client, - FlightCallOptions call_options, + arrow::flight::FlightClientOptions client_options, FlightCallOptions call_options, const odbcabstraction::MetadataSettings& metadata_settings) : diagnostics_("Apache Arrow", diagnostics.GetDataSourceComponent(), diagnostics.GetOdbcVersion()), sql_client_(sql_client), + client_options_(std::move(client_options)), call_options_(std::move(call_options)), metadata_settings_(metadata_settings) { attribute_[METADATA_ID] = static_cast(SQL_FALSE); @@ -135,8 +136,8 @@ bool FlightSqlStatement::ExecutePrepared() { ThrowIfNotOK(result.status()); current_result_set_ = std::make_shared( - sql_client_, call_options_, result.ValueOrDie(), nullptr, diagnostics_, - metadata_settings_); + sql_client_, client_options_, call_options_, result.ValueOrDie(), nullptr, + diagnostics_, metadata_settings_); return true; } @@ -148,8 +149,8 @@ bool FlightSqlStatement::Execute(const std::string& query) { ThrowIfNotOK(result.status()); current_result_set_ = std::make_shared( - sql_client_, call_options_, result.ValueOrDie(), nullptr, diagnostics_, - metadata_settings_); + sql_client_, client_options_, call_options_, result.ValueOrDie(), nullptr, + diagnostics_, metadata_settings_); return true; } @@ -170,27 +171,29 @@ std::shared_ptr FlightSqlStatement::GetTables( if ((catalog_name && *catalog_name == "%") && (schema_name && schema_name->empty()) && (table_name && table_name->empty())) { - current_result_set_ = GetTablesForSQLAllCatalogs( - column_names, call_options_, sql_client_, diagnostics_, metadata_settings_); + current_result_set_ = + GetTablesForSQLAllCatalogs(column_names, client_options_, call_options_, + sql_client_, diagnostics_, metadata_settings_); } else if ((catalog_name && catalog_name->empty()) && (schema_name && *schema_name == "%") && (table_name && table_name->empty())) { - current_result_set_ = - GetTablesForSQLAllDbSchemas(column_names, call_options_, sql_client_, schema_name, - diagnostics_, metadata_settings_); + current_result_set_ = GetTablesForSQLAllDbSchemas( + column_names, client_options_, call_options_, sql_client_, schema_name, + diagnostics_, metadata_settings_); } else if ((catalog_name && catalog_name->empty()) && (schema_name && schema_name->empty()) && (table_name && table_name->empty()) && (table_type && *table_type == "%")) { - current_result_set_ = GetTablesForSQLAllTableTypes( - column_names, call_options_, sql_client_, diagnostics_, metadata_settings_); + current_result_set_ = + GetTablesForSQLAllTableTypes(column_names, client_options_, call_options_, + sql_client_, diagnostics_, metadata_settings_); } else { if (table_type) { ParseTableTypes(*table_type, table_types); } current_result_set_ = GetTablesForGenericUse( - column_names, call_options_, sql_client_, catalog_name, schema_name, table_name, - table_types, diagnostics_, metadata_settings_); + column_names, client_options_, call_options_, sql_client_, catalog_name, + schema_name, table_name, table_types, diagnostics_, metadata_settings_); } return current_result_set_; @@ -228,9 +231,9 @@ std::shared_ptr FlightSqlStatement::GetColumns_V2( auto transformer = std::make_shared( metadata_settings_, odbcabstraction::V_2, column_name); - current_result_set_ = - std::make_shared(sql_client_, call_options_, flight_info, - transformer, diagnostics_, metadata_settings_); + current_result_set_ = std::make_shared( + sql_client_, client_options_, call_options_, flight_info, transformer, diagnostics_, + metadata_settings_); return current_result_set_; } @@ -249,9 +252,9 @@ std::shared_ptr FlightSqlStatement::GetColumns_V3( auto transformer = std::make_shared( metadata_settings_, odbcabstraction::V_3, column_name); - current_result_set_ = - std::make_shared(sql_client_, call_options_, flight_info, - transformer, diagnostics_, metadata_settings_); + current_result_set_ = std::make_shared( + sql_client_, client_options_, call_options_, flight_info, transformer, diagnostics_, + metadata_settings_); return current_result_set_; } @@ -267,9 +270,9 @@ std::shared_ptr FlightSqlStatement::GetTypeInfo_V2(int16_t data_type) auto transformer = std::make_shared( metadata_settings_, odbcabstraction::V_2, data_type); - current_result_set_ = - std::make_shared(sql_client_, call_options_, flight_info, - transformer, diagnostics_, metadata_settings_); + current_result_set_ = std::make_shared( + sql_client_, client_options_, call_options_, flight_info, transformer, diagnostics_, + metadata_settings_); return current_result_set_; } @@ -285,9 +288,9 @@ std::shared_ptr FlightSqlStatement::GetTypeInfo_V3(int16_t data_type) auto transformer = std::make_shared( metadata_settings_, odbcabstraction::V_3, data_type); - current_result_set_ = - std::make_shared(sql_client_, call_options_, flight_info, - transformer, diagnostics_, metadata_settings_); + current_result_set_ = std::make_shared( + sql_client_, client_options_, call_options_, flight_info, transformer, diagnostics_, + metadata_settings_); return current_result_set_; } diff --git a/cpp/src/arrow/flight/sql/odbc/flight_sql/flight_sql_statement.h b/cpp/src/arrow/flight/sql/odbc/flight_sql/flight_sql_statement.h index 56341223259d..00fe9137f51b 100644 --- a/cpp/src/arrow/flight/sql/odbc/flight_sql/flight_sql_statement.h +++ b/cpp/src/arrow/flight/sql/odbc/flight_sql/flight_sql_statement.h @@ -33,6 +33,7 @@ class FlightSqlStatement : public odbcabstraction::Statement { private: odbcabstraction::Diagnostics diagnostics_; std::map attribute_; + arrow::flight::FlightClientOptions client_options_; arrow::flight::FlightCallOptions call_options_; arrow::flight::sql::FlightSqlClient& sql_client_; std::shared_ptr current_result_set_; @@ -48,6 +49,7 @@ class FlightSqlStatement : public odbcabstraction::Statement { public: FlightSqlStatement(const odbcabstraction::Diagnostics& diagnostics, arrow::flight::sql::FlightSqlClient& sql_client, + arrow::flight::FlightClientOptions client_options, arrow::flight::FlightCallOptions call_options, const odbcabstraction::MetadataSettings& metadata_settings); ~FlightSqlStatement(); diff --git a/cpp/src/arrow/flight/sql/odbc/flight_sql/flight_sql_statement_get_tables.cc b/cpp/src/arrow/flight/sql/odbc/flight_sql/flight_sql_statement_get_tables.cc index a3cdf9768d20..ba3389ccba97 100644 --- a/cpp/src/arrow/flight/sql/odbc/flight_sql/flight_sql_statement_get_tables.cc +++ b/cpp/src/arrow/flight/sql/odbc/flight_sql/flight_sql_statement_get_tables.cc @@ -72,8 +72,9 @@ void ParseTableTypes(const std::string& table_type, } std::shared_ptr GetTablesForSQLAllCatalogs( - const ColumnNames& names, FlightCallOptions& call_options, - FlightSqlClient& sql_client, odbcabstraction::Diagnostics& diagnostics, + const ColumnNames& names, FlightClientOptions& client_options, + FlightCallOptions& call_options, FlightSqlClient& sql_client, + odbcabstraction::Diagnostics& diagnostics, const odbcabstraction::MetadataSettings& metadata_settings) { Result> result = sql_client.GetCatalogs(call_options); @@ -92,14 +93,15 @@ std::shared_ptr GetTablesForSQLAllCatalogs( .AddFieldOfNulls(names.remarks_column, arrow::utf8()) .Build(); - return std::make_shared( - sql_client, call_options, flight_info, transformer, diagnostics, metadata_settings); + return std::make_shared(sql_client, client_options, call_options, + flight_info, transformer, diagnostics, + metadata_settings); } std::shared_ptr GetTablesForSQLAllDbSchemas( - const ColumnNames& names, FlightCallOptions& call_options, - FlightSqlClient& sql_client, const std::string* schema_name, - odbcabstraction::Diagnostics& diagnostics, + const ColumnNames& names, FlightClientOptions& client_options, + FlightCallOptions& call_options, FlightSqlClient& sql_client, + const std::string* schema_name, odbcabstraction::Diagnostics& diagnostics, const odbcabstraction::MetadataSettings& metadata_settings) { Result> result = sql_client.GetDbSchemas(call_options, nullptr, schema_name); @@ -119,13 +121,15 @@ std::shared_ptr GetTablesForSQLAllDbSchemas( .AddFieldOfNulls(names.remarks_column, arrow::utf8()) .Build(); - return std::make_shared( - sql_client, call_options, flight_info, transformer, diagnostics, metadata_settings); + return std::make_shared(sql_client, client_options, call_options, + flight_info, transformer, diagnostics, + metadata_settings); } std::shared_ptr GetTablesForSQLAllTableTypes( - const ColumnNames& names, FlightCallOptions& call_options, - FlightSqlClient& sql_client, odbcabstraction::Diagnostics& diagnostics, + const ColumnNames& names, FlightClientOptions& client_options, + FlightCallOptions& call_options, FlightSqlClient& sql_client, + odbcabstraction::Diagnostics& diagnostics, const odbcabstraction::MetadataSettings& metadata_settings) { Result> result = sql_client.GetTableTypes(call_options); @@ -144,15 +148,16 @@ std::shared_ptr GetTablesForSQLAllTableTypes( .AddFieldOfNulls(names.remarks_column, arrow::utf8()) .Build(); - return std::make_shared( - sql_client, call_options, flight_info, transformer, diagnostics, metadata_settings); + return std::make_shared(sql_client, client_options, call_options, + flight_info, transformer, diagnostics, + metadata_settings); } std::shared_ptr GetTablesForGenericUse( - const ColumnNames& names, FlightCallOptions& call_options, - FlightSqlClient& sql_client, const std::string* catalog_name, - const std::string* schema_name, const std::string* table_name, - const std::vector& table_types, + const ColumnNames& names, FlightClientOptions& client_options, + FlightCallOptions& call_options, FlightSqlClient& sql_client, + const std::string* catalog_name, const std::string* schema_name, + const std::string* table_name, const std::vector& table_types, odbcabstraction::Diagnostics& diagnostics, const odbcabstraction::MetadataSettings& metadata_settings) { Result> result = sql_client.GetTables( @@ -173,8 +178,9 @@ std::shared_ptr GetTablesForGenericUse( .AddFieldOfNulls(names.remarks_column, arrow::utf8()) .Build(); - return std::make_shared( - sql_client, call_options, flight_info, transformer, diagnostics, metadata_settings); + return std::make_shared(sql_client, client_options, call_options, + flight_info, transformer, diagnostics, + metadata_settings); } } // namespace flight_sql diff --git a/cpp/src/arrow/flight/sql/odbc/flight_sql/flight_sql_statement_get_tables.h b/cpp/src/arrow/flight/sql/odbc/flight_sql/flight_sql_statement_get_tables.h index 8f0dc5fef6da..0f5ac461f3ff 100644 --- a/cpp/src/arrow/flight/sql/odbc/flight_sql/flight_sql_statement_get_tables.h +++ b/cpp/src/arrow/flight/sql/odbc/flight_sql/flight_sql_statement_get_tables.h @@ -30,6 +30,7 @@ namespace driver { namespace flight_sql { using arrow::flight::FlightCallOptions; +using arrow::flight::FlightClientOptions; using arrow::flight::sql::FlightSqlClient; using odbcabstraction::MetadataSettings; using odbcabstraction::ResultSet; @@ -46,26 +47,28 @@ void ParseTableTypes(const std::string& table_type, std::vector& table_types); std::shared_ptr GetTablesForSQLAllCatalogs( - const ColumnNames& column_names, FlightCallOptions& call_options, - FlightSqlClient& sql_client, odbcabstraction::Diagnostics& diagnostics, + const ColumnNames& column_names, FlightClientOptions& client_options, + FlightCallOptions& call_options, FlightSqlClient& sql_client, + odbcabstraction::Diagnostics& diagnostics, const odbcabstraction::MetadataSettings& metadata_settings); std::shared_ptr GetTablesForSQLAllDbSchemas( - const ColumnNames& column_names, FlightCallOptions& call_options, - FlightSqlClient& sql_client, const std::string* schema_name, - odbcabstraction::Diagnostics& diagnostics, + const ColumnNames& column_names, FlightClientOptions& client_options, + FlightCallOptions& call_options, FlightSqlClient& sql_client, + const std::string* schema_name, odbcabstraction::Diagnostics& diagnostics, const odbcabstraction::MetadataSettings& metadata_settings); std::shared_ptr GetTablesForSQLAllTableTypes( - const ColumnNames& column_names, FlightCallOptions& call_options, - FlightSqlClient& sql_client, odbcabstraction::Diagnostics& diagnostics, + const ColumnNames& column_names, FlightClientOptions& client_options, + FlightCallOptions& call_options, FlightSqlClient& sql_client, + odbcabstraction::Diagnostics& diagnostics, const odbcabstraction::MetadataSettings& metadata_settings); std::shared_ptr GetTablesForGenericUse( - const ColumnNames& column_names, FlightCallOptions& call_options, - FlightSqlClient& sql_client, const std::string* catalog_name, - const std::string* schema_name, const std::string* table_name, - const std::vector& table_types, + const ColumnNames& column_names, FlightClientOptions& client_options, + FlightCallOptions& call_options, FlightSqlClient& sql_client, + const std::string* catalog_name, const std::string* schema_name, + const std::string* table_name, const std::vector& table_types, odbcabstraction::Diagnostics& diagnostics, const odbcabstraction::MetadataSettings& metadata_settings); } // namespace flight_sql diff --git a/cpp/src/arrow/flight/sql/odbc/flight_sql/flight_sql_stream_chunk_buffer.cc b/cpp/src/arrow/flight/sql/odbc/flight_sql/flight_sql_stream_chunk_buffer.cc index 093a46dfe838..7da2d6ca89de 100644 --- a/cpp/src/arrow/flight/sql/odbc/flight_sql/flight_sql_stream_chunk_buffer.cc +++ b/cpp/src/arrow/flight/sql/odbc/flight_sql/flight_sql_stream_chunk_buffer.cc @@ -21,38 +21,71 @@ namespace driver { namespace flight_sql { +using arrow::flight::FlightClient; using arrow::flight::FlightEndpoint; FlightStreamChunkBuffer::FlightStreamChunkBuffer( FlightSqlClient& flight_sql_client, + const arrow::flight::FlightClientOptions& client_options, const arrow::flight::FlightCallOptions& call_options, const std::shared_ptr& flight_info, size_t queue_capacity) : queue_(queue_capacity) { - // FIXME: Endpoint iteration should consider endpoints may be at different hosts for (const auto& endpoint : flight_info->endpoints()) { const arrow::flight::Ticket& ticket = endpoint.ticket; - auto result = flight_sql_client.DoGet(call_options, ticket); + 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 + result = flight_sql_client.DoGet(call_options, ticket); + } else { + // If it is non-empty, the driver should create a FlightSqlClient to connect to one + // of the specified Locations directly. + + // GH-47117: Currently a new FlightClient will be made for each partition that + // returns a non-empty Location, which is then disposed of. It may be better to + // cache clients because a server may report the same Locations. It would also be + // good to identify when the reported Location is the same as the original + // connection's Location and skip creating a FlightClient in that scenario. + + std::unique_ptr temp_flight_client; + ThrowIfNotOK(FlightClient::Connect(endpoint_locations[0], client_options) + .Value(&temp_flight_client)); + temp_flight_sql_client.reset(new FlightSqlClient(std::move(temp_flight_client))); + + result = temp_flight_sql_client->DoGet(call_options, ticket); + } + ThrowIfNotOK(result.status()); std::shared_ptr stream_reader_ptr(std::move(result.ValueOrDie())); - BlockingQueue>::Supplier supplier = [=] { + BlockingQueue, + std::shared_ptr>>::Supplier supplier = [=] { auto result = stream_reader_ptr->Next(); bool isNotOk = !result.ok(); bool isNotEmpty = result.ok() && (result.ValueOrDie().data != nullptr); - return boost::make_optional(isNotOk || isNotEmpty, std::move(result)); + // 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. + // After all data is fetched from reader, the temp client is closed. + return boost::make_optional( + isNotOk || isNotEmpty, + std::make_pair(std::move(result), temp_flight_sql_client)); }; queue_.AddProducer(std::move(supplier)); } } bool FlightStreamChunkBuffer::GetNext(FlightStreamChunk* chunk) { - Result result; - if (!queue_.Pop(&result)) { + std::pair, std::shared_ptr> + closeableEndpointStreamPair; + if (!queue_.Pop(&closeableEndpointStreamPair)) { return false; } + Result result = closeableEndpointStreamPair.first; if (!result.status().ok()) { Close(); throw odbcabstraction::DriverException(result.status().message()); diff --git a/cpp/src/arrow/flight/sql/odbc/flight_sql/flight_sql_stream_chunk_buffer.h b/cpp/src/arrow/flight/sql/odbc/flight_sql/flight_sql_stream_chunk_buffer.h index 864c025d8b37..5d5616a4f027 100644 --- a/cpp/src/arrow/flight/sql/odbc/flight_sql/flight_sql_stream_chunk_buffer.h +++ b/cpp/src/arrow/flight/sql/odbc/flight_sql/flight_sql_stream_chunk_buffer.h @@ -32,10 +32,12 @@ using arrow::flight::sql::FlightSqlClient; using driver::odbcabstraction::BlockingQueue; class FlightStreamChunkBuffer { - BlockingQueue> queue_; + BlockingQueue, std::shared_ptr>> + queue_; public: FlightStreamChunkBuffer(FlightSqlClient& flight_sql_client, + const arrow::flight::FlightClientOptions& client_options, const arrow::flight::FlightCallOptions& call_options, const std::shared_ptr& flight_info, size_t queue_capacity = 5); diff --git a/cpp/src/arrow/flight/sql/odbc/flight_sql/flight_sql_stream_chunk_buffer_test.cc b/cpp/src/arrow/flight/sql/odbc/flight_sql/flight_sql_stream_chunk_buffer_test.cc new file mode 100644 index 000000000000..6857b53f5c25 --- /dev/null +++ b/cpp/src/arrow/flight/sql/odbc/flight_sql/flight_sql_stream_chunk_buffer_test.cc @@ -0,0 +1,137 @@ +// 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. + +#include "arrow/array.h" + +#include "arrow/testing/gtest_util.h" + +#include "arrow/flight/sql/odbc/flight_sql/flight_sql_stream_chunk_buffer.h" +#include "arrow/flight/sql/odbc/flight_sql/json_converter.h" +#include "arrow/flight/test_flight_server.h" +#include "arrow/flight/test_util.h" +#include "gtest/gtest.h" + +namespace driver { +namespace flight_sql { + +using arrow::Array; +using arrow::flight::FlightCallOptions; +using arrow::flight::FlightClientOptions; +using arrow::flight::FlightDescriptor; +using arrow::flight::FlightEndpoint; +using arrow::flight::Location; +using arrow::flight::Ticket; +using arrow::flight::sql::FlightSqlClient; + +class FlightStreamChunkBufferTest : public ::testing::Test { + // Sets up two mock servers for each test case. + // This is for testing endpoint iteration only. + + protected: + void SetUp() override { + // Set up server 1 + server1 = std::make_shared(); + ASSERT_OK_AND_ASSIGN(auto location1, Location::ForGrpcTcp("0.0.0.0", 0)); + arrow::flight::FlightServerOptions options1(location1); + ASSERT_OK(server1->Init(options1)); + ASSERT_OK_AND_ASSIGN(server_location1, + Location::ForGrpcTcp("localhost", server1->port())); + + // Set up server 2 + server2 = std::make_shared(); + ASSERT_OK_AND_ASSIGN(auto location2, Location::ForGrpcTcp("0.0.0.0", 0)); + arrow::flight::FlightServerOptions options2(location2); + ASSERT_OK(server2->Init(options2)); + ASSERT_OK_AND_ASSIGN(server_location2, + Location::ForGrpcTcp("localhost", server2->port())); + + // Make SQL Client that is connected to server 1 + ASSERT_OK_AND_ASSIGN(auto client, arrow::flight::FlightClient::Connect(location1)); + sql_client.reset(new FlightSqlClient(std::move(client))); + } + + void TearDown() override { + ASSERT_OK(server1->Shutdown()); + ASSERT_OK(server2->Shutdown()); + } + + public: + arrow::flight::Location server_location1; + std::shared_ptr server1; + arrow::flight::Location server_location2; + std::shared_ptr server2; + std::shared_ptr sql_client; +}; + +FlightInfo MultipleEndpointsFlightInfo(Location location1, Location location2) { + // Sever will generate random data for `ticket-ints-1` + FlightEndpoint endpoint1({Ticket{"ticket-ints-1"}, {location1}, std::nullopt, {}}); + FlightEndpoint endpoint2({Ticket{"ticket-ints-1"}, {location2}, std::nullopt, {}}); + + FlightDescriptor descr1{FlightDescriptor::PATH, "", {"examples", "ints"}}; + + auto schema1 = arrow::flight::ExampleIntSchema(); + + return arrow::flight::MakeFlightInfo(*schema1, descr1, {endpoint1, endpoint2}, 1000, + 100000, false, ""); +} + +void verifyArraysContainIntsOnly(std::shared_ptr intArray) { + for (int64_t i = 0; i < intArray->length(); ++i) { + // null values are accepted + if (!intArray->IsNull(i)) { + auto scalar_data = intArray->GetScalar(i).ValueOrDie(); + std::string scalar_str = ConvertToJson(*scalar_data); + ASSERT_TRUE(std::all_of(scalar_str.begin(), scalar_str.end(), ::isdigit)); + } + } +} + +TEST_F(FlightStreamChunkBufferTest, TestMultipleEndpointsInt) { + FlightClientOptions client_options = FlightClientOptions::Defaults(); + FlightCallOptions options; + FlightInfo info = MultipleEndpointsFlightInfo(server_location1, server_location2); + std::shared_ptr info_ptr = std::make_shared(info); + + FlightStreamChunkBuffer chunk_buffer(*sql_client, client_options, options, info_ptr); + + FlightStreamChunk current_chunk; + + // Server returns 5 batch of results from each endpoints. + // Each batch contains 8 columns + int num_chunks = 0; + while (chunk_buffer.GetNext(¤t_chunk)) { + num_chunks++; + + int num_cols = current_chunk.data->num_columns(); + EXPECT_EQ(num_cols, 8); + + for (int i = 0; i < num_cols; i++) { + auto array = current_chunk.data->column(i); + // Each array has random length + EXPECT_GT(array->length(), 0); + + verifyArraysContainIntsOnly(array); + } + } + + // Verify 5 batches of data is returned by each of the two endpoints. + // In total 10 batches should be returned. + EXPECT_EQ(num_chunks, 10); +} +} // namespace flight_sql +} // namespace driver diff --git a/cpp/src/arrow/flight/sql/odbc/flight_sql/get_info_cache.cc b/cpp/src/arrow/flight/sql/odbc/flight_sql/get_info_cache.cc index cf10c658dbf4..1e4cdbeb65de 100644 --- a/cpp/src/arrow/flight/sql/odbc/flight_sql/get_info_cache.cc +++ b/cpp/src/arrow/flight/sql/odbc/flight_sql/get_info_cache.cc @@ -204,15 +204,20 @@ inline void SetDefaultIfMissing( namespace driver { namespace flight_sql { using arrow::flight::FlightCallOptions; +using arrow::flight::FlightClientOptions; using arrow::flight::sql::FlightSqlClient; using arrow::flight::sql::SqlInfoOptions; using driver::odbcabstraction::Connection; using driver::odbcabstraction::DriverException; -GetInfoCache::GetInfoCache(FlightCallOptions& call_options, +GetInfoCache::GetInfoCache(FlightClientOptions& client_options, + FlightCallOptions& call_options, std::unique_ptr& client, const std::string& driver_version) - : call_options_(call_options), sql_client_(client), has_server_info_(false) { + : client_options_(client_options), + call_options_(call_options), + sql_client_(client), + has_server_info_(false) { info_[SQL_DRIVER_NAME] = "Arrow Flight ODBC Driver"; info_[SQL_DRIVER_VER] = ConvertToDBMSVer(driver_version); @@ -294,7 +299,8 @@ bool GetInfoCache::LoadInfoFromServer() { arrow::Result> result = sql_client_->GetSqlInfo(call_options_, {}); ThrowIfNotOK(result.status()); - FlightStreamChunkBuffer chunk_iter(*sql_client_, call_options_, result.ValueOrDie()); + FlightStreamChunkBuffer chunk_iter(*sql_client_, client_options_, call_options_, + result.ValueOrDie()); FlightStreamChunk chunk; bool supports_correlation_name = false; diff --git a/cpp/src/arrow/flight/sql/odbc/flight_sql/get_info_cache.h b/cpp/src/arrow/flight/sql/odbc/flight_sql/get_info_cache.h index 819b095e6a63..547fb1cdf286 100644 --- a/cpp/src/arrow/flight/sql/odbc/flight_sql/get_info_cache.h +++ b/cpp/src/arrow/flight/sql/odbc/flight_sql/get_info_cache.h @@ -30,13 +30,15 @@ namespace flight_sql { class GetInfoCache { private: std::unordered_map info_; + arrow::flight::FlightClientOptions& client_options_; arrow::flight::FlightCallOptions& call_options_; std::unique_ptr& sql_client_; std::mutex mutex_; std::atomic has_server_info_; public: - GetInfoCache(arrow::flight::FlightCallOptions& call_options, + GetInfoCache(arrow::flight::FlightClientOptions& client_options, + arrow::flight::FlightCallOptions& call_options, std::unique_ptr& client, const std::string& driver_version); void SetProperty(uint16_t property, driver::odbcabstraction::Connection::Info value); diff --git a/cpp/src/arrow/flight/sql/odbc/odbcabstraction/odbc_impl/odbc_connection.cc b/cpp/src/arrow/flight/sql/odbc/odbcabstraction/odbc_impl/odbc_connection.cc index 15d60ccf1ac3..029477564799 100644 --- a/cpp/src/arrow/flight/sql/odbc/odbcabstraction/odbc_impl/odbc_connection.cc +++ b/cpp/src/arrow/flight/sql/odbc/odbcabstraction/odbc_impl/odbc_connection.cc @@ -52,7 +52,7 @@ namespace { // characters such as semi-colons and equals signs. NOTE: This can be optimized to be // built statically. const boost::xpressive::sregex CONNECTION_STR_REGEX( - boost::xpressive::sregex::compile("([^=;]+)=({.+}|[^=;]+|[^;])")); + boost::xpressive::sregex::compile("([^=;]+)=({.+}|[^;]+|[^;])")); // Load properties from the given DSN. The properties loaded do _not_ overwrite existing // entries in the properties.