KanchanShu · Oct 22, 2026 · Oct 23, 2026 · shivamd-gpartner · Oct 23, 2026 · KanchanShu
@@ -741,66 +741,190 @@ 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("transport: lost connection") != std::string::npos) {
+        return true;
+      }
+      return false;
+    }
+    default:
+      return false;
+  }
+  return false;
+}
+
 StatusRecord ReadNextResultsFromStream(StatementHandle& stmt_handle) {
-  std::optional<StreamRange<ReadRowsResponse>>& 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<ODBCBQClient> bq_client = conn_handle->GetClient();
+  if (!bq_client) {
+    return StatusRecord{
+        SQLStates::k_HY000(),
+        "Invalid or null BQ Client within the connection handle"};
   }
-  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());
+
+  int max_retries = static_cast<int>(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<StreamRange<ReadRowsResponse>>& 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<MaxRetriesOption>(conn_handle->GetDsn().max_retries);
+      StreamRange<ReadRowsResponse> new_stream =
+          bq_client->GetReadRowsStream(read_rows_request, options);
+      stmt_handle.SetReadRowsStream(std::move(new_stream));
+      stmt_handle.ClearReadRowsIterator();
+    }
+
+    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);
     }
-    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<arrow::Schema> schema = stmt_handle.GetArrowSchema();
-      StatusRecordOr<std::shared_ptr<arrow::RecordBatch>> record_batch_status =
-          GetArrowRecordBatch(row.arrow_record_batch(), schema);
-      if (!record_batch_status) {
-        return record_batch_status.GetStatusRecord();
+    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<std::chrono::milliseconds>(
+                              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<arrow::Schema> schema = stmt_handle.GetArrowSchema();
+        StatusRecordOr<std::shared_ptr<arrow::RecordBatch>>
+            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 +968,14 @@ StatusRecord FetchBQDataReadArrow(StatementHandle& stmt_handle,
     std::shared_ptr<arrow::Schema> 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<google::cloud::bigquery::storage::v1::ReadRowsResponse>

@@ -62,9 +62,17 @@ odbc_internal::StatusRecordOr<DSResults> 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);
 

@@ -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 {

@@ -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_;

@@ -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<ODBCHandles>();
+  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<ODBCHandles>();