google_cloud_bigquery/query/
client.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::builder::bigquery::Query;
17use crate::client::Read;
18use crate::error::QueryError;
19use crate::query::client_builder::ClientBuilder;
20use crate::query::execution::check_job_status;
21use crate::query::{Query as QueryHandle, Result as QueryResult};
22use google_cloud_bigquery_v2::client::JobService;
23use google_cloud_bigquery_v2::model::JobReference;
24use google_cloud_gax::client_builder::Result as BuilderResult;
25use std::sync::Arc;
26
27/// A high-level BigQuery client for executing queries and managing jobs.
28///
29/// # Configuration
30///
31/// To configure a `BigQuery` client, use the `with_*` methods on the [`ClientBuilder`] returned
32/// by [`BigQuery::builder()`]. The default configuration uses Application Default Credentials (ADC)
33/// and connects to the global default endpoint, which works for most applications.
34///
35/// Common configuration customizations include:
36///
37/// - [`with_project_id()`][ClientBuilder::with_project_id]: Sets the default Google Cloud project ID for the client.
38/// - [`with_endpoint()`][ClientBuilder::with_endpoint]: Overrides the default API endpoint (`https://bigquery.googleapis.com`). Useful when testing against mock servers or running in restricted network environments (for example, with VPC Service Controls).
39/// - [`with_credentials()`][ClientBuilder::with_credentials]: Overrides the default Application Default Credentials with explicit or custom authentication credentials.
40///
41/// # Pooling and Cloning
42///
43/// `BigQuery` holds an internal gRPC/HTTP client and connection pool wrapped in an [`Arc`].
44/// You should create a single `BigQuery` client instance upon application initialization and reuse it across multiple tasks or requests.
45/// Cloning a `BigQuery` instance is cheap and does not duplicate underlying connections or thread pools, so you do not need to wrap `BigQuery` in an additional `Arc`.
46///
47/// # Example: Basic Setup and Query Execution
48///
49/// ```
50/// # use google_cloud_bigquery::client::BigQuery;
51/// # async fn sample() -> anyhow::Result<()> {
52/// let client = BigQuery::builder().build().await?;
53/// let mut rows = client
54///     .query("SELECT name, number FROM `bigquery-public-data.usa_names.usa_1910_2013` WHERE state = 'WA' ORDER BY number DESC LIMIT 5")
55///     .with_project_id("my-project-id")
56///     .until_done()
57///     .await?
58///     .read();
59///
60/// while let Some(row) = rows.next().await.transpose()? {
61///     let name: String = row.get("name")?;
62///     let number: i64 = row.get("number")?;
63///     println!("{name}: {number}");
64/// }
65/// # Ok(()) }
66/// ```
67#[derive(Clone, Debug)]
68pub struct BigQuery {
69    job_service: Arc<JobService>,
70    // TODO(#7037): Thread read_client through Query and CompleteQuery.
71    read_client: Option<Arc<Read>>,
72    project_id: Option<String>,
73}
74
75pub(super) mod info {
76    pub(crate) const NAME: &str = env!("CARGO_PKG_NAME");
77    pub(crate) const VERSION: &str = env!("CARGO_PKG_VERSION");
78}
79
80impl BigQuery {
81    /// Returns a new [`ClientBuilder`] for configuring and instantiating a [`BigQuery`] client.
82    ///
83    /// # Example
84    /// ```
85    /// # use google_cloud_bigquery::client::BigQuery;
86    /// # async fn sample() -> anyhow::Result<()> {
87    /// let client = BigQuery::builder()
88    ///     .with_endpoint("https://bigquery.googleapis.com")
89    ///     .build()
90    ///     .await?;
91    /// # Ok(()) }
92    /// ```
93    pub fn builder() -> ClientBuilder {
94        ClientBuilder::new()
95    }
96
97    pub(crate) async fn new(builder: ClientBuilder) -> BuilderResult<Self> {
98        let mut job_service_builder = JobService::builder();
99        if let Some(creds) = builder.config.cred.clone() {
100            job_service_builder = job_service_builder.with_credentials(creds);
101        }
102        if let Some(endpoint) = builder.config.endpoint {
103            job_service_builder = job_service_builder.with_endpoint(endpoint);
104        }
105        if let Some(universe_domain) = builder.config.universe_domain.clone() {
106            job_service_builder = job_service_builder.with_universe_domain(universe_domain);
107        }
108        if builder.config.tracing {
109            job_service_builder = job_service_builder.with_tracing();
110        }
111        let retry_policy = builder
112            .config
113            .retry_policy
114            .unwrap_or_else(crate::query::retry_policy::default_retry_policy);
115        job_service_builder = job_service_builder.with_retry_policy(retry_policy);
116
117        let backoff_policy = builder
118            .config
119            .backoff_policy
120            .unwrap_or_else(crate::query::retry_policy::default_backoff_policy);
121        job_service_builder = job_service_builder.with_backoff_policy(backoff_policy);
122        job_service_builder =
123            job_service_builder.with_retry_throttler(builder.config.retry_throttler);
124
125        job_service_builder =
126            job_service_builder.with_extension(gaxi::api_header::XGoogApiClient {
127                name: info::NAME,
128                version: info::VERSION,
129                library_type: gaxi::api_header::GCCL,
130            });
131
132        let job_service = Arc::new(job_service_builder.build().await?);
133
134        let read_client = if builder.storage_read_enabled {
135            let mut read_builder = Read::builder();
136            if let Some(creds) = builder.config.cred {
137                read_builder = read_builder.with_credentials(creds);
138            }
139            if let Some(endpoint) = builder.storage_read_endpoint {
140                read_builder = read_builder.with_endpoint(endpoint);
141            }
142            if let Some(universe_domain) = builder.config.universe_domain {
143                read_builder = read_builder.with_universe_domain(universe_domain);
144            }
145            if builder.config.tracing {
146                read_builder = read_builder.with_tracing();
147            }
148            Some(Arc::new(read_builder.build().await?))
149        } else {
150            None
151        };
152
153        Ok(BigQuery {
154            job_service,
155            read_client,
156            project_id: builder.project_id,
157        })
158    }
159
160    /// Returns `true` if BigQuery Storage Read API acceleration is enabled on this client.
161    #[cfg_attr(not(test), allow(dead_code))]
162    pub(crate) fn is_storage_read_enabled(&self) -> bool {
163        self.read_client.is_some()
164    }
165
166    /// Creates a request builder to configure and execute a SQL query.
167    ///
168    /// This method returns a [`Query`] builder used to set parameters, specify options,
169    /// and execute the query.
170    ///
171    /// If you configured a default project ID on the client via
172    /// [`ClientBuilder::with_project_id`],
173    /// the returned query builder inherits it automatically.
174    ///
175    /// Call [`Query::send()`] to start query execution, or [`Query::until_done()`]
176    /// to start execution and wait for results.
177    ///
178    /// # Example
179    ///
180    /// ```
181    /// # async fn sample() -> anyhow::Result<()> {
182    /// use google_cloud_bigquery::client::BigQuery;
183    ///
184    /// let client = BigQuery::builder().build().await?;
185    ///
186    /// // Execute a query and read the resulting rows.
187    /// let mut rows = client
188    ///     .query("SELECT name, count FROM `my-project.my_dataset.stats` LIMIT 50")
189    ///     .with_project_id("my-project-id")
190    ///     .set_location("US")
191    ///     .until_done()
192    ///     .await?
193    ///     .read();
194    ///
195    /// while let Some(row) = rows.next().await.transpose()? {
196    ///     let name: String = row.get("name")?;
197    ///     let count: i64 = row.get("count")?;
198    ///     println!("{name}: {count}");
199    /// }
200    /// # Ok(())
201    /// # }
202    /// ```
203    pub fn query<S: Into<String>>(&self, sql: S) -> Query {
204        let builder = Query::new(self.job_service.clone(), sql.into());
205        self.project_id
206            .as_deref()
207            .into_iter()
208            .fold(builder, |builder, project_id| {
209                builder.with_project_id(project_id)
210            })
211    }
212
213    /// Binds an existing out-of-process query job reference to a high-level [`Query`](QueryHandle) handle.
214    ///
215    /// Fetches the job metadata via [`JobService::get_job`] and initializes a
216    /// [`Query`](QueryHandle) handle.
217    /// If `job_ref.project_id` is empty, it defaults to the client's billing project ID.
218    ///
219    /// # Arguments
220    /// * `job_ref` - A [`JobReference`] identifying the job to attach to.
221    ///
222    /// # Example
223    /// ```no_run
224    /// # use google_cloud_bigquery::client::BigQuery;
225    /// # use google_cloud_bigquery_v2::model::JobReference;
226    /// # async fn sample(client: &BigQuery) -> anyhow::Result<()> {
227    /// let job_ref = JobReference::new()
228    ///     .set_project_id("my-project")
229    ///     .set_job_id("my_job_id")
230    ///     .set_location("us-central1");
231    /// let query = client.attach_job(job_ref).await?;
232    /// let mut results = query.until_done().await?.read();
233    /// while let Some(row) = results.next().await {
234    ///     let row = row?;
235    ///     // process row
236    /// }
237    /// # Ok(())
238    /// # }
239    /// ```
240    pub async fn attach_job(&self, mut job_ref: JobReference) -> QueryResult<QueryHandle> {
241        if job_ref.project_id.is_empty()
242            && let Some(proj) = &self.project_id
243        {
244            job_ref.project_id = proj.clone();
245        }
246
247        let req = self
248            .job_service
249            .get_job()
250            .set_job_id(job_ref.job_id.clone())
251            .set_project_id(job_ref.project_id.clone());
252
253        let req = job_ref
254            .location
255            .clone()
256            .into_iter()
257            .fold(req, |req, location| req.set_location(location));
258
259        let job = req.send().await?;
260
261        let is_query = job
262            .configuration
263            .as_ref()
264            .and_then(|c| c.query.as_ref())
265            .is_some();
266        if !is_query {
267            return Err(QueryError::UnsupportedJobType);
268        }
269
270        Ok(QueryHandle::from_job(
271            self.job_service.clone(),
272            check_job_status(job)?,
273            None,
274            None,
275            RequestOptions::default(),
276        ))
277    }
278}
279
280#[cfg(test)]
281mod tests {
282    use super::BigQuery;
283    use crate::error::QueryError;
284    use crate::query::tests::{MockJobService, create_job_service};
285    use google_cloud_auth::credentials::anonymous::Builder as Anonymous;
286    use google_cloud_bigquery_v2::client::JobService;
287    use google_cloud_bigquery_v2::model::{
288        Job, JobConfiguration, JobConfigurationQuery, JobReference,
289    };
290    use google_cloud_gax::response::Response;
291    use std::sync::Arc;
292
293    impl BigQuery {
294        fn from_job_service(job_service: Arc<JobService>, project_id: Option<String>) -> Self {
295            Self {
296                job_service,
297                read_client: None,
298                project_id,
299            }
300        }
301    }
302
303    #[tokio::test]
304    async fn test_bigquery_builder() -> anyhow::Result<()> {
305        let client = BigQuery::builder()
306            .with_credentials(Anonymous::new().build())
307            .build()
308            .await?;
309        assert!(client.project_id.is_none());
310        assert!(!client.is_storage_read_enabled());
311        Ok(())
312    }
313
314    #[cfg(google_cloud_unstable_bigquery_storage_read)]
315    #[tokio::test]
316    async fn test_bigquery_storage_read_forwards_config() -> anyhow::Result<()> {
317        use google_cloud_auth::credentials::mds;
318
319        let creds = mds::Builder::default()
320            .with_universe_domain("my-universe.com")
321            .build()?;
322        let client = BigQuery::builder()
323            .with_storage_read(true)
324            .with_storage_read_endpoint("https://custom.storage.endpoint:1234")
325            .with_credentials(creds)
326            .with_universe_domain("my-universe.com")
327            .with_tracing()
328            .build()
329            .await?;
330        assert!(client.is_storage_read_enabled());
331        let debug_str = format!("{:?}", client.read_client);
332        assert!(
333            debug_str.contains("custom.storage.endpoint"),
334            "expected custom storage endpoint in read_client debug output, got: {debug_str}"
335        );
336        assert!(
337            debug_str.contains("my-universe.com"),
338            "expected universe domain in read_client debug output, got: {debug_str}"
339        );
340        assert!(
341            debug_str.contains("tracing_attributes: Some"),
342            "expected tracing attributes in read_client debug output, got: {debug_str}"
343        );
344
345        Ok(())
346    }
347
348    #[tokio::test]
349    async fn test_bigquery_builder_with_project_id() -> anyhow::Result<()> {
350        let client = BigQuery::builder()
351            .with_project_id("test-proj")
352            .with_credentials(Anonymous::new().build())
353            .build()
354            .await?;
355        assert_eq!(client.project_id.as_deref(), Some("test-proj"));
356        Ok(())
357    }
358
359    #[tokio::test]
360    async fn test_bigquery_query_inherits_project_id() -> anyhow::Result<()> {
361        let client = BigQuery::builder()
362            .with_project_id("test-proj")
363            .with_credentials(Anonymous::new().build())
364            .build()
365            .await?;
366        let query_builder = client.query("SELECT 1");
367        assert_eq!(query_builder.project_id.as_deref(), Some("test-proj"));
368        assert_eq!(
369            query_builder.options.quota_project().as_deref(),
370            Some("test-proj")
371        );
372        Ok(())
373    }
374
375    #[tokio::test]
376    async fn test_bigquery_query_overrides_project_id() -> anyhow::Result<()> {
377        let client = BigQuery::builder()
378            .with_project_id("test-proj")
379            .with_credentials(Anonymous::new().build())
380            .build()
381            .await?;
382        let query_builder = client.query("SELECT 1").with_project_id("override-proj");
383        assert_eq!(query_builder.project_id.as_deref(), Some("override-proj"));
384        assert_eq!(
385            query_builder.options.quota_project().as_deref(),
386            Some("override-proj")
387        );
388        Ok(())
389    }
390
391    #[tokio::test]
392    async fn test_bigquery_query_without_project_id() -> anyhow::Result<()> {
393        let client = BigQuery::builder()
394            .with_credentials(Anonymous::new().build())
395            .build()
396            .await?;
397        let query_builder = client.query("SELECT 1");
398        assert!(query_builder.project_id.is_none());
399        assert!(query_builder.options.quota_project().is_none());
400        Ok(())
401    }
402
403    #[tokio::test]
404    async fn test_bigquery_attach_job() -> anyhow::Result<()> {
405        let mut mock = MockJobService::new();
406        mock.expect_get_job().returning(|req, _| {
407            assert_eq!(req.project_id, "test-proj");
408            assert_eq!(req.job_id, "job_123");
409            let job = Job::new()
410                .set_job_reference(
411                    JobReference::new()
412                        .set_project_id("test-proj")
413                        .set_job_id("job_123"),
414                )
415                .set_configuration(
416                    JobConfiguration::new()
417                        .set_query(JobConfigurationQuery::new().set_query("SELECT 1")),
418                );
419            Ok(Response::from(job))
420        });
421        let client = BigQuery::from_job_service(create_job_service(mock), None);
422        let job_ref = JobReference::new()
423            .set_project_id("test-proj")
424            .set_job_id("job_123");
425        let query = client.attach_job(job_ref).await?;
426        let job_ref = query
427            .metadata()
428            .job_reference
429            .as_ref()
430            .expect("job_reference should be set");
431        assert_eq!(job_ref.project_id, "test-proj");
432        assert_eq!(job_ref.job_id, "job_123");
433        Ok(())
434    }
435
436    #[tokio::test]
437    async fn test_bigquery_attach_job_inherits_project_id() -> anyhow::Result<()> {
438        let mut mock = MockJobService::new();
439        mock.expect_get_job().returning(|req, _| {
440            assert_eq!(req.project_id, "client-proj");
441            assert_eq!(req.job_id, "job_456");
442            let job = Job::new()
443                .set_job_reference(
444                    JobReference::new()
445                        .set_project_id("client-proj")
446                        .set_job_id("job_456"),
447                )
448                .set_configuration(
449                    JobConfiguration::new()
450                        .set_query(JobConfigurationQuery::new().set_query("SELECT 1")),
451                );
452            Ok(Response::from(job))
453        });
454        let client =
455            BigQuery::from_job_service(create_job_service(mock), Some("client-proj".to_string()));
456        let job_ref = JobReference::new().set_job_id("job_456");
457        let query = client.attach_job(job_ref).await?;
458        let job_ref = query
459            .metadata()
460            .job_reference
461            .as_ref()
462            .expect("job_reference should be set");
463        assert_eq!(job_ref.project_id, "client-proj");
464        assert_eq!(job_ref.job_id, "job_456");
465        Ok(())
466    }
467
468    #[tokio::test]
469    async fn test_bigquery_attach_job_missing_project_id() -> anyhow::Result<()> {
470        let client = BigQuery::builder()
471            .with_credentials(Anonymous::new().build())
472            .build()
473            .await?;
474        let job_ref = JobReference::new().set_job_id("job_789");
475        let err = client
476            .attach_job(job_ref)
477            .await
478            .expect_err("should return an error when project_id is missing");
479        assert!(
480            matches!(&err, QueryError::Rpc { source } if source.is_binding()),
481            "expected Binding error for missing project ID, got {err:?}"
482        );
483        Ok(())
484    }
485
486    #[tokio::test]
487    async fn test_bigquery_attach_job_empty_job_id() -> anyhow::Result<()> {
488        let client = BigQuery::builder()
489            .with_project_id("client-proj")
490            .with_credentials(Anonymous::new().build())
491            .build()
492            .await?;
493        let job_ref = JobReference::new();
494        let err = client
495            .attach_job(job_ref)
496            .await
497            .expect_err("should return an error when job_id is empty");
498        assert!(
499            matches!(&err, QueryError::Rpc { source } if source.is_binding()),
500            "expected Binding error for empty job ID, got {err:?}"
501        );
502        Ok(())
503    }
504
505    #[tokio::test]
506    async fn test_bigquery_attach_job_unsupported_job_type() -> anyhow::Result<()> {
507        let mut mock = MockJobService::new();
508        mock.expect_get_job().returning(|_, _| {
509            let job = Job::new().set_configuration(JobConfiguration::new());
510            Ok(Response::from(job))
511        });
512        let client =
513            BigQuery::from_job_service(create_job_service(mock), Some("client-proj".to_string()));
514        let job_ref = JobReference::new().set_job_id("job_extract");
515        let err = client
516            .attach_job(job_ref)
517            .await
518            .expect_err("should return an error for non-query job");
519        assert!(
520            matches!(&err, QueryError::UnsupportedJobType),
521            "expected UnsupportedJobType, got {err:?}"
522        );
523        Ok(())
524    }
525
526    #[tokio::test]
527    async fn test_bigquery_attach_job_failed_job() -> anyhow::Result<()> {
528        use google_cloud_bigquery_v2::model::{ErrorProto, JobConfigurationQuery, JobStatus};
529
530        let mut mock = MockJobService::new();
531        mock.expect_get_job().returning(|_, _| {
532            let err_proto = ErrorProto::new()
533                .set_reason("invalidQuery")
534                .set_message("Syntax error");
535            let job = Job::new()
536                .set_configuration(
537                    JobConfiguration::new()
538                        .set_query(JobConfigurationQuery::new().set_query("SELECT * FROM")),
539                )
540                .set_status(
541                    JobStatus::new()
542                        .set_state("DONE")
543                        .set_error_result(err_proto.clone())
544                        .set_errors(vec![err_proto]),
545                );
546            Ok(Response::from(job))
547        });
548        let client =
549            BigQuery::from_job_service(create_job_service(mock), Some("client-proj".to_string()));
550        let job_ref = JobReference::new().set_job_id("job_failed");
551        let err = client
552            .attach_job(job_ref)
553            .await
554            .expect_err("should return an error for failed query job");
555        assert!(
556            matches!(&err, QueryError::JobFailed { errors } if errors.len() == 1 && errors[0].reason == "invalidQuery"),
557            "expected JobFailed, got {err:?}"
558        );
559        Ok(())
560    }
561
562    #[tokio::test]
563    async fn test_bigquery_calls_send_veneer_header_not_gapic() -> anyhow::Result<()> {
564        use httptest::{Expectation, Server, all_of, matchers::*, responders::*};
565        use serde_json::json;
566
567        let server = Server::run();
568        server.expect(
569            Expectation::matching(all_of![
570                request::method_path("GET", "/bigquery/v2/projects/test-proj/jobs/job_123"),
571                request::headers(contains((
572                    "x-goog-api-client",
573                    matches(format!("gccl/{}", env!("CARGO_PKG_VERSION"))),
574                ))),
575                not(request::headers(contains((
576                    "x-goog-api-client",
577                    matches("gapic/"),
578                )))),
579            ])
580            .respond_with(json_encoded(json!({
581                "jobReference": {
582                    "projectId": "test-proj",
583                    "jobId": "job_123"
584                },
585                "configuration": {
586                    "query": {
587                        "query": "SELECT 1"
588                    }
589                },
590                "status": {
591                    "state": "DONE"
592                }
593            }))),
594        );
595
596        let client = BigQuery::builder()
597            .with_endpoint(server.url_str(""))
598            .with_credentials(Anonymous::new().build())
599            .with_project_id("test-proj")
600            .build()
601            .await?;
602
603        let job_ref = JobReference::new().set_job_id("job_123");
604        let query = client.attach_job(job_ref).await?;
605        let metadata = query.metadata();
606        assert_eq!(
607            metadata.job_reference.as_ref().map(|j| j.job_id.as_str()),
608            Some("job_123")
609        );
610
611        Ok(())
612    }
613}
google_cloud_bigquery/query/client.rs

google_cloud_bigquery/query/
client.rs
