diff --git a/src/bigquery/src/query/execution.rs b/src/bigquery/src/query/execution.rs index 701e033b32..38e31ea51c 100644 --- a/src/bigquery/src/query/execution.rs +++ b/src/bigquery/src/query/execution.rs @@ -16,7 +16,8 @@ use crate::error::QueryError; use crate::query::builder::{ QUERY_REQUEST_ID_PREFIX, Query, generate_job_reference, generate_prefixed_id, }; -use crate::query::retry_policy::JobRetryResult; +use crate::query::query_handle::build_get_job; +use crate::query::retry_policy::{JobRetryResult, is_duplicate_job_error}; use crate::query::{Query as QueryHandle, Result}; use google_cloud_bigquery_v2::client::JobService; use google_cloud_bigquery_v2::model::{ @@ -94,15 +95,7 @@ impl InsertJobExecutor { .send() .await?; - let job_status = res.status.as_ref(); - if let Some(status) = job_status - && status.error_result.is_some() - { - let errors = status.errors.clone(); - return Err(QueryError::JobFailed { errors }); - } - - Ok(res) + check_job_status(res) } } @@ -176,13 +169,37 @@ impl RetryContext { let job_ref = generate_job_reference(project_id, &self.template.request.location); let job = Job::new() .set_configuration(job_config) - .set_job_reference(job_ref); + .set_job_reference(job_ref.clone()); let req = InsertJobRequest::new() .set_job(job) .set_project_id(project_id); // Box heavy RPC call future to avoid large stack frames. - let job = Box::pin(InsertJobExecutor::new(job_service.clone(), req).execute()).await?; + let job = match Box::pin(InsertJobExecutor::new(job_service.clone(), req).execute()).await { + Ok(job) => job, + Err(err) if is_duplicate_job_error(&err) => { + // A fresh job ID is generated per attempt, so a duplicate means + // an earlier attempt of this request reached the service. Adopt + // the job it created rather than fail a running, billing query. + let existing_job = match build_get_job(&job_service, &job_ref) { + Some(get) => { + // The original error names the running job, and unlike a + // `jobs.get` failure it never makes the job retry loop + // reissue the query, so we discard the err if the job + // request fails with `.ok`. + Box::pin(get.send()).await.ok() + } + None => None, + }; + let Some(existing_job) = existing_job else { + // We were unable to successfully send a `jobs.get` RPC. + // Return the original error message. + return Err(err); + }; + check_job_status(existing_job)? + } + Err(err) => return Err(err), + }; Ok(QueryHandle::from_job( job_service, @@ -221,6 +238,18 @@ impl RetryContext { } } +/// Returns [`QueryError::JobFailed`] if the service reports the job as failed. +fn check_job_status(job: Job) -> Result { + if let Some(status) = job.status.as_ref() + && status.error_result.is_some() + { + let errors = status.errors.clone(); + return Err(QueryError::JobFailed { errors }); + } + + Ok(job) +} + #[cfg(test)] mod tests { use super::*; @@ -230,13 +259,31 @@ mod tests { QueryResponse, }; use google_cloud_gax::error::Error as GaxError; - use google_cloud_gax::error::rpc::{Code, Status}; + use google_cloud_gax::error::rpc::{Code, Status, StatusDetails}; use google_cloud_gax::response::Response; + use google_cloud_rpc::model::ErrorInfo; use serde_json::{Map, json}; + use std::sync::Mutex; use test_case::test_case; type TestResult = anyhow::Result<()>; + // The error BigQuery returns when a request tries to create a job that a + // previous, apparently failed, attempt of the same request already created. + fn duplicate_job_error(job_id: &str) -> GaxError { + let message = format!("Already Exists: Job my-project:US.{job_id}"); + let status = Status::default() + .set_code(Code::AlreadyExists) + .set_message(message.clone()) + .set_details(vec![StatusDetails::ErrorInfo( + ErrorInfo::new() + .set_reason("duplicate") + .set_domain("global") + .set_metadata([("message".to_string(), message)]), + )]); + GaxError::service(status) + } + #[tokio::test] async fn test_jobs_query_execute_success() -> TestResult { let mut mock = MockJobService::new(); @@ -468,6 +515,86 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_jobs_insert_duplicate_adopts_existing_job() -> TestResult { + let inserted = Arc::new(Mutex::new(None)); + + let mut mock = MockJobService::new(); + let captured = inserted.clone(); + mock.expect_insert_job().return_once(move |req, _| { + let job_ref = req.job.unwrap().job_reference.unwrap(); + let err = duplicate_job_error(&job_ref.job_id); + *captured.lock().unwrap() = Some(job_ref); + Err(err) + }); + mock.expect_get_job().return_once(move |req, _| { + let job_ref = JobReference::new() + .set_project_id(req.project_id) + .set_job_id(req.job_id) + .set_location(req.location); + let job = Job::new() + .set_configuration(JobConfiguration::new().set_query(JobConfigurationQuery::new())) + .set_job_reference(job_ref) + .set_status(JobStatus::new().set_state("RUNNING")); + Ok(Response::from(job)) + }); + + let job_service = create_job_service(mock); + let query = Query::new(job_service, "SELECT 1".to_string()) + .with_project_id("my-project") + .set_location("us-central1") + .set_priority("BATCH"); + + let retry_ctx = RetryContext::new(query); + assert!(retry_ctx.force_job_path(), "priority should force job path"); + + let handle = retry_ctx.execute_once("my-project").await?; + + // The recovered job must be the one the earlier attempt created, and it + // must be fetched from the location of the query. + let inserted = inserted.lock().unwrap().clone().expect("job inserted"); + let job_ref = handle.metadata.job_reference.expect("job reference"); + assert_eq!(job_ref, inserted, "{job_ref:?}"); + assert_eq!(job_ref.location.as_deref(), Some("us-central1")); + assert!(!handle.completed, "the recovered job is still running"); + Ok(()) + } + + #[tokio::test] + async fn test_jobs_insert_duplicate_reports_original_error() -> TestResult { + let mut mock = MockJobService::new(); + mock.expect_insert_job() + .return_once(move |_, _| Err(duplicate_job_error("job_123"))); + // The job cannot be fetched, for example because the caller lacks + // permissions to read it. + mock.expect_get_job().return_once(move |_, _| { + let status = Status::default() + .set_code(Code::PermissionDenied) + .set_message("simulated permission denied"); + Err(GaxError::service(status)) + }); + + let job_service = create_job_service(mock); + let query = Query::new(job_service, "SELECT 1".to_string()) + .with_project_id("my-project") + .set_priority("BATCH"); + + let err = RetryContext::new(query) + .execute_once("my-project") + .await + .unwrap_err(); + + // The duplicate error names the running job, so it is more useful than + // the error from `jobs.get`. + let QueryError::Rpc { source } = &err else { + panic!("expected QueryError::Rpc, got {err:?}"); + }; + let status = source.status().expect("status"); + assert_eq!(status.code, Code::AlreadyExists, "{status:?}"); + assert!(status.message.contains("Already Exists: Job"), "{status:?}"); + Ok(()) + } + #[tokio::test] async fn test_query_rpcs_are_idempotent() -> TestResult { let mut mock = MockJobService::new(); diff --git a/src/bigquery/src/query/query_handle.rs b/src/bigquery/src/query/query_handle.rs index 2282001d8a..9762e87a42 100644 --- a/src/bigquery/src/query/query_handle.rs +++ b/src/bigquery/src/query/query_handle.rs @@ -436,7 +436,7 @@ impl CompleteQuery { /// Builds a `jobs.get` request from a job reference, or `None` if the /// reference cannot identify a job. Dry-run queries return a job reference /// without a job ID. -fn build_get_job(job_service: &JobService, job_ref: &JobReference) -> Option { +pub(crate) fn build_get_job(job_service: &JobService, job_ref: &JobReference) -> Option { if job_ref.job_id.is_empty() { return None; } diff --git a/src/bigquery/src/query/retry_policy.rs b/src/bigquery/src/query/retry_policy.rs index 663a03091b..c9a27a5642 100644 --- a/src/bigquery/src/query/retry_policy.rs +++ b/src/bigquery/src/query/retry_policy.rs @@ -218,8 +218,6 @@ pub(crate) fn is_retryable_error_reason(reason: &str) -> bool { /// Returns true if `error` reports a conflict with a resource that already /// exists, such as `409 Already Exists: Job my-project:US.job_1234567890`. -// TODO(#6717): use this function on execution.rs and remove here -#[allow(dead_code)] pub(crate) fn is_duplicate_job_error(error: &QueryError) -> bool { let QueryError::Rpc { source } = error else { return false;