From 9b15d232183c61fe3df0598141dba2db20e02258 Mon Sep 17 00:00:00 2001 From: haphungw Date: Wed, 15 Jul 2026 20:21:24 +0000 Subject: [PATCH 1/2] feat(bigquery): generate lro::Poller for InsertJob This change enables standard LRO polling for InsertJob using BigQuery's custom job API. We implement DiscoveryOperation for Job to allow google-cloud-lro to wrap the get_job call seamlessly. Custom mutational retries for BigQuery will be handled independently via a wrapper JobPoller on top of this basic building block. --- Cargo.lock | 1 + librarian.yaml | 6 +++ src/generated/cloud/bigquery/v2/Cargo.toml | 1 + .../cloud/bigquery/v2/src/builder.rs | 48 +++++++++++++++++++ src/generated/cloud/bigquery/v2/src/lib.rs | 3 ++ src/generated/cloud/bigquery/v2/src/model.rs | 1 + .../cloud/bigquery/v2/src/operation.rs | 27 +++++++++++ 7 files changed, 87 insertions(+) create mode 100644 src/generated/cloud/bigquery/v2/src/operation.rs diff --git a/Cargo.lock b/Cargo.lock index c29b91a893..82380d5e1d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2118,6 +2118,7 @@ dependencies = [ "bytes", "google-cloud-gax", "google-cloud-gax-internal", + "google-cloud-lro", "google-cloud-type", "google-cloud-wkt", "serde", diff --git a/librarian.yaml b/librarian.yaml index 5ccd256262..dae82bb2ce 100644 --- a/librarian.yaml +++ b/librarian.yaml @@ -501,8 +501,14 @@ libraries: version: 1.0.0 copyright_year: "2025" skip_release: true + keep: + - src/operation.rs rust: package_name_override: google-cloud-bigquery-v2 + package_dependencies: + - name: google-cloud-lro + package: google-cloud-lro + force_used: true - name: google-cloud-bigquery-write copyright_year: "2026" output: src/bigquery-write diff --git a/src/generated/cloud/bigquery/v2/Cargo.toml b/src/generated/cloud/bigquery/v2/Cargo.toml index dd3763128b..94dbc04622 100644 --- a/src/generated/cloud/bigquery/v2/Cargo.toml +++ b/src/generated/cloud/bigquery/v2/Cargo.toml @@ -46,6 +46,7 @@ async-trait.workspace = true bytes.workspace = true gaxi = { workspace = true, features = ["_internal-http-client"] } google-cloud-gax.workspace = true +google-cloud-lro.workspace = true google-cloud-type.workspace = true serde.workspace = true serde_json.workspace = true diff --git a/src/generated/cloud/bigquery/v2/src/builder.rs b/src/generated/cloud/bigquery/v2/src/builder.rs index d41224d346..36af69558d 100644 --- a/src/generated/cloud/bigquery/v2/src/builder.rs +++ b/src/generated/cloud/bigquery/v2/src/builder.rs @@ -1016,6 +1016,54 @@ pub mod job_service { .map(crate::Response::into_body) } + /// Creates a [Poller][google_cloud_lro::Poller] to work with `InsertJob`. + pub fn poller(self) -> impl google_cloud_lro::Poller { + let req = self.0.request.clone(); + let stub = self.0.stub.clone(); + let options = self.0.options.clone(); + + let polling_error_policy = + std::sync::Arc::new(google_cloud_gax::polling_error_policy::AlwaysContinue); + let polling_backoff_policy = std::sync::Arc::new( + google_cloud_gax::exponential_backoff::ExponentialBackoff::default(), + ); + + let query = move |name: String| { + let stub_clone = stub.clone(); + let mut options_clone = options.clone(); + let req_clone = req.clone(); + options_clone.set_retry_policy(google_cloud_gax::retry_policy::NeverRetry); + async move { + let get_req = crate::model::GetJobRequest { + project_id: req_clone.project_id.clone(), + job_id: name, + location: req_clone + .job + .as_ref() + .and_then(|j| j.job_reference.as_ref()) + .and_then(|jr| jr.location.clone()) + .unwrap_or_default(), + _unknown_fields: std::default::Default::default(), + }; + let job = stub_clone + .get_job(get_req, options_clone.clone()) + .await + .map(crate::Response::into_body)?; + + Ok(job) + } + }; + + let start = move || async move { self.send().await }; + + google_cloud_lro::internal::new_discovery_poller( + polling_error_policy, + polling_backoff_policy, + start, + query, + ) + } + /// Sets the value of [project_id][crate::model::InsertJobRequest::project_id]. pub fn set_project_id>(mut self, v: T) -> Self { self.0.request.project_id = v.into(); diff --git a/src/generated/cloud/bigquery/v2/src/lib.rs b/src/generated/cloud/bigquery/v2/src/lib.rs index 47c2198ae8..ca89fdef26 100644 --- a/src/generated/cloud/bigquery/v2/src/lib.rs +++ b/src/generated/cloud/bigquery/v2/src/lib.rs @@ -110,3 +110,6 @@ pub(crate) use google_cloud_gax::client_builder::internal::new_builder as new_cl pub(crate) use google_cloud_gax::options::RequestOptions; pub(crate) use google_cloud_gax::options::internal::RequestBuilder; pub(crate) use google_cloud_gax::response::Response; + +#[allow(missing_docs)] +pub mod operation; diff --git a/src/generated/cloud/bigquery/v2/src/model.rs b/src/generated/cloud/bigquery/v2/src/model.rs index 2fffd25ae2..2297be4693 100644 --- a/src/generated/cloud/bigquery/v2/src/model.rs +++ b/src/generated/cloud/bigquery/v2/src/model.rs @@ -23,6 +23,7 @@ extern crate async_trait; extern crate bytes; extern crate gaxi; extern crate google_cloud_gax; +extern crate google_cloud_lro; extern crate google_cloud_type; extern crate serde; extern crate serde_json; diff --git a/src/generated/cloud/bigquery/v2/src/operation.rs b/src/generated/cloud/bigquery/v2/src/operation.rs new file mode 100644 index 0000000000..18cadf2fe0 --- /dev/null +++ b/src/generated/cloud/bigquery/v2/src/operation.rs @@ -0,0 +1,27 @@ +use crate::model::Job; +use google_cloud_lro::internal::DiscoveryOperation; + +impl DiscoveryOperation for Job { + fn done(&self) -> bool { + self.status + .as_ref() + .map(|s| s.state == "DONE") + .unwrap_or(false) + } + + fn name(&self) -> Option<&String> { + self.job_reference.as_ref().map(|r| &r.job_id) + } + + fn error(&self) -> Option { + self.status + .as_ref() + .and_then(|s| s.error_result.as_ref()) + .map(|err| { + let mut status = google_cloud_gax::error::rpc::Status::default(); + status.code = (google_cloud_gax::error::rpc::Code::Unknown as i32).into(); + status.message = err.message.clone(); + status + }) + } +} From 5c01e5bccd9b10f26d42cb0a8693576f93acc375 Mon Sep 17 00:00:00 2001 From: haphungw Date: Wed, 15 Jul 2026 20:34:16 +0000 Subject: [PATCH 2/2] test(bigquery): add tests for basic lro::Poller Includes unit tests for Job's DiscoveryOperation mapping. Updates integration tests to use `.poller().until_done().await` instead of `.send().await` to assert against live GCP servers that standard LRO polling works against BigQuery Jobs. --- .../cloud/bigquery/v2/src/operation.rs | 59 +++++++++++++++++++ tests/bigquery/src/lib.rs | 4 +- 2 files changed, 62 insertions(+), 1 deletion(-) diff --git a/src/generated/cloud/bigquery/v2/src/operation.rs b/src/generated/cloud/bigquery/v2/src/operation.rs index 18cadf2fe0..cdb5d02913 100644 --- a/src/generated/cloud/bigquery/v2/src/operation.rs +++ b/src/generated/cloud/bigquery/v2/src/operation.rs @@ -25,3 +25,62 @@ impl DiscoveryOperation for Job { }) } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::model::{ErrorProto, JobReference, JobStatus}; + + #[test] + fn test_done() { + let mut job = Job::default(); + assert!(!job.done(), "missing status should not be done"); + + job.status = Some(JobStatus { + state: "RUNNING".to_string(), + ..Default::default() + }); + assert!(!job.done(), "RUNNING should not be done"); + + job.status = Some(JobStatus { + state: "DONE".to_string(), + ..Default::default() + }); + assert!(job.done(), "DONE should be done"); + } + + #[test] + fn test_name() { + let mut job = Job::default(); + assert_eq!(job.name(), None, "missing job_reference should yield None"); + + job.job_reference = Some(JobReference { + job_id: "my_job".to_string(), + ..Default::default() + }); + assert_eq!(job.name(), Some(&"my_job".to_string()), "should return job_id"); + } + + #[test] + fn test_error() { + let mut job = Job::default(); + assert!(job.error().is_none(), "missing status should yield no error"); + + job.status = Some(JobStatus::default()); + assert!(job.error().is_none(), "missing error_result should yield no error"); + + job.status = Some(JobStatus { + error_result: Some(ErrorProto { + message: "some error".to_string(), + ..Default::default() + }), + ..Default::default() + }); + + let err = job.error().expect("should have error"); + assert_eq!(err.message, "some error"); + // Unknown code is 2. + assert_eq!(err.code, (google_cloud_gax::error::rpc::Code::Unknown as i32).into()); + } +} + diff --git a/tests/bigquery/src/lib.rs b/tests/bigquery/src/lib.rs index a6102f9ed6..8367b9eb8b 100644 --- a/tests/bigquery/src/lib.rs +++ b/tests/bigquery/src/lib.rs @@ -21,6 +21,7 @@ use google_cloud_bigquery_v2::model::{ Dataset, DatasetReference, Job, JobConfiguration, JobConfigurationQuery, JobReference, }; use google_cloud_gax::{error::rpc::Code, paginator::ItemPaginator}; +use google_cloud_lro::Poller; use google_cloud_test_utils::runtime_config::project_id; use rand::{RngExt, distr::Alphanumeric}; @@ -198,7 +199,8 @@ pub async fn job_service() -> Result<()> { .set_query(JobConfigurationQuery::new().set_query(query)), ), ) - .send() + .poller() + .until_done() .await?; println!("CREATE JOB = {job:?}");