From a65b26960e852784e426727612b251027dadad3c Mon Sep 17 00:00:00 2001 From: Alvaro Viebrantz Date: Wed, 2 Sep 2026 19:34:10 +0000 Subject: [PATCH 1/3] impl(bigquery): allow retry/backoff policy override and use default policies --- src/bigquery/src/query/client.rs | 17 ++++--- src/bigquery/src/query/client_builder.rs | 59 ++++++++++++++++++++++++ src/bigquery/src/query/retry_policy.rs | 3 -- 3 files changed, 70 insertions(+), 9 deletions(-) diff --git a/src/bigquery/src/query/client.rs b/src/bigquery/src/query/client.rs index 9be60d0349..1f71646ec1 100644 --- a/src/bigquery/src/query/client.rs +++ b/src/bigquery/src/query/client.rs @@ -98,12 +98,17 @@ impl BigQuery { if builder.config.tracing { job_service_builder = job_service_builder.with_tracing(); } - if let Some(retry_policy) = builder.config.retry_policy { - job_service_builder = job_service_builder.with_retry_policy(retry_policy); - } - if let Some(backoff_policy) = builder.config.backoff_policy { - job_service_builder = job_service_builder.with_backoff_policy(backoff_policy); - } + let retry_policy = builder + .config + .retry_policy + .unwrap_or_else(crate::query::retry_policy::default_retry_policy); + job_service_builder = job_service_builder.with_retry_policy(retry_policy); + + let backoff_policy = builder + .config + .backoff_policy + .unwrap_or_else(crate::query::retry_policy::default_backoff_policy); + job_service_builder = job_service_builder.with_backoff_policy(backoff_policy); job_service_builder = job_service_builder.with_retry_throttler(builder.config.retry_throttler); let job_service = Arc::new(job_service_builder.build().await?); diff --git a/src/bigquery/src/query/client_builder.rs b/src/bigquery/src/query/client_builder.rs index 80e57acb55..f0b0e47146 100644 --- a/src/bigquery/src/query/client_builder.rs +++ b/src/bigquery/src/query/client_builder.rs @@ -157,6 +157,57 @@ impl ClientBuilder { self } + /// Configure the retry policy. + /// + /// The client libraries can automatically retry operations that fail. The + /// retry policy controls what errors are considered retryable, sets limits + /// on the number of attempts or the time trying to make attempts. + /// + /// # Example + /// ``` + /// # use google_cloud_bigquery::client::BigQuery; + /// # async fn sample() -> anyhow::Result<()> { + /// use google_cloud_gax::retry_policy::{AlwaysRetry, RetryPolicyExt}; + /// let client = BigQuery::builder() + /// .with_retry_policy(AlwaysRetry.with_attempt_limit(3)) + /// .build() + /// .await?; + /// # Ok(()) } + /// ``` + pub fn with_retry_policy>( + mut self, + v: V, + ) -> Self { + self.config.retry_policy = Some(v.into().into()); + self + } + + /// Configure the retry backoff policy. + /// + /// The client libraries can automatically retry operations that fail. The + /// backoff policy controls how long to wait in between retry attempts. + /// + /// # Example + /// ``` + /// # use google_cloud_bigquery::client::BigQuery; + /// # async fn sample() -> anyhow::Result<()> { + /// use google_cloud_gax::exponential_backoff::ExponentialBackoff; + /// use std::time::Duration; + /// let policy = ExponentialBackoff::default(); + /// let client = BigQuery::builder() + /// .with_backoff_policy(policy) + /// .build() + /// .await?; + /// # Ok(()) } + /// ``` + pub fn with_backoff_policy>( + mut self, + v: V, + ) -> Self { + self.config.backoff_policy = Some(v.into().into()); + self + } + /// Creates a new [`BigQuery`] client. /// /// # Example @@ -175,6 +226,8 @@ impl ClientBuilder { mod tests { use super::*; use google_cloud_auth::credentials::anonymous::Builder as Anonymous; + use google_cloud_gax::exponential_backoff::ExponentialBackoff; + use google_cloud_gax::retry_policy::AlwaysRetry; #[test] fn defaults() -> anyhow::Result<()> { @@ -183,6 +236,8 @@ mod tests { assert!(builder.config.universe_domain.is_none(), "{builder:?}"); assert!(builder.config.cred.is_none(), "{builder:?}"); assert!(!builder.config.tracing); + assert!(builder.config.retry_policy.is_none(), "{builder:?}"); + assert!(builder.config.backoff_policy.is_none(), "{builder:?}"); assert!(builder.project_id.is_none(), "{builder:?}"); Ok(()) @@ -195,6 +250,8 @@ mod tests { .with_endpoint("test-endpoint.com") .with_universe_domain("test-universe.com") .with_credentials(Anonymous::new().build()) + .with_retry_policy(AlwaysRetry) + .with_backoff_policy(ExponentialBackoff::default()) .with_tracing(); assert_eq!(builder.project_id, Some("test-project".to_string())); @@ -208,6 +265,8 @@ mod tests { ); assert!(builder.config.cred.is_some(), "{builder:?}"); assert!(builder.config.tracing); + assert!(builder.config.retry_policy.is_some(), "{builder:?}"); + assert!(builder.config.backoff_policy.is_some(), "{builder:?}"); Ok(()) } diff --git a/src/bigquery/src/query/retry_policy.rs b/src/bigquery/src/query/retry_policy.rs index 6d7bf94dcf..a6dd224916 100644 --- a/src/bigquery/src/query/retry_policy.rs +++ b/src/bigquery/src/query/retry_policy.rs @@ -29,7 +29,6 @@ use std::time::Duration; /// Follows the RPC retry strategy recommended by the BigQuery guides on error handling. #[derive(Clone, Debug)] -#[allow(dead_code)] pub(crate) struct RetryableErrors; impl RetryPolicy for RetryableErrors { @@ -58,7 +57,6 @@ impl RetryPolicy for RetryableErrors { } } -#[allow(dead_code)] pub(crate) fn default_retry_policy() -> Arc { Arc::new(RetryableErrors) } @@ -166,7 +164,6 @@ pub(crate) fn is_retryable_errors(errors: &[ErrorProto]) -> bool { !errors.is_empty() && errors.iter().all(|e| is_retryable_error_reason(&e.reason)) } -#[allow(dead_code)] pub(crate) fn is_retryable_error_reason(reason: &str) -> bool { matches!( reason, From 99de16b1789752a879ba94d3404523f4facf6126 Mon Sep 17 00:00:00 2001 From: Alvaro Viebrantz Date: Wed, 2 Sep 2026 19:39:30 +0000 Subject: [PATCH 2/3] fix: address ai review comments --- src/bigquery/src/query/client_builder.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/src/bigquery/src/query/client_builder.rs b/src/bigquery/src/query/client_builder.rs index f0b0e47146..5834c6b722 100644 --- a/src/bigquery/src/query/client_builder.rs +++ b/src/bigquery/src/query/client_builder.rs @@ -192,7 +192,6 @@ impl ClientBuilder { /// # use google_cloud_bigquery::client::BigQuery; /// # async fn sample() -> anyhow::Result<()> { /// use google_cloud_gax::exponential_backoff::ExponentialBackoff; - /// use std::time::Duration; /// let policy = ExponentialBackoff::default(); /// let client = BigQuery::builder() /// .with_backoff_policy(policy) From 28464421645af77f2ba0c36a3da375bb8869cbfc Mon Sep 17 00:00:00 2001 From: Alvaro Viebrantz Date: Thu, 10 Sep 2026 15:58:42 +0000 Subject: [PATCH 3/3] impl: expose RetryableErrors --- src/bigquery/src/query.rs | 2 +- src/bigquery/src/query/client_builder.rs | 10 ++++++---- src/bigquery/src/query/retry_policy.rs | 24 ++++++++++++++++++++++-- 3 files changed, 29 insertions(+), 7 deletions(-) diff --git a/src/bigquery/src/query.rs b/src/bigquery/src/query.rs index 2333d70468..ae031e27f4 100644 --- a/src/bigquery/src/query.rs +++ b/src/bigquery/src/query.rs @@ -19,7 +19,7 @@ mod execution; pub(super) mod from_sql; mod iterator; mod query_handle; -mod retry_policy; +pub mod retry_policy; mod row; mod schema; diff --git a/src/bigquery/src/query/client_builder.rs b/src/bigquery/src/query/client_builder.rs index 5834c6b722..f363c8f0e0 100644 --- a/src/bigquery/src/query/client_builder.rs +++ b/src/bigquery/src/query/client_builder.rs @@ -167,9 +167,10 @@ impl ClientBuilder { /// ``` /// # use google_cloud_bigquery::client::BigQuery; /// # async fn sample() -> anyhow::Result<()> { - /// use google_cloud_gax::retry_policy::{AlwaysRetry, RetryPolicyExt}; + /// use google_cloud_bigquery::query::retry_policy::RetryableErrors; + /// use google_cloud_gax::retry_policy::RetryPolicyExt; /// let client = BigQuery::builder() - /// .with_retry_policy(AlwaysRetry.with_attempt_limit(3)) + /// .with_retry_policy(RetryableErrors.with_attempt_limit(3)) /// .build() /// .await?; /// # Ok(()) } @@ -192,6 +193,7 @@ impl ClientBuilder { /// # use google_cloud_bigquery::client::BigQuery; /// # async fn sample() -> anyhow::Result<()> { /// use google_cloud_gax::exponential_backoff::ExponentialBackoff; + /// use std::time::Duration; /// let policy = ExponentialBackoff::default(); /// let client = BigQuery::builder() /// .with_backoff_policy(policy) @@ -224,9 +226,9 @@ impl ClientBuilder { #[cfg(test)] mod tests { use super::*; + use crate::query::retry_policy::RetryableErrors; use google_cloud_auth::credentials::anonymous::Builder as Anonymous; use google_cloud_gax::exponential_backoff::ExponentialBackoff; - use google_cloud_gax::retry_policy::AlwaysRetry; #[test] fn defaults() -> anyhow::Result<()> { @@ -249,7 +251,7 @@ mod tests { .with_endpoint("test-endpoint.com") .with_universe_domain("test-universe.com") .with_credentials(Anonymous::new().build()) - .with_retry_policy(AlwaysRetry) + .with_retry_policy(RetryableErrors) .with_backoff_policy(ExponentialBackoff::default()) .with_tracing(); diff --git a/src/bigquery/src/query/retry_policy.rs b/src/bigquery/src/query/retry_policy.rs index af17340c4b..7cdd049eb7 100644 --- a/src/bigquery/src/query/retry_policy.rs +++ b/src/bigquery/src/query/retry_policy.rs @@ -27,9 +27,29 @@ use google_cloud_gax::retry_state::RetryState; use std::sync::Arc; use std::time::Duration; -/// Follows the RPC retry strategy recommended by the BigQuery guides on error handling. +/// Follows the RPC retry strategy recommended by the BigQuery guides on +/// [error handling]. +/// +/// ``` +/// # async fn sample() -> anyhow::Result<()> { +/// # use google_cloud_bigquery::client::BigQuery; +/// # use google_cloud_bigquery::query::retry_policy::RetryableErrors; +/// # use google_cloud_gax::retry_policy::RetryPolicyExt; +/// let policy = RetryableErrors.with_time_limit(std::time::Duration::from_secs(60)); +/// let client = BigQuery::builder() +/// .with_retry_policy(policy) +/// .build() +/// .await?; +/// # Ok(()) +/// # } +/// ``` +/// +/// This policy must be decorated to limit the duration of the retry loop or +/// the number of attempts. +/// +/// [error handling]: https://cloud.google.com/bigquery/docs/error-messages #[derive(Clone, Debug)] -pub(crate) struct RetryableErrors; +pub struct RetryableErrors; impl RetryPolicy for RetryableErrors { fn on_error(&self, _state: &RetryState, error: GaxError) -> RetryResult {