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(60))),
37            backoff_policy: Arc::new(ExponentialBackoff::default()),
38            attempt_timeout: Some(Duration::from_secs(30)),
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_case(Code::Aborted)]
89    #[test_case(Code::Cancelled)]
90    #[test_case(Code::DeadlineExceeded)]
91    #[test_case(Code::Internal)]
92    #[test_case(Code::ResourceExhausted)]
93    #[test_case(Code::Unavailable)]
94    fn retryable_status_codes(code: Code) {
95        let err = Error::service(Status::default().set_code(code).set_message("try again"));
96        assert!(matches!(
97            RetryableErrors.on_error(&RetryState::default(), err),
98            RetryResult::Continue(_)
99        ));
100    }
101
102    #[test_case(Code::Unknown)]
103    #[test_case(Code::InvalidArgument)]
104    #[test_case(Code::NotFound)]
105    #[test_case(Code::PermissionDenied)]
106    #[test_case(Code::FailedPrecondition)]
107    #[test_case(Code::DataLoss)]
108    fn permanent_status_codes(code: Code) {
109        let err = Error::service(Status::default().set_code(code).set_message("fail"));
110        assert!(matches!(
111            RetryableErrors.on_error(&RetryState::default(), err),
112            RetryResult::Permanent(_)
113        ));
114    }
115
116    #[test_case(429)]
117    #[test_case(500)]
118    #[test_case(502)]
119    #[test_case(503)]
120    #[test_case(504)]
121    fn retryable_http_status_codes(code: u16) {
122        let err = Error::http(code, HeaderMap::new(), bytes::Bytes::new());
123        assert!(matches!(
124            RetryableErrors.on_error(&RetryState::default(), err),
125            RetryResult::Continue(_)
126        ));
127    }
128
129    #[test_case(400)]
130    #[test_case(401)]
131    #[test_case(403)]
132    #[test_case(404)]
133    #[test_case(408)]
134    #[test_case(409)]
135    #[test_case(501)]
136    fn permanent_http_status_codes(code: u16) {
137        let err = Error::http(code, HeaderMap::new(), bytes::Bytes::new());
138        assert!(matches!(
139            RetryableErrors.on_error(&RetryState::default(), err),
140            RetryResult::Permanent(_)
141        ));
142    }
143
144    #[test]
145    fn retry_transient_before_rpc() {
146        let err = Error::authentication(CredentialsError::from_msg(true, "try again"));
147        assert!(matches!(
148            RetryableErrors.on_error(&RetryState::default(), err),
149            RetryResult::Continue(_)
150        ));
151
152        let err = Error::authentication(CredentialsError::from_msg(false, "fail"));
153        assert!(matches!(
154            RetryableErrors.on_error(&RetryState::default(), err),
155            RetryResult::Permanent(_)
156        ));
157    }
158
159    #[test]
160    fn retry_io() {
161        let err = Error::io("try again");
162        assert!(matches!(
163            RetryableErrors.on_error(&RetryState::default(), err),
164            RetryResult::Continue(_)
165        ));
166    }
167
168    #[test]
169    fn retry_timeout() {
170        let err = Error::timeout("deadline exceeded");
171        assert!(matches!(
172            RetryableErrors.on_error(&RetryState::default(), err),
173            RetryResult::Continue(_)
174        ));
175    }
176
177    #[test]
178    fn retry_connect() {
179        let err = Error::connect("connection failed");
180        assert!(matches!(
181            RetryableErrors.on_error(&RetryState::default(), err),
182            RetryResult::Continue(_)
183        ));
184    }
185
186    #[test]
187    fn retry_transport_without_http_status() {
188        let err = Error::transport(HeaderMap::new(), "connection reset");
189        assert!(matches!(
190            RetryableErrors.on_error(&RetryState::default(), err),
191            RetryResult::Continue(_)
192        ));
193    }
194
195    #[test]
196    fn permanent_serde() {
197        let err = Error::ser("fail");
198        assert!(matches!(
199            RetryableErrors.on_error(&RetryState::default(), err),
200            RetryResult::Permanent(_)
201        ));
202
203        let err = Error::deser("fail");
204        assert!(matches!(
205            RetryableErrors.on_error(&RetryState::default(), err),
206            RetryResult::Permanent(_)
207        ));
208    }
209}
google_cloud_bigquery/write/retry_policy.rs

google_cloud_bigquery/write/
retry_policy.rs
