feat(bigquery-v2): add lro polling support for bq jobs - WIP - prototype

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

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

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<crate::model::Job, crate::model::Job> {
+            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<T: Into<std::string::String>>(mut self, v: T) -> Self {
             self.0.request.project_id = v.into();

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;

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;

src/generated/cloud/bigquery/v2/src/operation.rs

@@ -0,0 +1,86 @@
+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<google_cloud_gax::error::rpc::Status> {
+        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
+            })
+    }
+}
+
+#[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());
+    }
+}
+

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:?}");