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 80e57acb55..f363c8f0e0 100644 --- a/src/bigquery/src/query/client_builder.rs +++ b/src/bigquery/src/query/client_builder.rs @@ -157,6 +157,58 @@ 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_bigquery::query::retry_policy::RetryableErrors; + /// use google_cloud_gax::retry_policy::RetryPolicyExt; + /// let client = BigQuery::builder() + /// .with_retry_policy(RetryableErrors.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 @@ -174,7 +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; #[test] fn defaults() -> anyhow::Result<()> { @@ -183,6 +237,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 +251,8 @@ mod tests { .with_endpoint("test-endpoint.com") .with_universe_domain("test-universe.com") .with_credentials(Anonymous::new().build()) + .with_retry_policy(RetryableErrors) + .with_backoff_policy(ExponentialBackoff::default()) .with_tracing(); assert_eq!(builder.project_id, Some("test-project".to_string())); @@ -208,6 +266,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 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 {