From e1c1bb8a5646586ac349fccc8260a2e2498a8b65 Mon Sep 17 00:00:00 2001 From: KanchanShu Date: Thu, 22 Oct 2026 07:46:47 +0000 Subject: [PATCH 1/4] fix(bq_driver): transparently handle transient stream disconnections and resume HTAPI reads --- .../internal/odbc_sql_execute_utils.cc | 214 +++++++++++++----- .../internal/odbc_sql_execute_utils.h | 10 +- .../bq_driver/internal/odbc_stmt_handle.cc | 3 + .../bq_driver/internal/odbc_stmt_handle.h | 24 +- 4 files changed, 196 insertions(+), 55 deletions(-) diff --git a/google/cloud/odbc/bq_driver/internal/odbc_sql_execute_utils.cc b/google/cloud/odbc/bq_driver/internal/odbc_sql_execute_utils.cc index 264f90691a..529492faf6 100644 --- a/google/cloud/odbc/bq_driver/internal/odbc_sql_execute_utils.cc +++ b/google/cloud/odbc/bq_driver/internal/odbc_sql_execute_utils.cc @@ -778,66 +778,171 @@ StatusRecord ProcessRecordBatch( return StatusRecord::Ok(); } +bool IsTransientStreamError(google::cloud::Status const& status) { + using google::cloud::StatusCode; + switch (status.code()) { + case StatusCode::kUnavailable: + case StatusCode::kDeadlineExceeded: + case StatusCode::kAborted: + case StatusCode::kResourceExhausted: + return true; + case StatusCode::kInternal: + case StatusCode::kUnknown: { + std::string message = status.message(); + std::transform(message.begin(), message.end(), message.begin(), + [](unsigned char c) { return std::tolower(c); }); + if (message.find("rst_stream") != std::string::npos || + message.find("stream removed") != std::string::npos || + message.find("connection reset") != std::string::npos || + message.find("broken pipe") != std::string::npos || + message.find("connection closed") != std::string::npos || + message.find("socket closed") != std::string::npos || + message.find("timeout") != std::string::npos || + message.find("timed out") != std::string::npos || + message.find("proxy") != std::string::npos || + message.find("transport: lost connection") != std::string::npos) { + return true; + } + return false; + } + default: + return false; + } + return false; +} + StatusRecord ReadNextResultsFromStream(StatementHandle& stmt_handle) { - std::optional>& optional_stream = - stmt_handle.GetReadRowsStream(); - if (!optional_stream.has_value()) { - return StatusRecord{SQLStates::k_HY000(), - "Internal Error: No HTAPI read stream found!!"}; + ConnectionHandle* conn_handle = stmt_handle.GetConnectionHandle(); + if (!conn_handle) { + return StatusRecord{ + SQLStates::k_HY000(), + "StatementHandle does not have an associated ConnectionHandle"}; } - auto& read_rows_stream = *optional_stream; - auto& optional_it = stmt_handle.GetReadRowsIterator(); - if (!optional_it.has_value()) { - // Initialize iterator to begin() and store it. - auto it = read_rows_stream.begin(); - optional_it = std::move(it); - } else { - // Advance the stored iterator. - // The previous call processed the element at the iterator, - // so we increment it to move to the next element. - ++(*optional_it); - // Irrespective of any errors, if the iterator hasn't reached the end, we - // want to cache it - stmt_handle.SetReadRowsIterator(*optional_it); + std::shared_ptr bq_client = conn_handle->GetClient(); + if (!bq_client) { + return StatusRecord{ + SQLStates::k_HY000(), + "Invalid or null BQ Client within the connection handle"}; + } + + int max_retries = conn_handle->GetDsn().max_retries; + if (max_retries <= 0) { + max_retries = 3; } - auto& it = *optional_it; - if (it != read_rows_stream.end()) { - auto const& read_row_status = *it; - if (!read_row_status) { - return StatusRecord::ConvertFrom(read_row_status.status()); + + google::cloud::ExponentialBackoffPolicy backoff( + std::chrono::milliseconds(100), std::chrono::milliseconds(2000), 2.0); + + for (int attempt = 0; attempt <= max_retries; ++attempt) { + std::optional>& optional_stream = + stmt_handle.GetReadRowsStream(); + if (!optional_stream.has_value()) { + std::string const& stream_name = stmt_handle.GetReadStreamName(); + if (stream_name.empty()) { + return StatusRecord{SQLStates::k_HY000(), + "Internal Error: No HTAPI read stream found!!"}; + } + // Re-attach to the existing ReadSession stream starting at + // stream_row_offset_ + ReadRowsRequest read_rows_request; + read_rows_request.set_read_stream(stream_name); + read_rows_request.set_offset(stmt_handle.GetStreamRowOffset()); + + Options options; + options.set(conn_handle->GetDsn().max_retries); + StreamRange new_stream = + bq_client->GetReadRowsStream(read_rows_request, options); + stmt_handle.SetReadRowsStream(std::move(new_stream)); + stmt_handle.ClearReadRowsIterator(); } - ReadRowsResponse row = *read_row_status; - if (row.has_arrow_record_batch()) { - // The schema is coming from ResultSet cached in the statement handle. - // We don't want to generate the schema again for every batch since it - // will remain the same. - std::shared_ptr schema = stmt_handle.GetArrowSchema(); - StatusRecordOr> record_batch_status = - GetArrowRecordBatch(row.arrow_record_batch(), schema); - if (!record_batch_status) { - return record_batch_status.GetStatusRecord(); + + auto& read_rows_stream = *stmt_handle.GetReadRowsStream(); + auto& optional_it = stmt_handle.GetReadRowsIterator(); + if (!optional_it.has_value()) { + // Initialize iterator to begin() and store it. + auto it = read_rows_stream.begin(); + optional_it = std::move(it); + } else { + // Advance the stored iterator. + // The previous call processed the element at the iterator, + // so we increment it to move to the next element. + ++(*optional_it); + stmt_handle.SetReadRowsIterator(*optional_it); + } + auto& it = *optional_it; + if (it != read_rows_stream.end()) { + auto const& read_row_status = *it; + if (!read_row_status) { + google::cloud::Status status = read_row_status.status(); + if (IsTransientStreamError(status) && attempt < max_retries) { + auto delay = backoff.OnCompletion(); + LOG(WARNING) << "Transient error on HTAPI stream: " << status + << ". Re-attaching to stream '" + << stmt_handle.GetReadStreamName() + << "' from row offset " + << stmt_handle.GetStreamRowOffset() << " (attempt " + << (attempt + 1) << " of " << max_retries + << ") after delay " + << std::chrono::duration_cast( + delay) + .count() + << "ms"; + stmt_handle.ClearReadRowsStream(); + stmt_handle.ClearReadRowsIterator(); + std::this_thread::sleep_for(delay); + continue; + } + return StatusRecord::ConvertFrom(status); + } + ReadRowsResponse row = *read_row_status; + if (row.has_arrow_record_batch()) { + // The schema is coming from ResultSet cached in the statement handle. + // We don't want to generate the schema again for every batch since it + // will remain the same. + std::shared_ptr schema = stmt_handle.GetArrowSchema(); + StatusRecordOr> + record_batch_status = + GetArrowRecordBatch(row.arrow_record_batch(), schema); + if (!record_batch_status) { + return record_batch_status.GetStatusRecord(); + } + // We are reading ResultSet from the stmt_handle because we want to + // preserve the previous state like `num_rows_fetched_yet`. + ResultSet& result_set = stmt_handle.GetResultSet(); + // To have SQLFetch read the rows from the start, we are setting the + // cursor to default. + result_set.cursor = -1; + StatusRecord process_status = + ProcessRecordBatch(schema, *record_batch_status, result_set); + if (!process_status.ok()) { + return process_status; + } + int64_t num_rows = row.row_count() > 0 + ? row.row_count() + : (*record_batch_status)->num_rows(); + stmt_handle.SetStreamRowOffset(stmt_handle.GetStreamRowOffset() + + num_rows); + return StatusRecord::Ok(); + } else { + return StatusRecord{ + SQLStates::k_HY000(), + "Internal Error: cannot find arrow record batch to process!"}; } - // We are reading ResultSet from the stmt_handle because we want to - // preserve the previous state like `num_rows_fetched_yet`. - ResultSet& result_set = stmt_handle.GetResultSet(); - // To have SQLFetch read the rows from the start, we are setting the - // cursor to default. - result_set.cursor = -1; - return ProcessRecordBatch(schema, *record_batch_status, result_set); } else { - return StatusRecord{ - SQLStates::k_HY000(), - "Internal Error: cannot find arrow record batch to process!"}; + stmt_handle.ClearReadRowsStream(); + stmt_handle.ClearReadRowsIterator(); + stmt_handle.ClearReadStreamName(); + stmt_handle.ResetStreamRowOffset(); + // Empty result set. + stmt_handle.GetResultSet().rows.clear(); + stmt_handle.GetResultSet().cursor = -1; + LOG(INFO) << "FetchBQDataReadArrow:: Read stream ended."; + return StatusRecord::Ok(); } - } else { - stmt_handle.ClearReadRowsStream(); - stmt_handle.ClearReadRowsIterator(); - // Empty result set. - stmt_handle.GetResultSet().rows.clear(); - stmt_handle.GetResultSet().cursor = -1; - LOG(INFO) << "FetchBQDataReadArrow:: Read stream ended."; } - return StatusRecord::Ok(); + + return StatusRecord{SQLStates::k_HY000(), + "Failed to read from HTAPI stream: retries exhausted."}; } StatusRecord FetchBQDataReadArrow(StatementHandle& stmt_handle, @@ -881,9 +986,14 @@ StatusRecord FetchBQDataReadArrow(StatementHandle& stmt_handle, std::shared_ptr schema = *schema_status; stmt_handle.SetArrowSchema(schema); + stmt_handle.SetReadStreamName(read_stream_name); + stmt_handle.SetStreamRowOffset(0); + stmt_handle.SetWasHtapiEnabled(true); + // Create a ReadRowsRequest. ReadRowsRequest read_rows_request; read_rows_request.set_read_stream(read_stream_name); + read_rows_request.set_offset(0); // Before we call ReadNextResultsFromStream, we are caching the stream StreamRange diff --git a/google/cloud/odbc/bq_driver/internal/odbc_sql_execute_utils.h b/google/cloud/odbc/bq_driver/internal/odbc_sql_execute_utils.h index e0526e1747..806d197ff5 100644 --- a/google/cloud/odbc/bq_driver/internal/odbc_sql_execute_utils.h +++ b/google/cloud/odbc/bq_driver/internal/odbc_sql_execute_utils.h @@ -72,9 +72,17 @@ odbc_internal::StatusRecord PopulateScriptChildJobs( post_query_request); #if (!defined(_WIN32) || defined(_WIN64)) && !defined(NO_ARROW) +/* + * @brief Identifies whether a stream error is transient (e.g., network + * interruptions, proxy timeouts, RST_STREAM, connection resets). + */ +bool IsTransientStreamError(google::cloud::Status const& status); + /* * @brief Reads the next set of rows from the stream cached in the statement - * handle + * handle, reconnecting/re-attaching to the existing ReadSession and resuming + * from the last successfully consumed row offset if transient disconnections + * occur. */ StatusRecord ReadNextResultsFromStream(StatementHandle& stmt_handle); diff --git a/google/cloud/odbc/bq_driver/internal/odbc_stmt_handle.cc b/google/cloud/odbc/bq_driver/internal/odbc_stmt_handle.cc index b71a0e2535..52ac0902d1 100644 --- a/google/cloud/odbc/bq_driver/internal/odbc_stmt_handle.cc +++ b/google/cloud/odbc/bq_driver/internal/odbc_stmt_handle.cc @@ -570,6 +570,9 @@ void StatementHandle::CloseCursor() { #if (!defined(_WIN32) || defined(_WIN64)) && !defined(NO_ARROW) ClearReadRowsStream(); ClearReadRowsIterator(); + ClearReadStreamName(); + ResetStreamRowOffset(); + SetWasHtapiEnabled(false); #endif if (StatementPrepared()) { SetStmtState(StmtStates::kStatementPrepared); diff --git a/google/cloud/odbc/bq_driver/internal/odbc_stmt_handle.h b/google/cloud/odbc/bq_driver/internal/odbc_stmt_handle.h index 69c5df367d..ba57e571ec 100644 --- a/google/cloud/odbc/bq_driver/internal/odbc_stmt_handle.h +++ b/google/cloud/odbc/bq_driver/internal/odbc_stmt_handle.h @@ -134,8 +134,25 @@ class StatementHandle : public Handle { // Clearer for the iterator inline void ClearReadRowsIterator() { read_rows_iterator_.reset(); } - // Returns true if read_rows_stream_ is NOT null (i.e., has a value). - inline bool WasHtapiEnabled() const { return read_rows_stream_.has_value(); } + // Returns true if HTAPI was enabled for this query. + inline bool WasHtapiEnabled() const { + return was_htapi_enabled_ || read_rows_stream_.has_value(); + } + inline void SetWasHtapiEnabled(bool enabled) { was_htapi_enabled_ = enabled; } + + inline std::string const& GetReadStreamName() const { + return read_stream_name_; + } + inline void SetReadStreamName(std::string name) { + read_stream_name_ = std::move(name); + } + inline void ClearReadStreamName() { read_stream_name_.clear(); } + + inline std::int64_t GetStreamRowOffset() const { return stream_row_offset_; } + inline void SetStreamRowOffset(std::int64_t offset) { + stream_row_offset_ = offset; + } + inline void ResetStreamRowOffset() { stream_row_offset_ = 0; } void SetReadRowsStream( StreamRange<::google::cloud::bigquery::storage::v1::ReadRowsResponse> @@ -346,6 +363,9 @@ class StatementHandle : public Handle { std::optional< StreamRange<::google::cloud::bigquery::storage::v1::ReadRowsResponse>> read_rows_stream_; + std::string read_stream_name_; + std::int64_t stream_row_offset_{0}; + bool was_htapi_enabled_{false}; mutable std::mutex statement_handle_mutex_; ::google::cloud::bigquery_v2_minimal_internal::PostQueryRequest post_query_request_; From 43b914be20c2eef2afe21eedeea5b5a1fe628bd6 Mon Sep 17 00:00:00 2001 From: KanchanShu Date: Fri, 23 Oct 2026 13:09:26 +0000 Subject: [PATCH 2/4] addded test case and fixed review comments --- .../internal/odbc_sql_execute_utils.cc | 37 ++++++++--- .../odbc_driver_tests/statement_test.cc | 61 +++++++++++++++++++ 2 files changed, 89 insertions(+), 9 deletions(-) diff --git a/google/cloud/odbc/bq_driver/internal/odbc_sql_execute_utils.cc b/google/cloud/odbc/bq_driver/internal/odbc_sql_execute_utils.cc index 529492faf6..6265d3c350 100644 --- a/google/cloud/odbc/bq_driver/internal/odbc_sql_execute_utils.cc +++ b/google/cloud/odbc/bq_driver/internal/odbc_sql_execute_utils.cc @@ -797,9 +797,6 @@ bool IsTransientStreamError(google::cloud::Status const& status) { message.find("broken pipe") != std::string::npos || message.find("connection closed") != std::string::npos || message.find("socket closed") != std::string::npos || - message.find("timeout") != std::string::npos || - message.find("timed out") != std::string::npos || - message.find("proxy") != std::string::npos || message.find("transport: lost connection") != std::string::npos) { return true; } @@ -825,10 +822,7 @@ StatusRecord ReadNextResultsFromStream(StatementHandle& stmt_handle) { "Invalid or null BQ Client within the connection handle"}; } - int max_retries = conn_handle->GetDsn().max_retries; - if (max_retries <= 0) { - max_retries = 3; - } + int max_retries = static_cast(conn_handle->GetDsn().max_retries); google::cloud::ExponentialBackoffPolicy backoff( std::chrono::milliseconds(100), std::chrono::milliseconds(2000), 2.0); @@ -871,9 +865,34 @@ StatusRecord ReadNextResultsFromStream(StatementHandle& stmt_handle) { } auto& it = *optional_it; if (it != read_rows_stream.end()) { + char const* inject_env = + std::getenv("BIGQUERY_ODBC_TEST_INJECT_STREAM_ERROR"); + bool has_injected_error = false; + google::cloud::Status injected_status; + if (inject_env && std::string_view(inject_env) == "1" && + stmt_handle.GetStreamRowOffset() > 0) { + injected_status = google::cloud::Status( + google::cloud::StatusCode::kUnavailable, + "transport: stream received RST_STREAM with error code: 8"); + std::string record = + "RESUMED_AT_OFFSET=" + + std::to_string(stmt_handle.GetStreamRowOffset()) + ";ERROR=" + + google::cloud::StatusCodeToString(injected_status.code()) + + ": RST_STREAM"; +#ifdef _WIN32 + _putenv_s("BIGQUERY_ODBC_TEST_INJECT_STREAM_ERROR", record.c_str()); +#else + setenv("BIGQUERY_ODBC_TEST_INJECT_STREAM_ERROR", record.c_str(), 1); +#endif + LOG(WARNING) << "[TEST HOOK] Injected transient stream error " + "(RST_STREAM) at row offset " + << stmt_handle.GetStreamRowOffset(); + has_injected_error = true; + } auto const& read_row_status = *it; - if (!read_row_status) { - google::cloud::Status status = read_row_status.status(); + if (has_injected_error || !read_row_status) { + google::cloud::Status status = + has_injected_error ? injected_status : read_row_status.status(); if (IsTransientStreamError(status) && attempt < max_retries) { auto delay = backoff.OnCompletion(); LOG(WARNING) << "Transient error on HTAPI stream: " << status diff --git a/google/cloud/odbc/integration_tests/odbc_driver_tests/statement_test.cc b/google/cloud/odbc/integration_tests/odbc_driver_tests/statement_test.cc index c5d8385725..c829a2f902 100644 --- a/google/cloud/odbc/integration_tests/odbc_driver_tests/statement_test.cc +++ b/google/cloud/odbc/integration_tests/odbc_driver_tests/statement_test.cc @@ -829,6 +829,67 @@ TEST_P(HTAPIParameterizedTest, SQLExecDirect_pagination_with_row_array_size) { EXPECT_EQ(Disconnect(conn), SQL_SUCCESS); } +#ifdef BQ_DRIVER_INTEGRATION_TESTS +TEST(StatementTest, SQLExecDirect_HTAPI_StreamRetryAndResume) { +#ifdef _WIN32 + _putenv_s("BIGQUERY_ODBC_TEST_INJECT_STREAM_ERROR", "1"); +#else + setenv("BIGQUERY_ODBC_TEST_INJECT_STREAM_ERROR", "1", 1); +#endif + + auto conn = std::make_shared(); + std::string connection_string = + kDefaultConnectionString + + ";AllowHtapiForLargeResults=1;UseDefaultLargeResultsDataset=0;" + "LargeResultsDataSetId=_bqodbc_temp_tables_euwest1"; + EXPECT_EQ(Connect(connection_string, conn), SQL_SUCCESS); + + int const limit = 3000; + std::string query = + "SELECT * EXCEPT (index) FROM " + "ODBC_HTAPI_TESTING_EUROPE_WEST1.300_columns_string " + "ORDER BY index LIMIT " + + std::to_string(limit) + ";"; + + Table table("Random_table_name"); + RowWiseResults const& results = table.Fetch(conn, query); + + // Assert proof of issue reproduction, error caught, and reconnection: + // The driver captures the injected error and records the resumption offset. + char const* test_status = + std::getenv("BIGQUERY_ODBC_TEST_INJECT_STREAM_ERROR"); + ASSERT_NE(test_status, nullptr) + << "Verification failed: Transient stream error was never injected or " + "reconnection never fired!"; + EXPECT_THAT(test_status, HasSubstr("RESUMED_AT_OFFSET=")); + EXPECT_THAT(test_status, HasSubstr("ERROR=UNAVAILABLE: RST_STREAM")); + +#ifdef _WIN32 + _putenv_s("BIGQUERY_ODBC_TEST_INJECT_STREAM_ERROR", ""); +#else + unsetenv("BIGQUERY_ODBC_TEST_INJECT_STREAM_ERROR"); +#endif + + int const expected_num_cols = 300; + ASSERT_EQ(results.size(), limit) << "Row count mismatch."; + for (int i = 0; i < limit; ++i) { + Row const& row = results[i]; + ASSERT_EQ(row.size(), expected_num_cols) + << "Row " << i << ": Column count mismatch."; + for (int j = 0; j < expected_num_cols; ++j) { + std::string expected_value = + "col" + std::to_string(j) + "_row" + std::to_string(i); + ASSERT_TRUE(row.count(j)) + << "Row " << i << ": Missing expected column with index " << j; + ASSERT_EQ(row.at(j), expected_value) + << "Row " << i << ", Col " << j << ": Value mismatch."; + } + } + + EXPECT_EQ(Disconnect(conn), SQL_SUCCESS); +} +#endif // BQ_DRIVER_INTEGRATION_TESTS + TEST_P(HTAPIParameterizedTest, SQLExecDirect_with_empty_result_set) { bool is_htapi = GetParam(); auto conn = std::make_shared(); From 83ff965e3afba3f40a25de279c3fa32e42c63199 Mon Sep 17 00:00:00 2001 From: Sachin Purohit Date: Fri, 30 Oct 2026 19:22:00 +0000 Subject: [PATCH 3/4] test(resilience): move SQLExecDirect_HTAPI_StreamRetryAndResume to resilience_test.cc --- .../odbc_driver_tests/resilience_test.cc | 70 +++++++++++++++++++ .../odbc_driver_tests/statement_test.cc | 61 ---------------- 2 files changed, 70 insertions(+), 61 deletions(-) diff --git a/google/cloud/odbc/integration_tests/odbc_driver_tests/resilience_test.cc b/google/cloud/odbc/integration_tests/odbc_driver_tests/resilience_test.cc index 3598c43f68..6aeef40873 100644 --- a/google/cloud/odbc/integration_tests/odbc_driver_tests/resilience_test.cc +++ b/google/cloud/odbc/integration_tests/odbc_driver_tests/resilience_test.cc @@ -14,6 +14,7 @@ #include "google/cloud/odbc/testing/odbc_utils/connection.h" #include "google/cloud/odbc/testing/odbc_utils/statement.h" +#include #include #include #include @@ -25,6 +26,8 @@ #include #include +using ::testing::HasSubstr; + namespace google::cloud::odbc_tests { #ifdef BQ_DRIVER_INTEGRATION_TESTS @@ -291,6 +294,73 @@ TEST(ResilienceTest, ReadAPI_IdleProxy_ConnectionDroppedBetweenQueries) { EXPECT_EQ(Disconnect(conn), SQL_SUCCESS); } +struct ScopedEnvVar { + char const* name; + ScopedEnvVar(char const* env_name, char const* value) : name(env_name) { +#ifdef _WIN32 + _putenv_s(name, value); +#else + setenv(name, value, 1); +#endif + } + ~ScopedEnvVar() { +#ifdef _WIN32 + _putenv_s(name, ""); +#else + unsetenv(name); +#endif + } +}; + +TEST(StatementTest, SQLExecDirect_HTAPI_StreamRetryAndResume) { + ScopedEnvVar env_guard("BIGQUERY_ODBC_TEST_INJECT_STREAM_ERROR", "1"); + + auto conn = std::make_shared(); + std::string connection_string = + kDefaultConnectionString + + ";AllowHtapiForLargeResults=1;UseDefaultLargeResultsDataset=0;" + "LargeResultsDataSetId=_bqodbc_temp_tables_euwest1"; + EXPECT_EQ(Connect(connection_string, conn), SQL_SUCCESS); + + int const limit = 3000; + std::string query = + "SELECT * EXCEPT (index) FROM " + "ODBC_HTAPI_TESTING_EUROPE_WEST1.300_columns_string " + "ORDER BY index LIMIT " + + std::to_string(limit) + ";"; + + Table table("Random_table_name"); + RowWiseResults const& results = table.Fetch(conn, query); + + // Assert proof of issue reproduction, error caught, and reconnection: + // The driver captures the injected error and records the resumption offset. + char const* test_status = + std::getenv("BIGQUERY_ODBC_TEST_INJECT_STREAM_ERROR"); + ASSERT_NE(test_status, nullptr) + << "Verification failed: Transient stream error was never injected or " + "reconnection never fired!"; + EXPECT_THAT(test_status, HasSubstr("RESUMED_AT_OFFSET=")); + EXPECT_THAT(test_status, HasSubstr("ERROR=UNAVAILABLE: RST_STREAM")); + + int const expected_num_cols = 300; + ASSERT_EQ(results.size(), limit) << "Row count mismatch."; + for (int i = 0; i < limit; ++i) { + Row const& row = results[i]; + ASSERT_EQ(row.size(), expected_num_cols) + << "Row " << i << ": Column count mismatch."; + for (int j = 0; j < expected_num_cols; ++j) { + std::string expected_value = + "col" + std::to_string(j) + "_row" + std::to_string(i); + ASSERT_TRUE(row.count(j)) + << "Row " << i << ": Missing expected column with index " << j; + ASSERT_EQ(row.at(j), expected_value) + << "Row " << i << ", Col " << j << ": Value mismatch."; + } + } + + EXPECT_EQ(Disconnect(conn), SQL_SUCCESS); +} + #endif // BQ_DRIVER_INTEGRATION_TESTS } // namespace google::cloud::odbc_tests diff --git a/google/cloud/odbc/integration_tests/odbc_driver_tests/statement_test.cc b/google/cloud/odbc/integration_tests/odbc_driver_tests/statement_test.cc index c829a2f902..c5d8385725 100644 --- a/google/cloud/odbc/integration_tests/odbc_driver_tests/statement_test.cc +++ b/google/cloud/odbc/integration_tests/odbc_driver_tests/statement_test.cc @@ -829,67 +829,6 @@ TEST_P(HTAPIParameterizedTest, SQLExecDirect_pagination_with_row_array_size) { EXPECT_EQ(Disconnect(conn), SQL_SUCCESS); } -#ifdef BQ_DRIVER_INTEGRATION_TESTS -TEST(StatementTest, SQLExecDirect_HTAPI_StreamRetryAndResume) { -#ifdef _WIN32 - _putenv_s("BIGQUERY_ODBC_TEST_INJECT_STREAM_ERROR", "1"); -#else - setenv("BIGQUERY_ODBC_TEST_INJECT_STREAM_ERROR", "1", 1); -#endif - - auto conn = std::make_shared(); - std::string connection_string = - kDefaultConnectionString + - ";AllowHtapiForLargeResults=1;UseDefaultLargeResultsDataset=0;" - "LargeResultsDataSetId=_bqodbc_temp_tables_euwest1"; - EXPECT_EQ(Connect(connection_string, conn), SQL_SUCCESS); - - int const limit = 3000; - std::string query = - "SELECT * EXCEPT (index) FROM " - "ODBC_HTAPI_TESTING_EUROPE_WEST1.300_columns_string " - "ORDER BY index LIMIT " + - std::to_string(limit) + ";"; - - Table table("Random_table_name"); - RowWiseResults const& results = table.Fetch(conn, query); - - // Assert proof of issue reproduction, error caught, and reconnection: - // The driver captures the injected error and records the resumption offset. - char const* test_status = - std::getenv("BIGQUERY_ODBC_TEST_INJECT_STREAM_ERROR"); - ASSERT_NE(test_status, nullptr) - << "Verification failed: Transient stream error was never injected or " - "reconnection never fired!"; - EXPECT_THAT(test_status, HasSubstr("RESUMED_AT_OFFSET=")); - EXPECT_THAT(test_status, HasSubstr("ERROR=UNAVAILABLE: RST_STREAM")); - -#ifdef _WIN32 - _putenv_s("BIGQUERY_ODBC_TEST_INJECT_STREAM_ERROR", ""); -#else - unsetenv("BIGQUERY_ODBC_TEST_INJECT_STREAM_ERROR"); -#endif - - int const expected_num_cols = 300; - ASSERT_EQ(results.size(), limit) << "Row count mismatch."; - for (int i = 0; i < limit; ++i) { - Row const& row = results[i]; - ASSERT_EQ(row.size(), expected_num_cols) - << "Row " << i << ": Column count mismatch."; - for (int j = 0; j < expected_num_cols; ++j) { - std::string expected_value = - "col" + std::to_string(j) + "_row" + std::to_string(i); - ASSERT_TRUE(row.count(j)) - << "Row " << i << ": Missing expected column with index " << j; - ASSERT_EQ(row.at(j), expected_value) - << "Row " << i << ", Col " << j << ": Value mismatch."; - } - } - - EXPECT_EQ(Disconnect(conn), SQL_SUCCESS); -} -#endif // BQ_DRIVER_INTEGRATION_TESTS - TEST_P(HTAPIParameterizedTest, SQLExecDirect_with_empty_result_set) { bool is_htapi = GetParam(); auto conn = std::make_shared(); From d0f0f81a1f9e98eed2452cd18e17b36f99121054 Mon Sep 17 00:00:00 2001 From: Sachin Purohit Date: Fri, 30 Oct 2026 19:22:23 +0000 Subject: [PATCH 4/4] fix(bq_driver): use custom BigQueryReadRetryPolicy for HTAPI stream retry --- .../integration-production-bq-driver-dm.sh | 2 +- ci/cloudbuild/builds/resilience-tests.sh | 25 ++ .../internal/odbc_sql_execute_utils.cc | 215 ++++++------------ .../internal/odbc_sql_execute_utils.h | 50 +++- .../internal/odbc_sql_execute_utils_test.cc | 58 +++++ .../bq_driver/internal/odbc_stmt_handle.cc | 3 - .../bq_driver/internal/odbc_stmt_handle.h | 24 +- .../odbc_driver_tests/resilience_test.cc | 56 ++--- .../idle_connect_proxy/idle_connect_proxy.py | 63 +++-- 9 files changed, 270 insertions(+), 226 deletions(-) diff --git a/ci/cloudbuild/builds/integration-production-bq-driver-dm.sh b/ci/cloudbuild/builds/integration-production-bq-driver-dm.sh index 605d823eee..80b6d5edf7 100755 --- a/ci/cloudbuild/builds/integration-production-bq-driver-dm.sh +++ b/ci/cloudbuild/builds/integration-production-bq-driver-dm.sh @@ -128,7 +128,7 @@ fi # Copy the roots.pem file to the .so directory to run test cases. cp /opt/odbc-driver/roots.pem "cmake-out/google/cloud/odbc/roots.pem" mapfile -t ctest_args < <(ctest::common_args) -io::run env -C cmake-out ctest "${ctest_args[@]}" -E "IdleProxy" +io::run env -C cmake-out ctest "${ctest_args[@]}" -E "IdleProxy|SQLExecDirect_HTAPI_StreamRetryAndResume" if [[ "${UNIXODBC_INSTALLED}" == "true" ]]; then "${WORKSPACE_DIR}/ci/cloudbuild/builds/resilience-tests.sh" diff --git a/ci/cloudbuild/builds/resilience-tests.sh b/ci/cloudbuild/builds/resilience-tests.sh index e9b89693b7..6a9766ee34 100755 --- a/ci/cloudbuild/builds/resilience-tests.sh +++ b/ci/cloudbuild/builds/resilience-tests.sh @@ -79,3 +79,28 @@ if [[ -f "${PROXY_LOG}" ]] && grep -q "PING_FLOOD" "${PROXY_LOG}"; then echo "ERROR: Excessive keepalive ping flood detected in proxy log!" exit 1 fi + +kill "${PROXY_PID}" 2>/dev/null || true +wait "${PROXY_PID}" 2>/dev/null || true +PROXY_PID="" + +echo "Starting idle connect proxy with --kill-after-bytes for stream retry test..." +python3 "${PROJECT_ROOT}/tools/idle_connect_proxy/idle_connect_proxy.py" \ + --port "${PROXY_PORT}" \ + --idle-timeout 0 \ + --kill-after-bytes 4000000 \ + --reconnect-kill-after-bytes 100000 \ + --max-kills 2 \ + --mode corrupt \ + --kills-file "${PROXY_KILLS}" \ + --stats-file "${PROXY_STATS}" \ + >>"${PROXY_LOG}" 2>&1 & +PROXY_PID=$! +sleep 1 +if ! kill -0 "${PROXY_PID}" 2>/dev/null; then + echo "ERROR: Idle connect proxy failed to start for stream retry test." + exit 1 +fi + +echo "Running HTAPI stream retry and resume resilience test..." +io::run env -C cmake-out ctest --output-on-failure --progress -j 1 -R "SQLExecDirect_HTAPI_StreamRetryAndResume" diff --git a/google/cloud/odbc/bq_driver/internal/odbc_sql_execute_utils.cc b/google/cloud/odbc/bq_driver/internal/odbc_sql_execute_utils.cc index 6265d3c350..bf1fb23df4 100644 --- a/google/cloud/odbc/bq_driver/internal/odbc_sql_execute_utils.cc +++ b/google/cloud/odbc/bq_driver/internal/odbc_sql_execute_utils.cc @@ -20,6 +20,7 @@ #include "absl/time/civil_time.h" #include "absl/time/time.h" #if (!defined(_WIN32) || defined(_WIN64)) && !defined(NO_ARROW) +#include "google/cloud/bigquery/storage/v1/bigquery_read_options.h" #include #endif #include @@ -797,6 +798,8 @@ bool IsTransientStreamError(google::cloud::Status const& status) { message.find("broken pipe") != std::string::npos || message.find("connection closed") != std::string::npos || message.find("socket closed") != std::string::npos || + message.find("transport closed") != std::string::npos || + message.find("endpoint read failed") != std::string::npos || message.find("transport: lost connection") != std::string::npos) { return true; } @@ -809,159 +812,65 @@ bool IsTransientStreamError(google::cloud::Status const& status) { } StatusRecord ReadNextResultsFromStream(StatementHandle& stmt_handle) { - ConnectionHandle* conn_handle = stmt_handle.GetConnectionHandle(); - if (!conn_handle) { - return StatusRecord{ - SQLStates::k_HY000(), - "StatementHandle does not have an associated ConnectionHandle"}; + std::optional>& optional_stream = + stmt_handle.GetReadRowsStream(); + if (!optional_stream.has_value()) { + return StatusRecord{SQLStates::k_HY000(), + "Internal Error: No HTAPI read stream found!!"}; } - std::shared_ptr bq_client = conn_handle->GetClient(); - if (!bq_client) { - return StatusRecord{ - SQLStates::k_HY000(), - "Invalid or null BQ Client within the connection handle"}; + auto& read_rows_stream = *optional_stream; + auto& optional_it = stmt_handle.GetReadRowsIterator(); + if (!optional_it.has_value()) { + // Initialize iterator to begin() and store it. + auto it = read_rows_stream.begin(); + optional_it = std::move(it); + } else { + // Advance the stored iterator. + // The previous call processed the element at the iterator, + // so we increment it to move to the next element. + ++(*optional_it); + // Irrespective of any errors, if the iterator hasn't reached the end, we + // want to cache it + stmt_handle.SetReadRowsIterator(*optional_it); } - - int max_retries = static_cast(conn_handle->GetDsn().max_retries); - - google::cloud::ExponentialBackoffPolicy backoff( - std::chrono::milliseconds(100), std::chrono::milliseconds(2000), 2.0); - - for (int attempt = 0; attempt <= max_retries; ++attempt) { - std::optional>& optional_stream = - stmt_handle.GetReadRowsStream(); - if (!optional_stream.has_value()) { - std::string const& stream_name = stmt_handle.GetReadStreamName(); - if (stream_name.empty()) { - return StatusRecord{SQLStates::k_HY000(), - "Internal Error: No HTAPI read stream found!!"}; - } - // Re-attach to the existing ReadSession stream starting at - // stream_row_offset_ - ReadRowsRequest read_rows_request; - read_rows_request.set_read_stream(stream_name); - read_rows_request.set_offset(stmt_handle.GetStreamRowOffset()); - - Options options; - options.set(conn_handle->GetDsn().max_retries); - StreamRange new_stream = - bq_client->GetReadRowsStream(read_rows_request, options); - stmt_handle.SetReadRowsStream(std::move(new_stream)); - stmt_handle.ClearReadRowsIterator(); + auto& it = *optional_it; + if (it != read_rows_stream.end()) { + auto const& read_row_status = *it; + if (!read_row_status) { + return StatusRecord::ConvertFrom(read_row_status.status()); } - - auto& read_rows_stream = *stmt_handle.GetReadRowsStream(); - auto& optional_it = stmt_handle.GetReadRowsIterator(); - if (!optional_it.has_value()) { - // Initialize iterator to begin() and store it. - auto it = read_rows_stream.begin(); - optional_it = std::move(it); - } else { - // Advance the stored iterator. - // The previous call processed the element at the iterator, - // so we increment it to move to the next element. - ++(*optional_it); - stmt_handle.SetReadRowsIterator(*optional_it); - } - auto& it = *optional_it; - if (it != read_rows_stream.end()) { - char const* inject_env = - std::getenv("BIGQUERY_ODBC_TEST_INJECT_STREAM_ERROR"); - bool has_injected_error = false; - google::cloud::Status injected_status; - if (inject_env && std::string_view(inject_env) == "1" && - stmt_handle.GetStreamRowOffset() > 0) { - injected_status = google::cloud::Status( - google::cloud::StatusCode::kUnavailable, - "transport: stream received RST_STREAM with error code: 8"); - std::string record = - "RESUMED_AT_OFFSET=" + - std::to_string(stmt_handle.GetStreamRowOffset()) + ";ERROR=" + - google::cloud::StatusCodeToString(injected_status.code()) + - ": RST_STREAM"; -#ifdef _WIN32 - _putenv_s("BIGQUERY_ODBC_TEST_INJECT_STREAM_ERROR", record.c_str()); -#else - setenv("BIGQUERY_ODBC_TEST_INJECT_STREAM_ERROR", record.c_str(), 1); -#endif - LOG(WARNING) << "[TEST HOOK] Injected transient stream error " - "(RST_STREAM) at row offset " - << stmt_handle.GetStreamRowOffset(); - has_injected_error = true; - } - auto const& read_row_status = *it; - if (has_injected_error || !read_row_status) { - google::cloud::Status status = - has_injected_error ? injected_status : read_row_status.status(); - if (IsTransientStreamError(status) && attempt < max_retries) { - auto delay = backoff.OnCompletion(); - LOG(WARNING) << "Transient error on HTAPI stream: " << status - << ". Re-attaching to stream '" - << stmt_handle.GetReadStreamName() - << "' from row offset " - << stmt_handle.GetStreamRowOffset() << " (attempt " - << (attempt + 1) << " of " << max_retries - << ") after delay " - << std::chrono::duration_cast( - delay) - .count() - << "ms"; - stmt_handle.ClearReadRowsStream(); - stmt_handle.ClearReadRowsIterator(); - std::this_thread::sleep_for(delay); - continue; - } - return StatusRecord::ConvertFrom(status); - } - ReadRowsResponse row = *read_row_status; - if (row.has_arrow_record_batch()) { - // The schema is coming from ResultSet cached in the statement handle. - // We don't want to generate the schema again for every batch since it - // will remain the same. - std::shared_ptr schema = stmt_handle.GetArrowSchema(); - StatusRecordOr> - record_batch_status = - GetArrowRecordBatch(row.arrow_record_batch(), schema); - if (!record_batch_status) { - return record_batch_status.GetStatusRecord(); - } - // We are reading ResultSet from the stmt_handle because we want to - // preserve the previous state like `num_rows_fetched_yet`. - ResultSet& result_set = stmt_handle.GetResultSet(); - // To have SQLFetch read the rows from the start, we are setting the - // cursor to default. - result_set.cursor = -1; - StatusRecord process_status = - ProcessRecordBatch(schema, *record_batch_status, result_set); - if (!process_status.ok()) { - return process_status; - } - int64_t num_rows = row.row_count() > 0 - ? row.row_count() - : (*record_batch_status)->num_rows(); - stmt_handle.SetStreamRowOffset(stmt_handle.GetStreamRowOffset() + - num_rows); - return StatusRecord::Ok(); - } else { - return StatusRecord{ - SQLStates::k_HY000(), - "Internal Error: cannot find arrow record batch to process!"}; + ReadRowsResponse row = *read_row_status; + if (row.has_arrow_record_batch()) { + // The schema is coming from ResultSet cached in the statement handle. + // We don't want to generate the schema again for every batch since it + // will remain the same. + std::shared_ptr schema = stmt_handle.GetArrowSchema(); + StatusRecordOr> record_batch_status = + GetArrowRecordBatch(row.arrow_record_batch(), schema); + if (!record_batch_status) { + return record_batch_status.GetStatusRecord(); } + // We are reading ResultSet from the stmt_handle because we want to + // preserve the previous state like `num_rows_fetched_yet`. + ResultSet& result_set = stmt_handle.GetResultSet(); + // To have SQLFetch read the rows from the start, we are setting the + // cursor to default. + result_set.cursor = -1; + return ProcessRecordBatch(schema, *record_batch_status, result_set); } else { - stmt_handle.ClearReadRowsStream(); - stmt_handle.ClearReadRowsIterator(); - stmt_handle.ClearReadStreamName(); - stmt_handle.ResetStreamRowOffset(); - // Empty result set. - stmt_handle.GetResultSet().rows.clear(); - stmt_handle.GetResultSet().cursor = -1; - LOG(INFO) << "FetchBQDataReadArrow:: Read stream ended."; - return StatusRecord::Ok(); + return StatusRecord{ + SQLStates::k_HY000(), + "Internal Error: cannot find arrow record batch to process!"}; } + } else { + stmt_handle.ClearReadRowsStream(); + stmt_handle.ClearReadRowsIterator(); + // Empty result set. + stmt_handle.GetResultSet().rows.clear(); + stmt_handle.GetResultSet().cursor = -1; + LOG(INFO) << "FetchBQDataReadArrow:: Read stream ended."; } - - return StatusRecord{SQLStates::k_HY000(), - "Failed to read from HTAPI stream: retries exhausted."}; + return StatusRecord::Ok(); } StatusRecord FetchBQDataReadArrow(StatementHandle& stmt_handle, @@ -980,8 +889,17 @@ StatusRecord FetchBQDataReadArrow(StatementHandle& stmt_handle, read_session->set_data_format(ARROW); ConnectionHandle& conn_handle = *(stmt_handle.GetConnectionHandle()); + int const max_retries = static_cast(conn_handle.GetDsn().max_retries); Options options; options.set(conn_handle.GetDsn().max_retries); + options + .set<::google::cloud::bigquery_storage_v1::BigQueryReadRetryPolicyOption>( + std::make_shared(max_retries)); + options.set< + ::google::cloud::bigquery_storage_v1::BigQueryReadBackoffPolicyOption>( + google::cloud::ExponentialBackoffPolicy( + std::chrono::milliseconds(100), std::chrono::milliseconds(2000), 2.0) + .clone()); auto bq_client = stmt_handle.GetConnectionHandle()->GetClient(); auto read_session_status = bq_client->CreateReadSession(create_read_session_request, options); @@ -1005,14 +923,9 @@ StatusRecord FetchBQDataReadArrow(StatementHandle& stmt_handle, std::shared_ptr schema = *schema_status; stmt_handle.SetArrowSchema(schema); - stmt_handle.SetReadStreamName(read_stream_name); - stmt_handle.SetStreamRowOffset(0); - stmt_handle.SetWasHtapiEnabled(true); - // Create a ReadRowsRequest. ReadRowsRequest read_rows_request; read_rows_request.set_read_stream(read_stream_name); - read_rows_request.set_offset(0); // Before we call ReadNextResultsFromStream, we are caching the stream StreamRange diff --git a/google/cloud/odbc/bq_driver/internal/odbc_sql_execute_utils.h b/google/cloud/odbc/bq_driver/internal/odbc_sql_execute_utils.h index 806d197ff5..d0f175044a 100644 --- a/google/cloud/odbc/bq_driver/internal/odbc_sql_execute_utils.h +++ b/google/cloud/odbc/bq_driver/internal/odbc_sql_execute_utils.h @@ -18,6 +18,9 @@ #include "google/cloud/odbc/bq_driver/internal/data_translation_inv.h" #include "google/cloud/odbc/bq_driver/internal/odbc_desc_handle.h" #include "google/cloud/odbc/bq_driver/internal/odbc_stmt_handle.h" +#if (!defined(_WIN32) || defined(_WIN64)) && !defined(NO_ARROW) +#include "google/cloud/bigquery/storage/v1/bigquery_read_options.h" +#endif namespace google::cloud::odbc_bq_driver_internal { @@ -78,11 +81,52 @@ odbc_internal::StatusRecord PopulateScriptChildJobs( */ bool IsTransientStreamError(google::cloud::Status const& status); +/* + * @brief Custom retry policy for BigQuery Storage Read API (HTAPI) streams + * that classifies transient HTTP/2 and proxy disconnects (such as RST_STREAM, + * Stream removed, and socket/connection resets) as retryable so + * google-cloud-cpp's ResumableStreamingReadRpc automatically reconnects and + * resumes from the last read row offset. + */ +class HtapiStreamRetryPolicy + : public ::google::cloud::bigquery_storage_v1::BigQueryReadRetryPolicy { + public: + explicit HtapiStreamRetryPolicy(int maximum_failures) + : maximum_failures_(maximum_failures < 0 ? 0 : maximum_failures) {} + + bool OnFailure(google::cloud::Status const& status) override { + if (status.ok()) { + return true; + } + if (!IsTransientStreamError(status)) { + failure_count_ = maximum_failures_ + 1; + return false; + } + ++failure_count_; + return failure_count_ <= maximum_failures_; + } + + bool IsExhausted() const override { + return failure_count_ > maximum_failures_; + } + + bool IsPermanentFailure(google::cloud::Status const& status) const override { + return !status.ok() && !IsTransientStreamError(status); + } + + std::unique_ptr<::google::cloud::bigquery_storage_v1::BigQueryReadRetryPolicy> + clone() const override { + return std::make_unique(maximum_failures_); + } + + private: + int failure_count_ = 0; + int maximum_failures_; +}; + /* * @brief Reads the next set of rows from the stream cached in the statement - * handle, reconnecting/re-attaching to the existing ReadSession and resuming - * from the last successfully consumed row offset if transient disconnections - * occur. + * handle */ StatusRecord ReadNextResultsFromStream(StatementHandle& stmt_handle); diff --git a/google/cloud/odbc/bq_driver/internal/odbc_sql_execute_utils_test.cc b/google/cloud/odbc/bq_driver/internal/odbc_sql_execute_utils_test.cc index 198c5c3e84..6493174da0 100644 --- a/google/cloud/odbc/bq_driver/internal/odbc_sql_execute_utils_test.cc +++ b/google/cloud/odbc/bq_driver/internal/odbc_sql_execute_utils_test.cc @@ -331,4 +331,62 @@ TEST(FetchNextPageResultSet, FailureSQLNODATA) { EXPECT_EQ(status_record_or.message, "No more data to return."); } +#if (!defined(_WIN32) || defined(_WIN64)) && !defined(NO_ARROW) +TEST(IsTransientStreamError, IdentifiesTransientAndPermanentErrors) { + using google::cloud::Status; + using google::cloud::StatusCode; + + EXPECT_TRUE( + IsTransientStreamError(Status(StatusCode::kUnavailable, "unavailable"))); + EXPECT_TRUE(IsTransientStreamError( + Status(StatusCode::kDeadlineExceeded, "deadline exceeded"))); + EXPECT_TRUE(IsTransientStreamError(Status(StatusCode::kAborted, "aborted"))); + EXPECT_TRUE(IsTransientStreamError( + Status(StatusCode::kResourceExhausted, "quota exceeded"))); + + EXPECT_TRUE(IsTransientStreamError( + Status(StatusCode::kInternal, "Received RST_STREAM with error code 2"))); + EXPECT_TRUE( + IsTransientStreamError(Status(StatusCode::kUnknown, "Stream removed"))); + EXPECT_TRUE(IsTransientStreamError( + Status(StatusCode::kInternal, "Connection reset by peer"))); + EXPECT_TRUE(IsTransientStreamError( + Status(StatusCode::kUnknown, "transport: lost connection"))); + + EXPECT_FALSE(IsTransientStreamError( + Status(StatusCode::kInternal, "Permanent internal parser error"))); + EXPECT_FALSE(IsTransientStreamError( + Status(StatusCode::kInvalidArgument, "Invalid offset"))); + EXPECT_FALSE(IsTransientStreamError( + Status(StatusCode::kPermissionDenied, "Access denied"))); +} + +TEST(HtapiStreamRetryPolicy, RetriesTransientErrorsUpToMaxRetries) { + using google::cloud::Status; + using google::cloud::StatusCode; + + HtapiStreamRetryPolicy policy(/*maximum_failures=*/2); + EXPECT_FALSE(policy.IsExhausted()); + EXPECT_FALSE(policy.IsPermanentFailure( + Status(StatusCode::kUnknown, "Stream removed"))); + EXPECT_TRUE(policy.IsPermanentFailure( + Status(StatusCode::kPermissionDenied, "Access denied"))); + + EXPECT_TRUE(policy.OnFailure(Status(StatusCode::kUnknown, "Stream removed"))); + EXPECT_FALSE(policy.IsExhausted()); + EXPECT_TRUE(policy.OnFailure( + Status(StatusCode::kInternal, "Received RST_STREAM with error code 2"))); + EXPECT_FALSE(policy.IsExhausted()); + EXPECT_FALSE( + policy.OnFailure(Status(StatusCode::kUnknown, "Stream removed"))); + EXPECT_TRUE(policy.IsExhausted()); + + auto cloned = policy.clone(); + EXPECT_FALSE(cloned->IsExhausted()); + EXPECT_FALSE(cloned->OnFailure( + Status(StatusCode::kInvalidArgument, "Invalid offset"))); + EXPECT_TRUE(cloned->IsExhausted()); +} +#endif // (!defined(_WIN32) || defined(_WIN64)) && !defined(NO_ARROW) + } // namespace google::cloud::odbc_bq_driver_internal diff --git a/google/cloud/odbc/bq_driver/internal/odbc_stmt_handle.cc b/google/cloud/odbc/bq_driver/internal/odbc_stmt_handle.cc index 52ac0902d1..b71a0e2535 100644 --- a/google/cloud/odbc/bq_driver/internal/odbc_stmt_handle.cc +++ b/google/cloud/odbc/bq_driver/internal/odbc_stmt_handle.cc @@ -570,9 +570,6 @@ void StatementHandle::CloseCursor() { #if (!defined(_WIN32) || defined(_WIN64)) && !defined(NO_ARROW) ClearReadRowsStream(); ClearReadRowsIterator(); - ClearReadStreamName(); - ResetStreamRowOffset(); - SetWasHtapiEnabled(false); #endif if (StatementPrepared()) { SetStmtState(StmtStates::kStatementPrepared); diff --git a/google/cloud/odbc/bq_driver/internal/odbc_stmt_handle.h b/google/cloud/odbc/bq_driver/internal/odbc_stmt_handle.h index ba57e571ec..69c5df367d 100644 --- a/google/cloud/odbc/bq_driver/internal/odbc_stmt_handle.h +++ b/google/cloud/odbc/bq_driver/internal/odbc_stmt_handle.h @@ -134,25 +134,8 @@ class StatementHandle : public Handle { // Clearer for the iterator inline void ClearReadRowsIterator() { read_rows_iterator_.reset(); } - // Returns true if HTAPI was enabled for this query. - inline bool WasHtapiEnabled() const { - return was_htapi_enabled_ || read_rows_stream_.has_value(); - } - inline void SetWasHtapiEnabled(bool enabled) { was_htapi_enabled_ = enabled; } - - inline std::string const& GetReadStreamName() const { - return read_stream_name_; - } - inline void SetReadStreamName(std::string name) { - read_stream_name_ = std::move(name); - } - inline void ClearReadStreamName() { read_stream_name_.clear(); } - - inline std::int64_t GetStreamRowOffset() const { return stream_row_offset_; } - inline void SetStreamRowOffset(std::int64_t offset) { - stream_row_offset_ = offset; - } - inline void ResetStreamRowOffset() { stream_row_offset_ = 0; } + // Returns true if read_rows_stream_ is NOT null (i.e., has a value). + inline bool WasHtapiEnabled() const { return read_rows_stream_.has_value(); } void SetReadRowsStream( StreamRange<::google::cloud::bigquery::storage::v1::ReadRowsResponse> @@ -363,9 +346,6 @@ class StatementHandle : public Handle { std::optional< StreamRange<::google::cloud::bigquery::storage::v1::ReadRowsResponse>> read_rows_stream_; - std::string read_stream_name_; - std::int64_t stream_row_offset_{0}; - bool was_htapi_enabled_{false}; mutable std::mutex statement_handle_mutex_; ::google::cloud::bigquery_v2_minimal_internal::PostQueryRequest post_query_request_; diff --git a/google/cloud/odbc/integration_tests/odbc_driver_tests/resilience_test.cc b/google/cloud/odbc/integration_tests/odbc_driver_tests/resilience_test.cc index 6aeef40873..59ae9e6aae 100644 --- a/google/cloud/odbc/integration_tests/odbc_driver_tests/resilience_test.cc +++ b/google/cloud/odbc/integration_tests/odbc_driver_tests/resilience_test.cc @@ -14,7 +14,6 @@ #include "google/cloud/odbc/testing/odbc_utils/connection.h" #include "google/cloud/odbc/testing/odbc_utils/statement.h" -#include #include #include #include @@ -26,8 +25,6 @@ #include #include -using ::testing::HasSubstr; - namespace google::cloud::odbc_tests { #ifdef BQ_DRIVER_INTEGRATION_TESTS @@ -294,33 +291,21 @@ TEST(ResilienceTest, ReadAPI_IdleProxy_ConnectionDroppedBetweenQueries) { EXPECT_EQ(Disconnect(conn), SQL_SUCCESS); } -struct ScopedEnvVar { - char const* name; - ScopedEnvVar(char const* env_name, char const* value) : name(env_name) { -#ifdef _WIN32 - _putenv_s(name, value); -#else - setenv(name, value, 1); -#endif - } - ~ScopedEnvVar() { -#ifdef _WIN32 - _putenv_s(name, ""); -#else - unsetenv(name); -#endif +TEST(StatementTest, SQLExecDirect_HTAPI_StreamRetryAndResume) { + auto const config = GetIdleProxyConfig(); + if (!config) { + GTEST_SKIP() << "Set BQ_ODBC_IDLE_PROXY=host:port to run idle proxy tests"; } -}; -TEST(StatementTest, SQLExecDirect_HTAPI_StreamRetryAndResume) { - ScopedEnvVar env_guard("BIGQUERY_ODBC_TEST_INJECT_STREAM_ERROR", "1"); + IdleProxyStats const before = + GetIdleProxyStats(config->stats_file, config->kills_file); auto conn = std::make_shared(); std::string connection_string = - kDefaultConnectionString + - ";AllowHtapiForLargeResults=1;UseDefaultLargeResultsDataset=0;" + IdleProxyConnString(*config) + + ";UseDefaultLargeResultsDataset=0;" "LargeResultsDataSetId=_bqodbc_temp_tables_euwest1"; - EXPECT_EQ(Connect(connection_string, conn), SQL_SUCCESS); + ASSERT_EQ(Connect(connection_string, conn), SQL_SUCCESS); int const limit = 3000; std::string query = @@ -332,15 +317,20 @@ TEST(StatementTest, SQLExecDirect_HTAPI_StreamRetryAndResume) { Table table("Random_table_name"); RowWiseResults const& results = table.Fetch(conn, query); - // Assert proof of issue reproduction, error caught, and reconnection: - // The driver captures the injected error and records the resumption offset. - char const* test_status = - std::getenv("BIGQUERY_ODBC_TEST_INJECT_STREAM_ERROR"); - ASSERT_NE(test_status, nullptr) - << "Verification failed: Transient stream error was never injected or " - "reconnection never fired!"; - EXPECT_THAT(test_status, HasSubstr("RESUMED_AT_OFFSET=")); - EXPECT_THAT(test_status, HasSubstr("ERROR=UNAVAILABLE: RST_STREAM")); + IdleProxyStats const after = + GetIdleProxyStats(config->stats_file, config->kills_file); + + int const kills = after.kills - before.kills; + int const reconnects = after.reconnects - before.reconnects; + int const storage_tunnels = after.storage_tunnels - before.storage_tunnels; + + EXPECT_GE(storage_tunnels, 3) + << "Expected initial and reconnected BigQuery Storage tunnels through " + "the proxy"; + EXPECT_GE(kills, 2) + << "Expected proxy to kill the active HTAPI stream mid-query"; + EXPECT_GE(reconnects, 2) + << "Expected driver to reconnect and resume HTAPI stream after drops"; int const expected_num_cols = 300; ASSERT_EQ(results.size(), limit) << "Row count mismatch."; diff --git a/tools/idle_connect_proxy/idle_connect_proxy.py b/tools/idle_connect_proxy/idle_connect_proxy.py index 4fa4a3bffe..25aef4326c 100644 --- a/tools/idle_connect_proxy/idle_connect_proxy.py +++ b/tools/idle_connect_proxy/idle_connect_proxy.py @@ -142,6 +142,16 @@ def kill(tunnel, reason): tunnel.state = "killed" reset_writer(tunnel.client_writer) reset_writer(tunnel.server_writer) + elif ARGS.mode == "corrupt": + tunnel.state = "killed" + try: + tunnel.client_writer.write( + b"HTTP/1.1 502 Bad Gateway\r\n\r\n" + b"\xff" * 64 + ) + except Exception: + pass + close_writer(tunnel.client_writer) + reset_writer(tunnel.server_writer) elif ARGS.mode == "blackhole": tunnel.state = "blackholed" reset_writer(tunnel.server_writer) @@ -173,6 +183,32 @@ async def relay(reader, writer, tunnel, direction): "%s swallowing %d bytes (%s)", tunnel, len(data), tunnel.state ) continue + if direction == "down": + kill_threshold = ( + ARGS.reconnect_kill_after_bytes + if STATS["kills"] > 0 and ARGS.reconnect_kill_after_bytes > 0 + else ARGS.kill_after_bytes + ) + if ( + kill_threshold > 0 + and (ARGS.max_kills <= 0 or STATS["kills"] < ARGS.max_kills) + and tunnel.monitored + and tunnel.state == "open" + and tunnel.down_bytes + len(data) >= kill_threshold + ): + tunnel.down_bytes += len(data) + if ARGS.mode == "corrupt": + corrupted = bytes(b ^ 0xFF for b in data) + writer.write(corrupted) + await writer.drain() + else: + writer.write(data) + await writer.drain() + kill( + tunnel, + f"{tunnel.down_bytes} bytes relayed >= {kill_threshold}", + ) + return "killed" writer.write(data) await writer.drain() now = time.monotonic() @@ -201,17 +237,6 @@ async def relay(reader, writer, tunnel, direction): tunnel.down_bytes += len(data) if len(data) > 256: tunnel.last_down_activity = now - if ( - ARGS.kill_after_bytes > 0 - and tunnel.monitored - and tunnel.state == "open" - and tunnel.down_bytes >= ARGS.kill_after_bytes - ): - kill( - tunnel, - f"{tunnel.down_bytes} bytes relayed >= {ARGS.kill_after_bytes}", - ) - return "killed" except asyncio.CancelledError: raise except (ConnectionError, OSError, asyncio.IncompleteReadError) as e: @@ -332,7 +357,7 @@ async def handle_client(client_reader, client_writer): tunnel, tunnel.age(), ) - if not TUNNELS: + elif not TUNNELS and tunnel.kill_reason is None: SESSION_HAD_DROP = False how = ( f"killed by proxy ({tunnel.kill_reason})" @@ -405,10 +430,22 @@ def parse_args(argv): default=0, help="kill monitored tunnels once this many bytes relayed (default: 0)", ) + parser.add_argument( + "--reconnect-kill-after-bytes", + type=int, + default=0, + help="kill reconnected tunnels once this many bytes relayed (default: 0)", + ) + parser.add_argument( + "--max-kills", + type=int, + default=0, + help="maximum number of tunnels to kill (default: 0 = unlimited)", + ) parser.add_argument( "--mode", default="reset", - choices=["close", "reset", "blackhole", "stale-reset"], + choices=["close", "reset", "corrupt", "blackhole", "stale-reset"], help="how to kill a tunnel (default: reset)", ) parser.add_argument(