google_cloud_bigquery/query/
builder.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::RequestOptions;
16use crate::error::QueryError;
17use crate::generated::QueryRequest;
18use crate::query::execution::RetryContext;
19use crate::query::retry_policy::{JobRetryPolicy, default_job_retry_policy};
20use crate::query::{CompleteQuery, Query as QueryHandle, Result};
21use google_cloud_bigquery_v2::client::JobService;
22use google_cloud_bigquery_v2::model::JobReference;
23use google_cloud_bigquery_v2::model::query_request::JobCreationMode;
24use std::sync::Arc;
25use uuid::Uuid;
26
27pub(crate) const JOB_ID_PREFIX: &str = "job_";
28pub(crate) const QUERY_REQUEST_ID_PREFIX: &str = "req_";
29
30/// A builder for configuring and executing a SQL query.
31///
32/// [`BigQuery::query()`](crate::client::BigQuery::query) returns a [`Query`] builder.
33///
34/// Use this builder to configure query parameters, dataset defaults, geographic location,
35/// result limits, and caching before executing the query with [`send()`](Query::send)
36/// or [`until_done()`](Query::until_done).
37///
38/// # Example
39///
40/// ```
41/// # use google_cloud_bigquery::client::BigQuery;
42/// # async fn sample(client: BigQuery) -> anyhow::Result<()> {
43/// let mut rows = client
44///     .query("SELECT name FROM `bigquery-public-data.usa_names.usa_1910_2013` WHERE state = 'TX' LIMIT 100")
45///     .set_location("US")
46///     .set_page_size(50_u32)
47///     .until_done()
48///     .await?
49///     .read();
50///
51/// while let Some(row) = rows.next().await.transpose()? {
52///     let name: String = row.get("name")?;
53///     println!("Name: {name}");
54/// }
55/// # Ok(())
56/// # }
57/// ```
58#[derive(Clone, Debug)]
59pub struct Query {
60    pub(crate) job_service: Arc<JobService>,
61    pub(crate) request: QueryRequest,
62    pub(crate) project_id: Option<String>,
63    pub(crate) job_retry_policy: Arc<dyn JobRetryPolicy>,
64    pub(crate) options: RequestOptions,
65}
66
67impl Query {
68    /// Creates a new `Query` builder for the given SQL query.
69    pub(crate) fn new(job_service: Arc<JobService>, sql: String) -> Self {
70        Self {
71            job_service,
72            request: QueryRequest::default()
73                .set_query(sql)
74                .set_use_legacy_sql(wkt::BoolValue::from(false))
75                .set_job_creation_mode(JobCreationMode::JobCreationOptional),
76            project_id: None,
77            job_retry_policy: default_job_retry_policy(),
78            options: RequestOptions::default(),
79        }
80    }
81
82    /// Sets the project ID to override the default client project ID for query
83    /// execution and billing.
84    ///
85    /// If you configured a default project ID on the
86    /// [`BigQuery`][crate::client::BigQuery] client via
87    /// [`ClientBuilder::with_project_id`](crate::builder::bigquery::ClientBuilder::with_project_id),
88    /// the query inherits it automatically. Calling this method overrides the
89    /// project ID for this specific query.
90    ///
91    /// You must specify a project ID either on the client or on the query
92    /// builder before calling [`send()`](Query::send).
93    ///
94    /// # Example
95    ///
96    /// ```
97    /// # use google_cloud_bigquery::client::BigQuery;
98    /// # async fn sample(client: BigQuery) -> anyhow::Result<()> {
99    /// let query_handle = client
100    ///     .query("SELECT 1 AS count")
101    ///     .with_project_id("my-project-id")
102    ///     .send()
103    ///     .await?;
104    /// # Ok(())
105    /// # }
106    /// ```
107    pub fn with_project_id<S: Into<String>>(mut self, project_id: S) -> Self {
108        let project_id = project_id.into();
109        self.project_id = Some(project_id.clone());
110        self.options.set_quota_project(project_id);
111        self
112    }
113
114    /// Sets the User-Agent header to amend the default client header.
115    ///
116    /// # Example
117    ///
118    /// ```
119    /// # use google_cloud_bigquery::client::BigQuery;
120    /// # async fn sample(client: BigQuery) -> anyhow::Result<()> {
121    /// let query_handle = client
122    ///     .query("SELECT 1 AS count")
123    ///     .with_user_agent("my-app/1.2.3")
124    ///     .send()
125    ///     .await?;
126    /// # Ok(())
127    /// # }
128    /// ```
129    pub fn with_user_agent(mut self, user_agent: impl Into<String>) -> Self {
130        self.options.set_user_agent(user_agent);
131        self
132    }
133
134    /// Executes the SQL query.
135    ///
136    /// Returns a [`Query`](QueryHandle) handle representing the query execution.
137    /// You can call [`until_done()`](QueryHandle::until_done) on the returned handle
138    /// to wait for results, or inspect the initial [`metadata()`](QueryHandle::metadata).
139    ///
140    /// You must configure a target project ID on either the
141    /// [`BigQuery`][crate::client::BigQuery] client via
142    /// [`ClientBuilder::with_project_id`](crate::builder::bigquery::ClientBuilder::with_project_id)
143    /// or on this builder via [`with_project_id()`](Query::with_project_id).
144    ///
145    /// # Example
146    ///
147    /// ```
148    /// # use google_cloud_bigquery::client::BigQuery;
149    /// # async fn sample(client: BigQuery) -> anyhow::Result<()> {
150    /// let query_handle = client
151    ///     .query("SELECT CURRENT_TIMESTAMP() AS now")
152    ///     .send()
153    ///     .await?;
154    ///
155    /// let completed_query = query_handle
156    ///     .until_done()
157    ///     .await?;
158    ///
159    /// let mut rows = completed_query.read();
160    /// if let Some(row) = rows.next().await.transpose()? {
161    ///     let now: String = row.get("now")?;
162    ///     println!("Current time: {now}");
163    /// }
164    /// # Ok(())
165    /// # }
166    /// ```
167    pub async fn send(self) -> Result<QueryHandle> {
168        // Box heavy RPC call future to avoid large stack frames.
169        Box::pin(RetryContext::new(self).execute()).await
170    }
171
172    /// Sends the query execution request and waits until execution completes.
173    ///
174    /// This is a convenience method equivalent to calling
175    /// [`.send().await?.until_done().await`](Query::send).
176    ///
177    /// # Errors
178    ///
179    /// Returns [`QueryError::DryRun`](crate::error::QueryError::DryRun) if the query was configured as a dry run.
180    ///
181    /// Returns an error if a remote service or network failure happens during
182    /// execution or polling, or if the BigQuery job fails due to runtime execution errors.
183    ///
184    /// # Example
185    ///
186    /// ```
187    /// # use google_cloud_bigquery::client::BigQuery;
188    /// # async fn sample(client: BigQuery) -> anyhow::Result<()> {
189    /// let completed_query = client
190    ///     .query("SELECT CURRENT_TIMESTAMP() AS now")
191    ///     .until_done()
192    ///     .await?;
193    ///
194    /// let mut rows = completed_query.read();
195    /// if let Some(row) = rows.next().await.transpose()? {
196    ///     let now: String = row.get("now")?;
197    ///     println!("Current time: {now}");
198    /// }
199    /// # Ok(())
200    /// # }
201    /// ```
202    pub async fn until_done(self) -> Result<CompleteQuery> {
203        if self.request.dry_run {
204            return Err(QueryError::DryRun);
205        }
206        // Box heavy RPC call future to avoid large stack frames.
207        Box::pin(async move { self.send().await?.until_done().await }).await
208    }
209}
210
211// Create a job reference with a generated job ID.
212//
213// BigQuery does not strictly define a format for job IDs, just a limit in size.
214// See https://docs.cloud.google.com/bigquery/docs/reference/rest/v2/JobReference
215pub(crate) fn generate_job_reference(project_id: &str, location: &str) -> JobReference {
216    let job_id = generate_prefixed_id(JOB_ID_PREFIX);
217    let mut job_ref = JobReference::new()
218        .set_project_id(project_id.to_string())
219        .set_job_id(job_id);
220
221    if !location.is_empty() {
222        job_ref = job_ref.set_location(location.to_string());
223    }
224    job_ref
225}
226
227// Create a random ID with the given prefix and a UUID.
228//
229// BigQuery does not strictly define a format for request and job IDs, just a limit in size.
230// However, request IDs are more restrictive and have a limit of 36 characters.
231// UUID v4 simple format is used to reduce length.
232//
233// See https://docs.cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query#QueryRequest for request ID
234// and https://docs.cloud.google.com/bigquery/docs/reference/rest/v2/JobReference for job ID.
235pub(crate) fn generate_prefixed_id(prefix: &str) -> String {
236    format!("{prefix}{}", Uuid::new_v4().simple())
237}
238
239include!("../generated/builder.rs");
240
241#[cfg(test)]
242mod tests {
243    use super::*;
244    use crate::client::BigQuery;
245    use crate::error::QueryError;
246    use crate::query::tests::{MockJobService, create_job_service};
247    use google_cloud_auth::credentials::anonymous::Builder as Anonymous;
248    use google_cloud_bigquery_v2::model::query_request::JobCreationMode;
249    use google_cloud_bigquery_v2::model::{
250        ErrorProto, Job, JobConfiguration, JobReference, JobStatus,
251        QueryRequest as JobsQueryRequest, QueryResponse,
252    };
253    use google_cloud_gax::error::Error as GaxError;
254    use google_cloud_gax::error::rpc::Status;
255    use google_cloud_gax::response::Response;
256
257    // bigquery limits request id to 36 characters
258    const BIGQUERY_REQ_ID_LIMIT: usize = 36;
259
260    type TestResult = anyhow::Result<()>;
261
262    #[test]
263    fn test_new() {
264        let job_service = create_job_service(MockJobService::new());
265        let sql = "SELECT 1".to_string();
266        let query_builder = Query::new(job_service, sql.clone());
267        assert_eq!(query_builder.request.query, sql);
268        assert_eq!(
269            query_builder.request.use_legacy_sql,
270            Some(wkt::BoolValue::from(false))
271        );
272        assert_eq!(
273            query_builder.request.job_creation_mode,
274            JobCreationMode::JobCreationOptional
275        );
276        assert_eq!(query_builder.project_id, None);
277        assert_eq!(query_builder.options.quota_project(), &None);
278        assert_eq!(query_builder.options.user_agent(), &None);
279    }
280
281    #[test]
282    fn test_with_project_id() {
283        let job_service = create_job_service(MockJobService::new());
284        let query_builder =
285            Query::new(job_service, "SELECT 1".to_string()).with_project_id("my-project");
286        assert_eq!(query_builder.project_id.unwrap(), "my-project");
287        assert_eq!(
288            query_builder.options.quota_project().as_deref(),
289            Some("my-project")
290        );
291    }
292
293    #[test]
294    fn test_with_user_agent() {
295        let job_service = create_job_service(MockJobService::new());
296        let query_builder =
297            Query::new(job_service, "SELECT 1".to_string()).with_user_agent("my-app/1.2.3");
298        assert_eq!(
299            query_builder.options.user_agent().as_deref(),
300            Some("my-app/1.2.3")
301        );
302    }
303
304    #[tokio::test]
305    async fn test_run_missing_project_id() -> anyhow::Result<()> {
306        let client = BigQuery::builder()
307            .with_credentials(Anonymous::new().build())
308            .build()
309            .await?;
310        let query_builder = client.query("SELECT 1");
311        let err = query_builder
312            .send()
313            .await
314            .expect_err("should return an error when project_id is missing");
315        assert!(
316            matches!(&err, QueryError::Rpc { source } if source.is_binding()),
317            "expected Binding error for missing project ID, got {err:?}"
318        );
319        Ok(())
320    }
321
322    #[tokio::test]
323    async fn test_run_missing_project_id_force_job_path() -> anyhow::Result<()> {
324        let client = BigQuery::builder()
325            .with_credentials(Anonymous::new().build())
326            .build()
327            .await?;
328        let query_builder = client.query("SELECT 1").set_allow_large_results(true);
329        let err = query_builder
330            .send()
331            .await
332            .expect_err("should return an error when project_id is missing");
333        assert!(
334            matches!(&err, QueryError::Rpc { source } if source.is_binding()),
335            "expected Binding error for missing project ID on job path, got {err:?}"
336        );
337        Ok(())
338    }
339
340    #[tokio::test]
341    async fn test_run_until_done_missing_project_id() -> anyhow::Result<()> {
342        let client = BigQuery::builder()
343            .with_credentials(Anonymous::new().build())
344            .build()
345            .await?;
346        let query_builder = client.query("SELECT 1");
347        let err = query_builder
348            .until_done()
349            .await
350            .expect_err("should return an error when project_id is missing");
351        assert!(
352            matches!(&err, QueryError::Rpc { source } if source.is_binding()),
353            "expected Binding error for missing project ID, got {err:?}"
354        );
355        Ok(())
356    }
357
358    #[test]
359    fn test_generate_prefixed_id() {
360        let job_id = generate_prefixed_id(JOB_ID_PREFIX);
361        assert!(job_id.starts_with(JOB_ID_PREFIX), "{job_id:?}");
362        assert!(
363            Uuid::parse_str(&job_id[JOB_ID_PREFIX.len()..]).is_ok(),
364            "{job_id:?}"
365        );
366
367        let req_id = generate_prefixed_id(QUERY_REQUEST_ID_PREFIX);
368        assert!(req_id.starts_with(QUERY_REQUEST_ID_PREFIX), "{req_id:?}");
369        assert!(req_id.len() <= BIGQUERY_REQ_ID_LIMIT, "{req_id:?}");
370        assert!(
371            Uuid::parse_str(&req_id[QUERY_REQUEST_ID_PREFIX.len()..]).is_ok(),
372            "{req_id:?}"
373        );
374    }
375
376    #[test]
377    fn test_generate_job_reference() {
378        let job_ref = generate_job_reference("my-project", "us-central1");
379        assert_eq!(job_ref.project_id, "my-project");
380        assert!(job_ref.job_id.starts_with(JOB_ID_PREFIX), "{job_ref:?}");
381        assert_eq!(job_ref.location.as_deref(), Some("us-central1"));
382    }
383
384    #[tokio::test]
385    async fn test_run_jobs_insert() -> TestResult {
386        let mut mock = MockJobService::new();
387        mock.expect_insert_job().returning(|req, options| {
388            assert_eq!(options.quota_project().as_deref(), Some("my-project"));
389            let job_ref = req.job.as_ref().unwrap().job_reference.as_ref().unwrap();
390            assert!(job_ref.job_id.starts_with(JOB_ID_PREFIX), "{job_ref:?}");
391            let job_ref = JobReference::new()
392                .set_job_id("test-job")
393                .set_project_id("my-project");
394            let job = Job::new()
395                .set_job_reference(job_ref)
396                .set_status(JobStatus::new().set_state("DONE"));
397            Ok(google_cloud_gax::response::Response::from(job))
398        });
399        mock.expect_query().never();
400
401        let job_service = create_job_service(mock);
402
403        let query_builder = Query::new(job_service, "SELECT 1".to_string())
404            .with_project_id("my-project")
405            .set_allow_large_results(true);
406        let query = query_builder.send().await?;
407        assert!(query.completed, "{query:?}");
408        assert_eq!(query.options.quota_project().as_deref(), Some("my-project"));
409
410        Ok(())
411    }
412
413    #[tokio::test]
414    async fn test_run_jobs_query() -> TestResult {
415        let mut mock = MockJobService::new();
416        mock.expect_query().returning(move |req, options| {
417            assert_eq!(options.quota_project().as_deref(), Some("my-project"));
418            let req_id = &req.query_request.as_ref().unwrap().request_id;
419            assert!(req_id.starts_with(QUERY_REQUEST_ID_PREFIX), "{req_id:?}");
420            assert!(req_id.len() <= BIGQUERY_REQ_ID_LIMIT, "{req_id:?}");
421            Ok(Response::from(
422                QueryResponse::new().set_query_id("some_query_id"),
423            ))
424        });
425        mock.expect_insert_job().never();
426
427        let job_service = create_job_service(mock);
428        let query_builder =
429            Query::new(job_service, "SELECT 1".to_string()).with_project_id("my-project");
430        let query = query_builder.send().await?;
431        assert!(!query.completed, "{query:?}");
432        assert_eq!(query.metadata.query_id, "some_query_id");
433        assert_eq!(query.options.quota_project().as_deref(), Some("my-project"));
434
435        Ok(())
436    }
437
438    #[test]
439    fn test_force_job_path() {
440        let job_service = create_job_service(MockJobService::new());
441        let mut query_builder = Query::new(job_service, "SELECT 1".to_string());
442        assert!(!query_builder.request.force_job_path());
443
444        // setting a jobs.insert exclusive field
445        query_builder = query_builder.set_allow_large_results(true);
446        assert!(query_builder.request.force_job_path());
447    }
448
449    #[test]
450    fn test_request_conversions() {
451        let req = QueryRequest::default()
452            .set_query("SELECT 1".to_string())
453            .set_dry_run(true)
454            .set_use_legacy_sql(true);
455
456        let query_request: JobsQueryRequest = req.clone().into();
457        assert_eq!(query_request.query, "SELECT 1");
458        assert!(query_request.dry_run);
459        assert_eq!(
460            query_request.use_legacy_sql,
461            Some(wkt::BoolValue::from(true))
462        );
463
464        let job_config: JobConfiguration = req.into();
465        let job_query = job_config.query.as_ref().unwrap();
466        assert_eq!(job_query.query, "SELECT 1");
467        assert_eq!(job_query.use_legacy_sql, Some(wkt::BoolValue::from(true)));
468    }
469
470    #[tokio::test(start_paused = true)]
471    async fn test_run_reissue_on_retryable_job_failed() -> TestResult {
472        let mut mock = MockJobService::new();
473        let mut seq = mockall::Sequence::new();
474        let first_request_id = std::sync::Arc::new(std::sync::Mutex::new(String::new()));
475        let first_request_id_clone = first_request_id.clone();
476
477        mock.expect_query()
478            .in_sequence(&mut seq)
479            .times(1)
480            .returning(move |req, _| {
481                let req_id = req.query_request.as_ref().unwrap().request_id.clone();
482                assert!(req_id.starts_with(QUERY_REQUEST_ID_PREFIX));
483                *first_request_id_clone.lock().unwrap() = req_id;
484                let err_proto = ErrorProto::new()
485                    .set_reason("backendError")
486                    .set_message("temporary server issue");
487                Ok(Response::from(
488                    QueryResponse::new().set_errors(vec![err_proto]),
489                ))
490            });
491
492        mock.expect_query()
493            .in_sequence(&mut seq)
494            .times(1)
495            .returning(move |req, _| {
496                let req_id = req.query_request.as_ref().unwrap().request_id.clone();
497                assert!(req_id.starts_with(QUERY_REQUEST_ID_PREFIX));
498                assert_ne!(
499                    req_id,
500                    *first_request_id.lock().unwrap(),
501                    "reissued query must generate a fresh request_id to avoid 409 duplicate ID conflicts"
502                );
503                Ok(Response::from(
504                    QueryResponse::new().set_query_id("q_success"),
505                ))
506            });
507
508        let job_service = create_job_service(mock);
509        let query_builder =
510            Query::new(job_service, "SELECT 1".to_string()).with_project_id("my-project");
511        let query = query_builder.send().await?;
512        assert_eq!(query.metadata.query_id, "q_success");
513
514        Ok(())
515    }
516
517    #[tokio::test(start_paused = true)]
518    async fn test_run_reissue_on_retryable_rpc_error() -> TestResult {
519        let mut mock = MockJobService::new();
520        let mut seq = mockall::Sequence::new();
521        let first_request_id = std::sync::Arc::new(std::sync::Mutex::new(String::new()));
522        let first_request_id_clone = first_request_id.clone();
523
524        mock.expect_query()
525            .in_sequence(&mut seq)
526            .times(1)
527            .returning(move |req, _| {
528                let req_id = req.query_request.as_ref().unwrap().request_id.clone();
529                assert!(req_id.starts_with(QUERY_REQUEST_ID_PREFIX));
530                *first_request_id_clone.lock().unwrap() = req_id;
531                const BQ_REST_PAYLOAD: &[u8] = br#"{
532  "error": {
533    "code": 400,
534    "message": "The job encountered an error during execution. Retrying the job may solve the problem.",
535    "errors": [
536      {
537        "message": "The job encountered an error during execution. Retrying the job may solve the problem.",
538        "domain": "global",
539        "reason": "backendError"
540      }
541    ],
542    "status": "INVALID_ARGUMENT"
543  }
544}"#;
545                let status = Status::try_from(&bytes::Bytes::from_static(BQ_REST_PAYLOAD)).unwrap();
546                Err(GaxError::service(status))
547            });
548
549        mock.expect_query()
550            .in_sequence(&mut seq)
551            .times(1)
552            .returning(move |req, _| {
553                let req_id = req.query_request.as_ref().unwrap().request_id.clone();
554                assert!(req_id.starts_with(QUERY_REQUEST_ID_PREFIX));
555                assert_ne!(
556                    req_id,
557                    *first_request_id.lock().unwrap(),
558                    "reissued query must generate a fresh request_id to avoid 409 duplicate ID conflicts"
559                );
560                Ok(Response::from(
561                    QueryResponse::new().set_query_id("q_success"),
562                ))
563            });
564
565        let job_service = create_job_service(mock);
566        let query_builder =
567            Query::new(job_service, "SELECT 1".to_string()).with_project_id("my-project");
568        let query = query_builder.send().await?;
569        assert_eq!(query.metadata.query_id, "q_success");
570
571        Ok(())
572    }
573
574    #[tokio::test]
575    async fn test_run_jobs_query_with_page_size() -> TestResult {
576        let mut mock = MockJobService::new();
577        mock.expect_query().returning(move |req, _| {
578            assert_eq!(
579                req.query_request.as_ref().and_then(|r| r.max_results),
580                Some(100)
581            );
582            Ok(Response::from(QueryResponse::new()))
583        });
584        let job_service = create_job_service(mock);
585        let query_builder = Query::new(job_service, "SELECT 1".to_string())
586            .with_project_id("my-project")
587            .set_page_size(100_u32);
588        let query = query_builder.send().await?;
589        assert_eq!(query.page_size, Some(100));
590
591        Ok(())
592    }
593
594    #[tokio::test]
595    async fn test_run_jobs_insert_with_page_size() -> TestResult {
596        let mut mock = MockJobService::new();
597        mock.expect_insert_job().returning(|_, _| {
598            let job_ref = JobReference::new()
599                .set_job_id("test-job")
600                .set_project_id("my-project");
601            let job = Job::new()
602                .set_job_reference(job_ref)
603                .set_status(JobStatus::new().set_state("DONE"));
604            Ok(google_cloud_gax::response::Response::from(job))
605        });
606        let job_service = create_job_service(mock);
607
608        let query_builder = Query::new(job_service, "SELECT 1".to_string())
609            .with_project_id("my-project")
610            .set_allow_large_results(true)
611            .set_page_size(50_u32);
612        let query = query_builder.send().await?;
613        assert_eq!(query.page_size, Some(50));
614
615        Ok(())
616    }
617
618    #[tokio::test]
619    async fn test_until_done_jobs_query() -> TestResult {
620        let mut mock = MockJobService::new();
621        mock.expect_query().returning(move |_, options| {
622            assert_eq!(options.quota_project().as_deref(), Some("my-project"));
623            Ok(Response::from(
624                QueryResponse::new()
625                    .set_job_complete(true)
626                    .set_query_id("some_query_id"),
627            ))
628        });
629        let job_service = create_job_service(mock);
630        let query = Query::new(job_service, "SELECT 1".to_string()).with_project_id("my-project");
631        let complete = query.until_done().await?;
632        assert_eq!(complete.metadata().query_id, "some_query_id");
633        assert_eq!(
634            complete.options.quota_project().as_deref(),
635            Some("my-project")
636        );
637
638        Ok(())
639    }
640
641    #[tokio::test]
642    async fn test_until_done_dry_run_returns_error() -> TestResult {
643        let mut mock = MockJobService::new();
644        mock.expect_insert_job().never();
645        mock.expect_query().never();
646        let job_service = create_job_service(mock);
647        let query = Query::new(job_service, "SELECT 1".to_string())
648            .with_project_id("my-project")
649            .set_dry_run(true);
650        let err = query.until_done().await.unwrap_err();
651        assert!(
652            matches!(err, QueryError::DryRun),
653            "expected DryRun error, got {err:?}"
654        );
655
656        Ok(())
657    }
658
659    #[tokio::test]
660    async fn test_run_jobs_query_with_user_agent() -> TestResult {
661        let mut mock = MockJobService::new();
662        mock.expect_query().returning(move |_, options| {
663            assert_eq!(options.user_agent().as_deref(), Some("my-app/1.2.3"));
664            Ok(Response::from(QueryResponse::new()))
665        });
666        let job_service = create_job_service(mock);
667        let query_builder = Query::new(job_service, "SELECT 1".to_string())
668            .with_project_id("my-project")
669            .with_user_agent("my-app/1.2.3");
670        let query = query_builder.send().await?;
671        assert_eq!(query.options.user_agent().as_deref(), Some("my-app/1.2.3"));
672
673        Ok(())
674    }
675
676    #[tokio::test]
677    async fn test_run_jobs_insert_with_user_agent() -> TestResult {
678        let mut mock = MockJobService::new();
679        mock.expect_insert_job().returning(|_, options| {
680            assert_eq!(options.user_agent().as_deref(), Some("my-app/1.2.3"));
681            let job_ref = JobReference::new()
682                .set_job_id("test-job")
683                .set_project_id("my-project");
684            let job = Job::new()
685                .set_job_reference(job_ref)
686                .set_status(JobStatus::new().set_state("DONE"));
687            Ok(Response::from(job))
688        });
689        let job_service = create_job_service(mock);
690
691        let query_builder = Query::new(job_service, "SELECT 1".to_string())
692            .with_project_id("my-project")
693            .set_allow_large_results(true)
694            .with_user_agent("my-app/1.2.3");
695        let query = query_builder.send().await?;
696        assert_eq!(query.options.user_agent().as_deref(), Some("my-app/1.2.3"));
697
698        Ok(())
699    }
700}
google_cloud_bigquery/query/builder.rs

google_cloud_bigquery/query/
builder.rs
