google-cloud-bigquery 0.18.0

// 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 a retry policy for BigQuery.

use crate::error::QueryError;
use google_cloud_bigquery_v2::model::ErrorProto;
use google_cloud_gax::backoff_policy::BackoffPolicy;
use google_cloud_gax::backoff_policy::BackoffPolicyArg;
use google_cloud_gax::error::Error as GaxError;
use google_cloud_gax::error::rpc::{Code, StatusDetails};
use google_cloud_gax::exponential_backoff::ExponentialBackoffBuilder;
use google_cloud_gax::retry_policy::{RetryPolicy, RetryPolicyExt};
use google_cloud_gax::retry_result::RetryResult;
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].
///
/// ```
/// # 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_attempt_limit(6);
/// 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 struct RetryableErrors;

impl RetryPolicy for RetryableErrors {
    fn on_error(&self, state: &RetryState, error: GaxError) -> RetryResult {
        if error.is_transient_and_before_rpc() {
            return RetryResult::Continue(error);
        }
        if !state.idempotent {
            return RetryResult::Permanent(error);
        }
        if error.is_io() || error.is_timeout() {
            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
                | Code::Unknown => RetryResult::Continue(error),
                _ => RetryResult::Permanent(error),
            };
        }
        RetryResult::Permanent(error)
    }
}

pub(crate) fn default_retry_policy() -> Arc<dyn RetryPolicy> {
    Arc::new(RetryableErrors.with_attempt_limit(6))
}

pub(crate) fn default_backoff_policy() -> Arc<dyn BackoffPolicy> {
    Arc::new(
        ExponentialBackoffBuilder::default()
            .with_initial_delay(Duration::from_secs(1))
            .with_maximum_delay(Duration::from_secs(32))
            .with_scaling(2.0)
            .build()
            .expect("valid backoff configuration"),
    )
}

/// The result of evaluating a BigQuery job error against a [`JobRetryPolicy`].
#[derive(Debug)]
pub(crate) enum JobRetryResult {
    Continue(Duration, #[allow(dead_code)] QueryError),
    Exhausted(QueryError),
    Permanent(QueryError),
}

impl JobRetryResult {
    #[allow(dead_code)]
    pub(crate) fn is_continue(&self) -> bool {
        matches!(self, Self::Continue(_, _))
    }

    #[allow(dead_code)]
    pub(crate) fn is_exhausted(&self) -> bool {
        matches!(self, Self::Exhausted(_))
    }

    #[allow(dead_code)]
    pub(crate) fn is_permanent(&self) -> bool {
        matches!(self, Self::Permanent(_))
    }
}

/// A policy trait for handling BigQuery job execution retries.
///
/// Note: This trait is kept crate-internal for now. It is a copy/adaptation of GAX's
/// `RetryPolicy` trait specifically tailored for BigQuery job-level re-issuance.
/// In the future, we plan to refine this trait and discuss how to make it public
/// so customers can provide their own custom job retry policies.
pub(crate) trait JobRetryPolicy<S = RetryState>: Send + Sync + std::fmt::Debug {
    fn on_error(&self, state: &S, error: QueryError) -> JobRetryResult;
}

#[derive(Clone, Debug)]
pub(crate) struct RetryableJobErrors {
    attempt_limit: u32,
    backoff: Arc<dyn BackoffPolicy>,
}

impl Default for RetryableJobErrors {
    fn default() -> Self {
        Self {
            attempt_limit: 3,
            backoff: default_backoff_policy(),
        }
    }
}

impl RetryableJobErrors {
    #[allow(dead_code)]
    pub fn with_attempt_limit(mut self, attempt_limit: u32) -> Self {
        self.attempt_limit = attempt_limit;
        self
    }

    #[allow(dead_code)]
    pub fn with_backoff_policy<V: Into<BackoffPolicyArg>>(mut self, v: V) -> Self {
        self.backoff = v.into().into();
        self
    }
}

impl JobRetryPolicy for RetryableJobErrors {
    fn on_error(&self, state: &RetryState, error: QueryError) -> JobRetryResult {
        if !is_query_error_retryable(&error) {
            return JobRetryResult::Permanent(error);
        }
        if state.attempt_count >= self.attempt_limit {
            return JobRetryResult::Exhausted(error);
        }
        let delay = self.backoff.on_failure(state);
        JobRetryResult::Continue(delay, error)
    }
}

pub(crate) fn default_job_retry_policy() -> Arc<dyn JobRetryPolicy> {
    Arc::new(RetryableJobErrors::default())
}

pub(crate) fn is_query_error_retryable(err: &QueryError) -> bool {
    match err {
        QueryError::JobFailed { errors } => is_retryable_errors(errors),
        QueryError::Rpc { source } => is_rpc_error_retryable(source),
        _ => false,
    }
}

pub(crate) fn is_rpc_error_retryable(error: &GaxError) -> bool {
    if let Some(status) = error.status() {
        for detail in &status.details {
            if let StatusDetails::ErrorInfo(info) = detail
                && is_retryable_error_reason(&info.reason)
            {
                return true;
            }
        }
    }
    false
}

pub(crate) fn is_retryable_errors(errors: &[ErrorProto]) -> bool {
    !errors.is_empty() && errors.iter().all(|e| is_retryable_error_reason(&e.reason))
}

pub(crate) fn is_retryable_error_reason(reason: &str) -> bool {
    matches!(
        reason,
        "backendError"
            | "jobBackendError"
            | "rateLimitExceeded"
            | "jobRateLimitExceeded"
            | "internalError"
            | "jobInternalError"
    )
}

/// Returns true if `error` reports a conflict with a resource that already
/// exists, such as `409 Already Exists: Job my-project:US.job_1234567890`.
pub(crate) fn is_duplicate_job_error(error: &QueryError) -> bool {
    let QueryError::Rpc { source } = error else {
        return false;
    };
    source.http_status_code() == Some(409)
        || source
            .status()
            .is_some_and(|s| s.code == Code::AlreadyExists)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::query::tests::create_test_backoff_policy;
    use google_cloud_bigquery_v2::model::ErrorProto;
    use google_cloud_gax::error::CredentialsError;
    use google_cloud_gax::error::rpc::{Code, Status};
    use google_cloud_gax::retry_state::RetryState;
    use google_cloud_rpc::model::ErrorInfo;
    use http::HeaderMap;
    use test_case::test_case;

    #[test_case("backendError", true)]
    #[test_case("jobBackendError", true)]
    #[test_case("rateLimitExceeded", true)]
    #[test_case("jobRateLimitExceeded", true)]
    #[test_case("internalError", true)]
    #[test_case("jobInternalError", true)]
    #[test_case("invalidQuery", false)]
    #[test_case("notFound", false)]
    fn test_is_retryable_error_reason(reason: &str, expected: bool) {
        assert_eq!(is_retryable_error_reason(reason), expected);
    }

    #[test]
    fn test_is_retryable_errors() {
        assert!(!is_retryable_errors(&[]));

        let non_retryable = vec![ErrorProto::new().set_reason("invalidQuery")];
        assert!(!is_retryable_errors(&non_retryable));

        let mixed = vec![
            ErrorProto::new().set_reason("invalidQuery"),
            ErrorProto::new().set_reason("backendError"),
        ];
        assert!(!is_retryable_errors(&mixed));

        let retryable = vec![
            ErrorProto::new().set_reason("rateLimitExceeded"),
            ErrorProto::new().set_reason("backendError"),
        ];
        assert!(is_retryable_errors(&retryable));
    }

    #[test]
    fn test_retryable_errors_on_error() {
        let p = RetryableErrors;
        let idempotent = RetryState::new(true);
        let non_idempotent = RetryState::new(false);

        let retryable_codes = [
            Code::Aborted,
            Code::DeadlineExceeded,
            Code::Internal,
            Code::ResourceExhausted,
            Code::Unavailable,
            Code::Unknown,
        ];
        for code in retryable_codes {
            let err = || GaxError::service(Status::default().set_code(code));
            assert!(p.on_error(&idempotent, err()).is_continue(), "{code:?}");
            assert!(
                p.on_error(&non_idempotent, err()).is_permanent(),
                "{code:?}"
            );
        }

        let permanent_codes = [
            Code::NotFound,
            Code::PermissionDenied,
            Code::InvalidArgument,
        ];
        for code in permanent_codes {
            let err = || GaxError::service(Status::default().set_code(code));
            assert!(p.on_error(&idempotent, err()).is_permanent(), "{code:?}");
            assert!(
                p.on_error(&non_idempotent, err()).is_permanent(),
                "{code:?}"
            );
        }

        let retryable_http = [429, 500, 502, 503, 504];
        for code in retryable_http {
            let err = || GaxError::http(code, HeaderMap::new(), bytes::Bytes::new());
            assert!(p.on_error(&idempotent, err()).is_continue(), "HTTP {code}");
            assert!(
                p.on_error(&non_idempotent, err()).is_permanent(),
                "HTTP {code}"
            );
        }

        let permanent_http = [400, 404, 408, 409, 501];
        for code in permanent_http {
            let err = || GaxError::http(code, HeaderMap::new(), bytes::Bytes::new());
            assert!(p.on_error(&idempotent, err()).is_permanent(), "HTTP {code}");
            assert!(
                p.on_error(&non_idempotent, err()).is_permanent(),
                "HTTP {code}"
            );
        }

        let io = || GaxError::io("connection reset");
        assert!(p.on_error(&idempotent, io()).is_continue());
        assert!(p.on_error(&non_idempotent, io()).is_permanent());

        let timeout = || GaxError::timeout("deadline");
        assert!(p.on_error(&idempotent, timeout()).is_continue());
        assert!(p.on_error(&non_idempotent, timeout()).is_permanent());

        // Failures before the request is sent cannot have reached the service,
        // so they are retried either way.
        let before_rpc =
            || GaxError::authentication(CredentialsError::from_msg(true, "token refresh failed"));
        assert!(p.on_error(&idempotent, before_rpc()).is_continue());
        assert!(p.on_error(&non_idempotent, before_rpc()).is_continue());
    }

    #[test]
    fn test_is_duplicate_job_error() {
        let rpc = |source| QueryError::Rpc { source };
        let status = |code| GaxError::service(Status::default().set_code(code));
        let http = |code| GaxError::http(code, HeaderMap::new(), bytes::Bytes::new());

        assert!(is_duplicate_job_error(&rpc(http(409))));
        assert!(is_duplicate_job_error(&rpc(status(Code::AlreadyExists))));

        assert!(!is_duplicate_job_error(&rpc(status(Code::Aborted))));
        assert!(!is_duplicate_job_error(&rpc(http(500))));
        assert!(!is_duplicate_job_error(&QueryError::JobFailed {
            errors: vec![ErrorProto::new().set_reason("duplicate")],
        }));
    }

    #[test]
    fn test_duplicate_job_error_is_not_retryable() {
        const BQ_DUPLICATE_PAYLOAD: &[u8] = br#"{
  "error": {
    "code": 409,
    "message": "Already Exists: Job my-project:US.job_1234567890",
    "errors": [
      {
        "message": "Already Exists: Job my-project:US.job_1234567890",
        "domain": "global",
        "reason": "duplicate"
      }
    ],
    "status": "ALREADY_EXISTS"
  }
}"#;
        let status = Status::try_from(&bytes::Bytes::from_static(BQ_DUPLICATE_PAYLOAD))
            .expect("should deserialize BigQuery REST error");
        let err = QueryError::Rpc {
            source: GaxError::service(status),
        };
        assert!(is_duplicate_job_error(&err), "{err:?}");
        assert!(!is_query_error_retryable(&err), "{err:?}");
    }

    #[test]
    fn test_is_rpc_error_retryable() {
        use google_cloud_rpc::model::ErrorInfo;

        // BigQuery REST API JSON error with backendError
        const BQ_REST_PAYLOAD: &[u8] = br#"{
  "error": {
    "code": 400,
    "message": "The job encountered an error during execution. Retrying the job may solve the problem.",
    "errors": [
      {
        "message": "The job encountered an error during execution. Retrying the job may solve the problem.",
        "domain": "global",
        "reason": "backendError"
      }
    ],
    "status": "INVALID_ARGUMENT"
  }
}"#;
        let status = Status::try_from(&bytes::Bytes::from_static(BQ_REST_PAYLOAD))
            .expect("should deserialize BigQuery REST error");
        let err = GaxError::service(status);
        assert!(is_rpc_error_retryable(&err));

        // ErrorInfo detail with retryable reason
        let status = Status::default()
            .set_code(Code::InvalidArgument)
            .set_message("Error occurred")
            .set_details(vec![StatusDetails::ErrorInfo(
                ErrorInfo::new().set_reason("backendError"),
            )]);
        let err = GaxError::service(status);
        assert!(is_rpc_error_retryable(&err));

        // ErrorInfo detail with non-retryable reason
        let status = Status::default()
            .set_code(Code::InvalidArgument)
            .set_message("Error occurred")
            .set_details(vec![StatusDetails::ErrorInfo(
                ErrorInfo::new().set_reason("invalidQuery"),
            )]);
        let err = GaxError::service(status);
        assert!(!is_rpc_error_retryable(&err));

        // Status without ErrorInfo details
        let status = Status::default()
            .set_code(Code::InvalidArgument)
            .set_message("Syntax error: Unexpected identifier");
        let err = GaxError::service(status);
        assert!(!is_rpc_error_retryable(&err));
    }

    #[test]
    fn test_job_retryable_errors() {
        let policy = RetryableJobErrors::default();
        let state = RetryState::default();

        let retryable_err = QueryError::JobFailed {
            errors: vec![ErrorProto::new().set_reason("backendError")],
        };
        assert!(policy.on_error(&state, retryable_err).is_continue());

        let permanent_err = QueryError::JobFailed {
            errors: vec![ErrorProto::new().set_reason("invalidQuery")],
        };
        assert!(policy.on_error(&state, permanent_err).is_permanent());

        let rpc_retryable_err = QueryError::Rpc {
            source: GaxError::service(
                Status::default()
                    .set_code(Code::InvalidArgument)
                    .set_details(vec![StatusDetails::ErrorInfo(
                        ErrorInfo::new().set_reason("backendError"),
                    )]),
            ),
        };
        assert!(policy.on_error(&state, rpc_retryable_err).is_continue());

        let rpc_permanent_err = QueryError::Rpc {
            source: GaxError::service(
                Status::default()
                    .set_code(Code::InvalidArgument)
                    .set_message("Syntax error: Unexpected identifier"),
            ),
        };
        assert!(policy.on_error(&state, rpc_permanent_err).is_permanent());
    }

    #[test]
    fn test_job_attempt_limit() {
        let policy = default_job_retry_policy(); // default attempt_limit is 3
        let retryable_err = || QueryError::JobFailed {
            errors: vec![ErrorProto::new().set_reason("backendError")],
        };

        let mut state = RetryState::default();
        state.attempt_count = 1;
        assert!(policy.on_error(&state, retryable_err()).is_continue());

        state.attempt_count = 2;
        assert!(policy.on_error(&state, retryable_err()).is_continue());

        state.attempt_count = 3;
        assert!(policy.on_error(&state, retryable_err()).is_exhausted());
    }

    #[test]
    fn test_job_backoff_policy() {
        let mut backoff = create_test_backoff_policy();
        backoff
            .expect_on_failure()
            .return_const(Duration::from_secs(5));

        let policy = RetryableJobErrors::default().with_backoff_policy(backoff);
        let retryable_err = QueryError::JobFailed {
            errors: vec![ErrorProto::new().set_reason("backendError")],
        };

        let state = RetryState::default();
        if let JobRetryResult::Continue(delay, _) = policy.on_error(&state, retryable_err) {
            assert_eq!(delay, Duration::from_secs(5));
        } else {
            panic!("expected Continue with 5s delay");
        }
    }

    #[test]
    fn test_default_retry_policy_is_bounded() {
        let policy = default_retry_policy();
        let err = || GaxError::service(Status::default().set_code(Code::Unavailable));

        for attempt in 1u32..6u32 {
            let state = RetryState::new(true).set_attempt_count(attempt);
            assert!(
                policy.on_error(&state, err()).is_continue(),
                "attempt {attempt} should continue"
            );
            assert_eq!(policy.remaining_time(&state), None);
        }

        let state = RetryState::new(true).set_attempt_count(6u32);
        assert!(
            policy.on_error(&state, err()).is_exhausted(),
            "attempt 6 should be exhausted"
        );
    }
}