google_cloud_bigquery/read/
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
15//! Defines the retry policy for the BigQuery Storage Read API.
16
17use google_cloud_gax::error::Error;
18use google_cloud_gax::error::rpc::Code;
19use google_cloud_gax::retry_policy::RetryPolicy;
20use google_cloud_gax::retry_result::RetryResult;
21use google_cloud_gax::retry_state::RetryState;
22
23/// The default RPC retry/reconnect strategy recommended for the BigQuery
24/// Storage Read API.
25///
26/// This policy is used for when to reconnect a read stream, such as after
27/// transient network/transport errors, HTTP/2 stream resets, rate-limit errors
28/// (`ResourceExhausted`), and transient server statuses (`Unavailable`,
29/// `Internal`, `Aborted`, `Cancelled`, `DeadlineExceeded`). Reconnection is
30/// always assumed to be idempotent.
31///
32/// This policy must be decorated to limit the duration of the retry loop or
33/// the number of attempts.
34///
35/// # Example
36/// ```
37/// # use google_cloud_bigquery::client::Read;
38/// # use google_cloud_bigquery::read::retry_policy::RetryableErrors;
39/// # use google_cloud_gax::retry_policy::RetryPolicyExt;
40/// # async fn sample(client: Read) -> anyhow::Result<()> {
41/// let mut reader = client
42///     .read_rows()
43///     .set_read_stream("projects/my-project/locations/us/sessions/s1/streams/st1")
44///     .into_reader()
45///     .with_retry_policy(RetryableErrors.with_attempt_limit(5));
46/// # Ok(())
47/// # }
48/// ```
49#[derive(Clone, Debug, Default)]
50pub struct RetryableErrors;
51
52impl RetryPolicy for RetryableErrors {
53    fn on_error(&self, _state: &RetryState, error: Error) -> RetryResult {
54        if error.is_transient_and_before_rpc() {
55            return RetryResult::Continue(error);
56        }
57        if error.is_io() || error.is_timeout() || error.is_connect() {
58            return RetryResult::Continue(error);
59        }
60        if error.is_transport() && error.http_status_code().is_none() {
61            return RetryResult::Continue(error);
62        }
63        if let Some(429 | 500 | 502 | 503 | 504) = error.http_status_code() {
64            return RetryResult::Continue(error);
65        }
66        if let Some(status) = error.status() {
67            return match status.code {
68                Code::Aborted
69                | Code::Cancelled
70                | Code::DeadlineExceeded
71                | Code::Internal
72                | Code::ResourceExhausted
73                | Code::Unavailable => RetryResult::Continue(error),
74                _ => RetryResult::Permanent(error),
75            };
76        }
77        RetryResult::Permanent(error)
78    }
79}
80
81#[cfg(test)]
82mod tests {
83    use super::*;
84    use google_cloud_gax::error::CredentialsError;
85    use google_cloud_gax::error::rpc::Status;
86    use http::HeaderMap;
87    use test_case::test_case;
88
89    #[test_case(Code::Aborted)]
90    #[test_case(Code::Cancelled)]
91    #[test_case(Code::DeadlineExceeded)]
92    #[test_case(Code::Internal)]
93    #[test_case(Code::ResourceExhausted)]
94    #[test_case(Code::Unavailable)]
95    fn retryable_status_codes(code: Code) {
96        let err = || Error::service(Status::default().set_code(code).set_message("try again"));
97        assert!(matches!(
98            RetryableErrors.on_error(&RetryState::new(true), err()),
99            RetryResult::Continue(_)
100        ));
101        // Reads are always considered idempotent.
102        assert!(matches!(
103            RetryableErrors.on_error(&RetryState::new(false), err()),
104            RetryResult::Continue(_)
105        ));
106    }
107
108    #[test_case(Code::Unknown)]
109    #[test_case(Code::InvalidArgument)]
110    #[test_case(Code::NotFound)]
111    #[test_case(Code::PermissionDenied)]
112    #[test_case(Code::FailedPrecondition)]
113    #[test_case(Code::DataLoss)]
114    fn permanent_status_codes(code: Code) {
115        let err = Error::service(Status::default().set_code(code).set_message("fail"));
116        assert!(matches!(
117            RetryableErrors.on_error(&RetryState::new(true), err),
118            RetryResult::Permanent(_)
119        ));
120    }
121
122    #[test_case(429)]
123    #[test_case(500)]
124    #[test_case(502)]
125    #[test_case(503)]
126    #[test_case(504)]
127    fn retryable_http_status_codes(code: u16) {
128        let err = Error::http(code, HeaderMap::new(), bytes::Bytes::new());
129        assert!(matches!(
130            RetryableErrors.on_error(&RetryState::new(true), err),
131            RetryResult::Continue(_)
132        ));
133    }
134
135    #[test_case(400)]
136    #[test_case(401)]
137    #[test_case(403)]
138    #[test_case(404)]
139    #[test_case(408)]
140    #[test_case(409)]
141    #[test_case(501)]
142    fn permanent_http_status_codes(code: u16) {
143        let err = Error::http(code, HeaderMap::new(), bytes::Bytes::new());
144        assert!(matches!(
145            RetryableErrors.on_error(&RetryState::new(true), err),
146            RetryResult::Permanent(_)
147        ));
148    }
149
150    #[test]
151    fn retry_transient_before_rpc() {
152        let err = Error::authentication(CredentialsError::from_msg(true, "try again"));
153        assert!(matches!(
154            RetryableErrors.on_error(&RetryState::new(false), err),
155            RetryResult::Continue(_)
156        ));
157
158        let err = Error::authentication(CredentialsError::from_msg(false, "fail"));
159        assert!(matches!(
160            RetryableErrors.on_error(&RetryState::new(true), err),
161            RetryResult::Permanent(_)
162        ));
163    }
164
165    #[test]
166    fn retry_io_timeout_connect_transport() {
167        let state = RetryState::new(true);
168        assert!(
169            RetryableErrors
170                .on_error(&state, Error::io("io"))
171                .is_continue()
172        );
173        assert!(
174            RetryableErrors
175                .on_error(&state, Error::timeout("timeout"))
176                .is_continue()
177        );
178        assert!(
179            RetryableErrors
180                .on_error(&state, Error::connect("connect"))
181                .is_continue()
182        );
183        assert!(
184            RetryableErrors
185                .on_error(&state, Error::transport(HeaderMap::new(), "rst_stream"))
186                .is_continue()
187        );
188    }
189
190    #[test]
191    fn permanent_serde() {
192        let state = RetryState::new(true);
193        assert!(
194            RetryableErrors
195                .on_error(&state, Error::ser("fail"))
196                .is_permanent()
197        );
198        assert!(
199            RetryableErrors
200                .on_error(&state, Error::deser("fail"))
201                .is_permanent()
202        );
203    }
204}
google_cloud_bigquery/read/retry_policy.rs

google_cloud_bigquery/read/
retry_policy.rs
