From 237bbd89658370e5c697aee1b145d71ef72f0136 Mon Sep 17 00:00:00 2001 From: KanchanShu Date: Thu, 22 Oct 2026 07:46:47 +0000 Subject: [PATCH 1/2] 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 | 5 + .../bq_driver/internal/odbc_stmt_handle.h | 24 +- 4 files changed, 198 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 dbfb02792f..4ac2a2df39 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 @@ -741,66 +741,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, @@ -844,9 +949,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 ffd2665145..e9f46fe547 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 @@ -62,9 +62,17 @@ odbc_internal::StatusRecordOr ExecuteScript( 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 c2dfdd162e..9ac4e8138b 100644 --- a/google/cloud/odbc/bq_driver/internal/odbc_stmt_handle.cc +++ b/google/cloud/odbc/bq_driver/internal/odbc_stmt_handle.cc @@ -557,6 +557,11 @@ StatusRecord StatementHandle::PopulateIpd(DescriptorHandle& handle, void StatementHandle::CloseCursor() { ResultSet result_set; result_set_ = result_set; + ClearReadRowsStream(); + ClearReadRowsIterator(); + ClearReadStreamName(); + ResetStreamRowOffset(); + SetWasHtapiEnabled(false); if (StatementPrepared()) { SetStmtState(StmtStates::kStatementPrepared); } else { 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 08e70a3d5e..af0d938518 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> @@ -345,6 +362,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 0f7d73a5154ebb160f80775f5672e9f8e03233a6 Mon Sep 17 00:00:00 2001 From: KanchanShu Date: Fri, 23 Oct 2026 13:09:26 +0000 Subject: [PATCH 2/2] 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 4ac2a2df39..5c5d3cbbe9 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 @@ -760,9 +760,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; } @@ -788,10 +785,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); @@ -834,9 +828,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 7e1db5cda9..30f2d430d5 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 @@ -731,6 +731,67 @@ TEST_P(HTAPIParameterizedTest, SQLExecDirect_with_pagination) { 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();