From a29536a061600cfecfabe819d401f34603f37944 Mon Sep 17 00:00:00 2001 From: Darren Bolduc Date: Tue, 15 Sep 2026 02:10:17 -0400 Subject: [PATCH 1/2] impl(bigquery): default write retry policy --- src/bigquery/src/write.rs | 2 + src/bigquery/src/write/retry_policy.rs | 189 +++++++++++++++++++++++++ 2 files changed, 191 insertions(+) create mode 100644 src/bigquery/src/write/retry_policy.rs diff --git a/src/bigquery/src/write.rs b/src/bigquery/src/write.rs index 7d25e140f6..014a94b845 100644 --- a/src/bigquery/src/write.rs +++ b/src/bigquery/src/write.rs @@ -32,6 +32,8 @@ mod dispatcher; mod entry; mod pool; mod proto_schema; +#[allow(dead_code)] +mod retry_policy; mod runner; mod stream; mod transport; diff --git a/src/bigquery/src/write/retry_policy.rs b/src/bigquery/src/write/retry_policy.rs new file mode 100644 index 0000000000..297642bb45 --- /dev/null +++ b/src/bigquery/src/write/retry_policy.rs @@ -0,0 +1,189 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Defines the retry policy for the BigQuery Storage Write API. + +use google_cloud_gax::error::Error; +use google_cloud_gax::error::rpc::Code; +use google_cloud_gax::retry_policy::RetryPolicy; +use google_cloud_gax::retry_result::RetryResult; +use google_cloud_gax::retry_state::RetryState; + +/// Follows the RPC retry strategy recommended for BigQuery Storage Write API. +/// +/// This policy must be decorated to limit the duration of the retry loop or +/// the number of attempts. +#[derive(Clone, Debug)] +#[allow(dead_code)] +pub struct RetryableErrors; + +impl RetryPolicy for RetryableErrors { + fn on_error(&self, _state: &RetryState, error: Error) -> RetryResult { + if error.is_transient_and_before_rpc() + || error.is_io() + || error.is_timeout() + || error.is_connect() + { + return RetryResult::Continue(error); + } + if error.is_transport() && error.http_status_code().is_none() { + return RetryResult::Continue(error); + } + if let Some(429 | 500 | 502 | 503 | 504) = error.http_status_code() { + return RetryResult::Continue(error); + } + if let Some(status) = error.status() { + return match status.code { + Code::Aborted + | Code::DeadlineExceeded + | Code::Internal + | Code::ResourceExhausted + | Code::Unavailable => RetryResult::Continue(error), + _ => RetryResult::Permanent(error), + }; + } + RetryResult::Permanent(error) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use google_cloud_gax::error::CredentialsError; + use google_cloud_gax::error::rpc::Status; + use http::HeaderMap; + use test_case::test_case; + + #[test_case(Code::Aborted)] + #[test_case(Code::DeadlineExceeded)] + #[test_case(Code::Internal)] + #[test_case(Code::ResourceExhausted)] + #[test_case(Code::Unavailable)] + fn retryable_status_codes(code: Code) { + let err = Error::service(Status::default().set_code(code).set_message("try again")); + assert!(matches!( + RetryableErrors.on_error(&RetryState::default(), err), + RetryResult::Continue(_) + )); + } + + #[test_case(Code::Cancelled)] + #[test_case(Code::Unknown)] + #[test_case(Code::InvalidArgument)] + #[test_case(Code::NotFound)] + #[test_case(Code::PermissionDenied)] + #[test_case(Code::FailedPrecondition)] + #[test_case(Code::DataLoss)] + fn permanent_status_codes(code: Code) { + let err = Error::service(Status::default().set_code(code).set_message("fail")); + assert!(matches!( + RetryableErrors.on_error(&RetryState::default(), err), + RetryResult::Permanent(_) + )); + } + + #[test_case(429)] + #[test_case(500)] + #[test_case(502)] + #[test_case(503)] + #[test_case(504)] + fn retryable_http_status_codes(code: u16) { + let err = Error::http(code, HeaderMap::new(), bytes::Bytes::new()); + assert!(matches!( + RetryableErrors.on_error(&RetryState::default(), err), + RetryResult::Continue(_) + )); + } + + #[test_case(400)] + #[test_case(401)] + #[test_case(403)] + #[test_case(404)] + #[test_case(408)] + #[test_case(409)] + #[test_case(501)] + fn permanent_http_status_codes(code: u16) { + let err = Error::http(code, HeaderMap::new(), bytes::Bytes::new()); + assert!(matches!( + RetryableErrors.on_error(&RetryState::default(), err), + RetryResult::Permanent(_) + )); + } + + #[test] + fn retry_transient_before_rpc() { + let err = Error::authentication(CredentialsError::from_msg(true, "try again")); + assert!(matches!( + RetryableErrors.on_error(&RetryState::default(), err), + RetryResult::Continue(_) + )); + + let err = Error::authentication(CredentialsError::from_msg(false, "fail")); + assert!(matches!( + RetryableErrors.on_error(&RetryState::default(), err), + RetryResult::Permanent(_) + )); + } + + #[test] + fn retry_io() { + let err = Error::io("try again"); + assert!(matches!( + RetryableErrors.on_error(&RetryState::default(), err), + RetryResult::Continue(_) + )); + } + + #[test] + fn retry_timeout() { + let err = Error::timeout("deadline exceeded"); + assert!(matches!( + RetryableErrors.on_error(&RetryState::default(), err), + RetryResult::Continue(_) + )); + } + + #[test] + fn retry_connect() { + let err = Error::connect("connection failed"); + assert!(matches!( + RetryableErrors.on_error(&RetryState::default(), err), + RetryResult::Continue(_) + )); + } + + #[test] + fn retry_transport_without_http_status() { + let err = Error::transport(HeaderMap::new(), "connection reset"); + assert!(matches!( + RetryableErrors.on_error(&RetryState::default(), err), + RetryResult::Continue(_) + )); + } + + #[test] + fn permanent_serde() { + let err = Error::ser("fail"); + assert!(matches!( + RetryableErrors.on_error(&RetryState::default(), err), + RetryResult::Permanent(_) + )); + + let err = Error::deser("fail"); + assert!(matches!( + RetryableErrors.on_error(&RetryState::default(), err), + RetryResult::Permanent(_) + )); + } +} From d8d12cacf00ad96a2331b772d2f7095926be4c09 Mon Sep 17 00:00:00 2001 From: Darren Bolduc Date: Tue, 15 Sep 2026 11:44:21 -0400 Subject: [PATCH 2/2] address review comments --- src/bigquery/src/write/retry_policy.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/bigquery/src/write/retry_policy.rs b/src/bigquery/src/write/retry_policy.rs index 297642bb45..7dfe814373 100644 --- a/src/bigquery/src/write/retry_policy.rs +++ b/src/bigquery/src/write/retry_policy.rs @@ -46,6 +46,7 @@ impl RetryPolicy for RetryableErrors { if let Some(status) = error.status() { return match status.code { Code::Aborted + | Code::Cancelled | Code::DeadlineExceeded | Code::Internal | Code::ResourceExhausted @@ -66,6 +67,7 @@ mod tests { use test_case::test_case; #[test_case(Code::Aborted)] + #[test_case(Code::Cancelled)] #[test_case(Code::DeadlineExceeded)] #[test_case(Code::Internal)] #[test_case(Code::ResourceExhausted)] @@ -78,7 +80,6 @@ mod tests { )); } - #[test_case(Code::Cancelled)] #[test_case(Code::Unknown)] #[test_case(Code::InvalidArgument)] #[test_case(Code::NotFound)]