google_cloud_bigquery_v2/
operation.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 crate::builder::job_service::InsertJob;
16use crate::model::Job;
17use google_cloud_gax::backoff_policy::BackoffPolicy;
18use google_cloud_gax::error::Error as GaxError;
19use google_cloud_gax::error::rpc::{Code, Status};
20use google_cloud_gax::exponential_backoff::ExponentialBackoff;
21use google_cloud_gax::retry_state::RetryState;
22use google_cloud_lro::Poller;
23
24impl google_cloud_lro::internal::DiscoveryOperation for Job {
25    fn name(&self) -> Option<&String> {
26        self.job_reference.as_ref().map(|r| &r.job_id)
27    }
28
29    fn done(&self) -> bool {
30        self.status
31            .as_ref()
32            .map(|s| s.state == "DONE")
33            .unwrap_or(false)
34    }
35
36    fn error(&self) -> Option<Status> {
37        self.status.as_ref().and_then(|s| {
38            s.error_result.as_ref().map(|e| {
39                Status::default()
40                    .set_code(Code::Unknown)
41                    .set_message(e.message.clone())
42            })
43        })
44    }
45}
46
47/// Determines if a BigQuery job failure reason is transient and eligible for
48/// job-level retry.
49///
50/// Returns `true` for retryable reasons (`jobBackendError`,
51/// `jobInternalError`, `jobRateLimitExceeded`, `tableUnavailable`) per
52/// BigQuery error handling specification.
53#[allow(dead_code)]
54pub(crate) fn is_retryable_job_error(reason: &str) -> bool {
55    matches!(
56        reason,
57        "jobBackendError" | "jobInternalError" | "jobRateLimitExceeded" | "tableUnavailable"
58    )
59}
60
61/// Prepares a `Job` instance for retry by assigning a new synthetic job ID
62/// and clearing existing execution status.
63///
64/// To preserve idempotency and avoid job execution collisions, each job-level
65/// retry must use a unique job ID while retaining original reference details
66/// (project ID, location) and configuration settings.
67#[allow(dead_code)]
68pub(crate) fn prepare_job_for_retry(mut job: Job) -> Job {
69    job.job_reference.get_or_insert_default().job_id = uuid::Uuid::new_v4().to_string();
70    job.status = None;
71    job
72}
73
74/// Configuration policy for BigQuery job-level retries.
75#[derive(Debug)]
76pub(crate) struct JobRetryPolicy {
77    /// Maximum number of general job-level attempts for retryable job errors.
78    pub job_level_attempt_limit: u32,
79    /// Backoff strategy between retry attempts.
80    pub backoff: ExponentialBackoff,
81}
82
83impl Default for JobRetryPolicy {
84    fn default() -> Self {
85        Self {
86            job_level_attempt_limit: 3,
87            backoff: ExponentialBackoff::default(),
88        }
89    }
90}
91
92/// Errors returned by the JobPoller.
93#[derive(Debug, thiserror::Error)]
94pub enum JobPollerError {
95    /// An error occurred during the RPC or LRO polling.
96    #[error(transparent)]
97    Rpc(#[from] GaxError),
98    /// The job completed, but the BigQuery service reported an internal error.
99    #[error("BigQuery job failed ({}): {}", .0.reason, .0.message)]
100    ErrorProto(crate::model::ErrorProto),
101}
102
103/// A poller that monitors the status of an inserted BigQuery job and handles retries.
104#[derive(Debug)]
105pub struct JobPoller {
106    policy: JobRetryPolicy,
107    // Because the builder holds a `Job` which is >7kB, we ought to store it
108    // on the heap.
109    //
110    // This is important across `await` points where stack variables
111    // are captured in the async fn's state machine. This bloats the size of
112    // the returned `Future`, which can potentially overflow the stack.
113    //
114    // See #6391.
115    builder: Box<InsertJob>,
116}
117
118impl JobPoller {
119    pub(crate) fn new(builder: InsertJob) -> Self {
120        Self {
121            policy: JobRetryPolicy::default(),
122            builder: Box::new(builder),
123        }
124    }
125
126    /// Sets the maximum number of job-level attempts.
127    pub fn with_attempt_limit(mut self, limit: u32) -> Self {
128        self.policy.job_level_attempt_limit = limit;
129        self
130    }
131
132    /// Sets the exponential backoff policy for job-level retries.
133    pub fn with_job_retry_backoff(mut self, backoff: ExponentialBackoff) -> Self {
134        self.policy.backoff = backoff;
135        self
136    }
137
138    /// Polls the job until it is done, returning the final Job status.
139    pub async fn until_done(self) -> Result<Job, JobPollerError> {
140        let mut attempts = 0_u32;
141        let mut builder = self.builder;
142        let backoff = self.policy.backoff;
143        let start_time = std::time::Instant::now();
144
145        loop {
146            // NOTE: the client library intercepts errors and retries internally
147            // according to the policies set on `builder`.
148            {
149                let poller = (*builder).clone().poller();
150                let job = poller.until_done().await?;
151                let Some(status) = &job.status else {
152                    return Ok(job);
153                };
154                let Some(err) = &status.error_result else {
155                    return Ok(job);
156                };
157
158                attempts += 1;
159                if !is_retryable_job_error(&err.reason)
160                    || attempts >= self.policy.job_level_attempt_limit
161                {
162                    return Err(JobPollerError::ErrorProto(err.clone()));
163                }
164
165                let job = prepare_job_for_retry(job);
166                *builder = (*builder).set_job(job);
167
168                // We use a block so that `job` (~7kB) is not allocated on the
169                // stack across the sleep `await` point.
170            }
171
172            let retry_state = RetryState::new(true)
173                .set_start(start_time)
174                .set_attempt_count(attempts);
175            let delay = backoff.on_failure(&retry_state);
176            tokio::time::sleep(delay).await;
177        }
178    }
179}
180
181impl InsertJob {
182    /// Returns a `JobPoller`, which can retry on [job-level errors].
183    ///
184    /// If the job fails with an internal error, the `JobPoller` will retry the
185    /// `InsertJob` operation. Note that the client library will supply a
186    /// synthetic job ID for any retries.
187    ///
188    /// ```no_run
189    /// # async fn example(builder: google_cloud_bigquery_v2::builder::job_service::InsertJob) -> Result<(), Box<dyn std::error::Error>> {
190    /// let job = builder.into_job_poller().until_done().await?;
191    /// # Ok(())
192    /// # }
193    /// ```
194    ///
195    /// [job-level errors]: https://docs.cloud.google.com/bigquery/docs/error-messages#errortable
196    pub fn into_job_poller(self) -> JobPoller {
197        JobPoller::new(self)
198    }
199}
200
201#[cfg(test)]
202mod tests {
203    use super::*;
204    use crate::model::{
205        ErrorProto, JobConfiguration, JobConfigurationQuery, JobReference, JobStatus,
206    };
207    use google_cloud_lro::internal::DiscoveryOperation;
208
209    #[test]
210    fn name_none() {
211        let job = Job::default();
212        assert_eq!(job.name(), None);
213    }
214
215    #[test]
216    fn name_some() {
217        let job = Job::new().set_job_reference(JobReference::new().set_job_id("test-id"));
218        assert_eq!(job.name().map(|s| s.as_str()), Some("test-id"));
219    }
220
221    #[test]
222    fn done_none() {
223        let job = Job::default();
224        assert!(!job.done());
225    }
226
227    #[test]
228    fn done_false() {
229        let job = Job::new().set_status(JobStatus::new().set_state("RUNNING"));
230        assert!(!job.done());
231    }
232
233    #[test]
234    fn done_true() {
235        let job = Job::new().set_status(JobStatus::new().set_state("DONE"));
236        assert!(job.done());
237    }
238
239    #[test]
240    fn error_none() {
241        let job = Job::default();
242        assert!(job.error().is_none());
243
244        let job_no_error = Job::new().set_status(JobStatus::new().set_state("DONE"));
245        assert!(job_no_error.error().is_none());
246    }
247
248    #[test]
249    fn error_some() {
250        let job = Job::new()
251            .set_status(JobStatus::new().set_error_result(ErrorProto::new().set_message("failed")));
252        let err = job.error().expect("should have error");
253        assert_eq!(err.code, Code::Unknown);
254        assert_eq!(err.message, "failed");
255    }
256
257    #[test]
258    fn retryable_job_errors() {
259        assert!(is_retryable_job_error("jobBackendError"));
260        assert!(is_retryable_job_error("jobInternalError"));
261        assert!(is_retryable_job_error("jobRateLimitExceeded"));
262        assert!(is_retryable_job_error("tableUnavailable"));
263
264        assert!(!is_retryable_job_error("invalidQuery"));
265        assert!(!is_retryable_job_error("accessDenied"));
266        assert!(!is_retryable_job_error("notFound"));
267        assert!(!is_retryable_job_error("backendError"));
268        assert!(!is_retryable_job_error(""));
269    }
270
271    #[test]
272    fn job_retry_policy_defaults() {
273        let policy = JobRetryPolicy::default();
274        assert_eq!(policy.job_level_attempt_limit, 3);
275    }
276
277    #[test]
278    fn prepare_job_for_retry_generates_new_id_and_resets_status() {
279        let original_job = Job::new()
280            .set_job_reference(
281                JobReference::new()
282                    .set_project_id("test-project")
283                    .set_job_id("original-job-id")
284                    .set_location("US"),
285            )
286            .set_status(
287                JobStatus::new().set_state("DONE").set_error_result(
288                    ErrorProto::new()
289                        .set_reason("jobBackendError")
290                        .set_message("backend failed"),
291                ),
292            );
293
294        let retried_job = prepare_job_for_retry(original_job);
295
296        assert!(retried_job.status.is_none());
297
298        let ref_data = retried_job
299            .job_reference
300            .expect("should have job reference");
301        assert_eq!(ref_data.project_id, "test-project");
302        assert_eq!(ref_data.location.as_deref(), Some("US"));
303        assert_ne!(ref_data.job_id, "original-job-id");
304        assert!(uuid::Uuid::parse_str(&ref_data.job_id).is_ok());
305    }
306
307    #[test]
308    fn prepare_job_for_retry_handles_none_job_reference() {
309        let original_job = Job::new().set_status(JobStatus::new().set_state("DONE"));
310
311        let retried_job = prepare_job_for_retry(original_job);
312        assert!(retried_job.status.is_none());
313
314        let ref_data = retried_job
315            .job_reference
316            .expect("should create job reference when missing");
317        assert!(uuid::Uuid::parse_str(&ref_data.job_id).is_ok());
318    }
319
320    #[test]
321    fn prepare_job_for_retry_preserves_job_configuration_and_metadata() {
322        let original_job = Job::new()
323            .set_job_reference(
324                JobReference::new()
325                    .set_project_id("my-project")
326                    .set_job_id("initial-id")
327                    .set_location("EU"),
328            )
329            .set_configuration(
330                JobConfiguration::new()
331                    .set_query(JobConfigurationQuery::new().set_query("SELECT 42"))
332                    .set_labels([("env".to_string(), "test".to_string())]),
333            )
334            .set_user_email("user@example.com")
335            .set_status(
336                JobStatus::new().set_state("DONE").set_error_result(
337                    ErrorProto::new()
338                        .set_reason("jobInternalError")
339                        .set_message("internal error"),
340                ),
341            );
342
343        let retried = prepare_job_for_retry(original_job);
344
345        // Status must be reset to None for retry submission
346        assert!(retried.status.is_none());
347
348        // Configuration and user_email must be preserved
349        assert_eq!(
350            retried
351                .configuration
352                .as_ref()
353                .and_then(|c| c.query.as_ref())
354                .map(|q| q.query.as_str()),
355            Some("SELECT 42")
356        );
357        assert_eq!(
358            retried
359                .configuration
360                .as_ref()
361                .and_then(|c| c.labels.get("env").map(|s| s.as_str())),
362            Some("test")
363        );
364        assert_eq!(retried.user_email.as_str(), "user@example.com");
365
366        // JobReference metadata preserved, but job_id replaced with a new valid UUID
367        let ref_data = retried.job_reference.expect("must have reference");
368        assert_eq!(ref_data.project_id, "my-project");
369        assert_eq!(ref_data.location.as_deref(), Some("EU"));
370        assert_ne!(ref_data.job_id, "initial-id");
371        assert!(uuid::Uuid::parse_str(&ref_data.job_id).is_ok());
372    }
373
374    #[test]
375    fn custom_retry_policy_builder() {
376        let mut policy = JobRetryPolicy::default();
377        assert_eq!(policy.job_level_attempt_limit, 3);
378
379        policy.job_level_attempt_limit = 5;
380        assert_eq!(policy.job_level_attempt_limit, 5);
381    }
382}
google_cloud_bigquery_v2/operation.rs

google_cloud_bigquery_v2/
operation.rs
