Htapi stream retry fix

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"

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"

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 <arrow/array/array_decimal.h>
 #endif
 #include <thread>
@@ -778,6 +779,38 @@ 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 closed") != std::string::npos ||
+          message.find("endpoint read failed") != 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();
@@ -856,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<int>(conn_handle.GetDsn().max_retries);
   Options options;
   options.set<MaxRetriesOption>(conn_handle.GetDsn().max_retries);
+  options
+      .set<::google::cloud::bigquery_storage_v1::BigQueryReadRetryPolicyOption>(
+          std::make_shared<HtapiStreamRetryPolicy>(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);

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 {
 
@@ -72,6 +75,55 @@ 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 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<HtapiStreamRetryPolicy>(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

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

google/cloud/odbc/integration_tests/odbc_driver_tests/resilience_test.cc

@@ -291,6 +291,66 @@ TEST(ResilienceTest, ReadAPI_IdleProxy_ConnectionDroppedBetweenQueries) {
   EXPECT_EQ(Disconnect(conn), SQL_SUCCESS);
 }
 
+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";
+  }
+
+  IdleProxyStats const before =
+      GetIdleProxyStats(config->stats_file, config->kills_file);
+
+  auto conn = std::make_shared<ODBCHandles>();
+  std::string connection_string =
+      IdleProxyConnString(*config) +
+      ";UseDefaultLargeResultsDataset=0;"
+      "LargeResultsDataSetId=_bqodbc_temp_tables_euwest1";
+  ASSERT_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);
+
+  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.";
+  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

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(