feat(bigquery): allow retry/backoff policy override and use default policies

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;
 

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<V: Into<google_cloud_gax::retry_policy::RetryPolicyArg>>(
+        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<V: Into<google_cloud_gax::backoff_policy::BackoffPolicyArg>>(
+        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(())
     }

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 {