google_cloud_bigquery/write/
retry_policy.rs

1// Copyright 2026 Google LLC
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     https://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use google_cloud_gax::backoff_policy::BackoffPolicy;
16use google_cloud_gax::error::Error;
17use google_cloud_gax::error::rpc::Code;
18use google_cloud_gax::exponential_backoff::ExponentialBackoff;
19use google_cloud_gax::retry_policy::{RetryPolicy, RetryPolicyExt};
20use google_cloud_gax::retry_result::RetryResult;
21use google_cloud_gax::retry_state::RetryState;
22use std::sync::Arc;
23use std::time::Duration;
24
25/// Configuration options for retrying writes.
26#[derive(Clone, Debug)]
27pub(crate) struct RetryOptions {
28    pub(crate) retry_policy: Arc<dyn RetryPolicy>,
29    pub(crate) backoff_policy: Arc<dyn BackoffPolicy>,
30    pub(crate) attempt_timeout: Option<Duration>,
31}
32
33impl Default for RetryOptions {
34    fn default() -> Self {
35        Self {
36            retry_policy: Arc::new(RetryableErrors.with_time_limit(Duration::from_secs(300))),
37            backoff_policy: Arc::new(ExponentialBackoff::default()),
38            attempt_timeout: None,
39        }
40    }
41}
42
43/// Follows the RPC retry strategy recommended for BigQuery Storage Write API.
44///
45/// This policy must be decorated to limit the duration of the retry loop or
46/// the number of attempts.
47#[derive(Clone, Debug)]
48pub struct RetryableErrors;
49
50impl RetryPolicy for RetryableErrors {
51    fn on_error(&self, _state: &RetryState, error: Error) -> RetryResult {
52        if error.is_transient_and_before_rpc()
53            || error.is_io()
54            || error.is_timeout()
55            || error.is_connect()
56        {
57            return RetryResult::Continue(error);
58        }
59        if error.is_transport() && error.http_status_code().is_none() {
60            return RetryResult::Continue(error);
61        }
62        if let Some(429 | 500 | 502 | 503 | 504) = error.http_status_code() {
63            return RetryResult::Continue(error);
64        }
65        if let Some(status) = error.status() {
66            return match status.code {
67                Code::Aborted
68                | Code::Cancelled
69                | Code::DeadlineExceeded
70                | Code::Internal
71                | Code::ResourceExhausted
72                | Code::Unavailable => RetryResult::Continue(error),
73                _ => RetryResult::Permanent(error),
74            };
75        }
76        RetryResult::Permanent(error)
77    }
78}
79
80#[cfg(test)]
81mod tests {
82    use super::*;
83    use google_cloud_gax::error::CredentialsError;
84    use google_cloud_gax::error::rpc::Status;
85    use http::HeaderMap;
86    use test_case::test_case;
87
88    #[test]
89    fn default_retry_options() {
90        let options = RetryOptions::default();
91        assert_eq!(options.attempt_timeout, None);
92        let fmt = format!("{options:?}");
93        assert!(fmt.contains("RetryableErrors"), "{fmt}");
94        assert!(fmt.contains("ExponentialBackoff"), "{fmt}");
95    }
96
97    #[test_case(Code::Aborted)]
98    #[test_case(Code::Cancelled)]
99    #[test_case(Code::DeadlineExceeded)]
100    #[test_case(Code::Internal)]
101    #[test_case(Code::ResourceExhausted)]
102    #[test_case(Code::Unavailable)]
103    fn retryable_status_codes(code: Code) {
104        let err = Error::service(Status::default().set_code(code).set_message("try again"));
105        assert!(matches!(
106            RetryableErrors.on_error(&RetryState::default(), err),
107            RetryResult::Continue(_)
108        ));
109    }
110
111    #[test_case(Code::Unknown)]
112    #[test_case(Code::InvalidArgument)]
113    #[test_case(Code::NotFound)]
114    #[test_case(Code::PermissionDenied)]
115    #[test_case(Code::FailedPrecondition)]
116    #[test_case(Code::DataLoss)]
117    fn permanent_status_codes(code: Code) {
118        let err = Error::service(Status::default().set_code(code).set_message("fail"));
119        assert!(matches!(
120            RetryableErrors.on_error(&RetryState::default(), err),
121            RetryResult::Permanent(_)
122        ));
123    }
124
125    #[test_case(429)]
126    #[test_case(500)]
127    #[test_case(502)]
128    #[test_case(503)]
129    #[test_case(504)]
130    fn retryable_http_status_codes(code: u16) {
131        let err = Error::http(code, HeaderMap::new(), bytes::Bytes::new());
132        assert!(matches!(
133            RetryableErrors.on_error(&RetryState::default(), err),
134            RetryResult::Continue(_)
135        ));
136    }
137
138    #[test_case(400)]
139    #[test_case(401)]
140    #[test_case(403)]
141    #[test_case(404)]
142    #[test_case(408)]
143    #[test_case(409)]
144    #[test_case(501)]
145    fn permanent_http_status_codes(code: u16) {
146        let err = Error::http(code, HeaderMap::new(), bytes::Bytes::new());
147        assert!(matches!(
148            RetryableErrors.on_error(&RetryState::default(), err),
149            RetryResult::Permanent(_)
150        ));
151    }
152
153    #[test]
154    fn retry_transient_before_rpc() {
155        let err = Error::authentication(CredentialsError::from_msg(true, "try again"));
156        assert!(matches!(
157            RetryableErrors.on_error(&RetryState::default(), err),
158            RetryResult::Continue(_)
159        ));
160
161        let err = Error::authentication(CredentialsError::from_msg(false, "fail"));
162        assert!(matches!(
163            RetryableErrors.on_error(&RetryState::default(), err),
164            RetryResult::Permanent(_)
165        ));
166    }
167
168    #[test]
169    fn retry_io() {
170        let err = Error::io("try again");
171        assert!(matches!(
172            RetryableErrors.on_error(&RetryState::default(), err),
173            RetryResult::Continue(_)
174        ));
175    }
176
177    #[test]
178    fn retry_timeout() {
179        let err = Error::timeout("deadline exceeded");
180        assert!(matches!(
181            RetryableErrors.on_error(&RetryState::default(), err),
182            RetryResult::Continue(_)
183        ));
184    }
185
186    #[test]
187    fn retry_connect() {
188        let err = Error::connect("connection failed");
189        assert!(matches!(
190            RetryableErrors.on_error(&RetryState::default(), err),
191            RetryResult::Continue(_)
192        ));
193    }
194
195    #[test]
196    fn retry_transport_without_http_status() {
197        let err = Error::transport(HeaderMap::new(), "connection reset");
198        assert!(matches!(
199            RetryableErrors.on_error(&RetryState::default(), err),
200            RetryResult::Continue(_)
201        ));
202    }
203
204    #[test]
205    fn permanent_serde() {
206        let err = Error::ser("fail");
207        assert!(matches!(
208            RetryableErrors.on_error(&RetryState::default(), err),
209            RetryResult::Permanent(_)
210        ));
211
212        let err = Error::deser("fail");
213        assert!(matches!(
214            RetryableErrors.on_error(&RetryState::default(), err),
215            RetryResult::Permanent(_)
216        ));
217    }
218}
google_cloud_bigquery/write/retry_policy.rs

google_cloud_bigquery/write/
retry_policy.rs
