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 264f90691a..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 @@ -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>& 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(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); 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..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 { @@ -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(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 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/integration_tests/odbc_driver_tests/resilience_test.cc b/google/cloud/odbc/integration_tests/odbc_driver_tests/resilience_test.cc index 3598c43f68..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 @@ -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(); + 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 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(