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::builder::bigquery::Query;
16use crate::error::QueryError;
17use crate::query::client_builder::ClientBuilder;
18use crate::query::{Query as QueryHandle, Result as QueryResult};
19use google_cloud_bigquery_v2::client::JobService;
20use google_cloud_bigquery_v2::model::JobReference;
21use google_cloud_gax::client_builder::Result as BuilderResult;
22use std::sync::Arc;
23
24/// A high-level BigQuery client for executing queries and managing jobs.
25///
26/// # Configuration
27///
28/// To configure a `BigQuery` client, use the `with_*` methods on the [`ClientBuilder`] returned
29/// by [`BigQuery::builder()`]. The default configuration uses Application Default Credentials (ADC)
30/// and connects to the global default endpoint, which works for most applications.
31///
32/// Common configuration customizations include:
33///
34/// - [`with_project_id()`][ClientBuilder::with_project_id]: Sets the default Google Cloud project ID for the client.
35/// - [`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).
36/// - [`with_credentials()`][ClientBuilder::with_credentials]: Overrides the default Application Default Credentials with explicit or custom authentication credentials.
37///
38/// # Pooling and Cloning
39///
40/// `BigQuery` holds an internal gRPC/HTTP client and connection pool wrapped in an [`Arc`].
41/// You should create a single `BigQuery` client instance upon application initialization and reuse it across multiple tasks or requests.
42/// 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`.
43///
44/// # Example: Basic Setup and Query Execution
45///
46/// ```
47/// # use google_cloud_bigquery::client::BigQuery;
48/// # async fn sample() -> anyhow::Result<()> {
49/// let client = BigQuery::builder().build().await?;
50/// let mut rows = client
51///     .query("SELECT name, count FROM `bigquery-public-data.usa_names.usa_1910_2013` WHERE state = 'WA' ORDER BY count DESC LIMIT 5")
52///     .with_project_id("my-project-id")
53///     .until_done()
54///     .await?
55///     .read();
56///
57/// while let Some(row) = rows.next().await.transpose()? {
58///     let name: String = row.get("name");
59///     let count: i64 = row.get("count");
60///     println!("{name}: {count}");
61/// }
62/// # Ok(()) }
63/// ```
64#[derive(Clone, Debug)]
65pub struct BigQuery {
66    job_service: Arc<JobService>,
67    project_id: Option<String>,
68}
69
70impl BigQuery {
71    /// Returns a new [`ClientBuilder`] for configuring and instantiating a [`BigQuery`] client.
72    ///
73    /// # Example
74    /// ```
75    /// # use google_cloud_bigquery::client::BigQuery;
76    /// # async fn sample() -> anyhow::Result<()> {
77    /// let client = BigQuery::builder()
78    ///     .with_endpoint("https://bigquery.googleapis.com")
79    ///     .build()
80    ///     .await?;
81    /// # Ok(()) }
82    /// ```
83    pub fn builder() -> ClientBuilder {
84        ClientBuilder::new()
85    }
86
87    pub(crate) async fn new(builder: ClientBuilder) -> BuilderResult<Self> {
88        let mut job_service_builder = JobService::builder();
89        if let Some(creds) = builder.config.cred {
90            job_service_builder = job_service_builder.with_credentials(creds);
91        }
92        if let Some(endpoint) = builder.config.endpoint {
93            job_service_builder = job_service_builder.with_endpoint(endpoint);
94        }
95        if let Some(universe_domain) = builder.config.universe_domain {
96            job_service_builder = job_service_builder.with_universe_domain(universe_domain);
97        }
98        if builder.config.tracing {
99            job_service_builder = job_service_builder.with_tracing();
100        }
101        if let Some(retry_policy) = builder.config.retry_policy {
102            job_service_builder = job_service_builder.with_retry_policy(retry_policy);
103        }
104        if let Some(backoff_policy) = builder.config.backoff_policy {
105            job_service_builder = job_service_builder.with_backoff_policy(backoff_policy);
106        }
107        job_service_builder =
108            job_service_builder.with_retry_throttler(builder.config.retry_throttler);
109        let job_service = Arc::new(job_service_builder.build().await?);
110
111        Ok(BigQuery {
112            job_service,
113            project_id: builder.project_id,
114        })
115    }
116
117    /// Creates a request builder to configure and execute a SQL query.
118    ///
119    /// This method returns a [`Query`] builder used to set parameters, specify options,
120    /// and execute the query.
121    ///
122    /// If you configured a default project ID on the client via
123    /// [`ClientBuilder::with_project_id`],
124    /// the returned query builder inherits it automatically.
125    ///
126    /// Call [`Query::send()`] to start query execution, or [`Query::until_done()`]
127    /// to start execution and wait for results.
128    ///
129    /// # Example
130    ///
131    /// ```
132    /// # async fn sample() -> anyhow::Result<()> {
133    /// use google_cloud_bigquery::client::BigQuery;
134    ///
135    /// let client = BigQuery::builder().build().await?;
136    ///
137    /// // Execute a query and read the resulting rows.
138    /// let mut rows = client
139    ///     .query("SELECT name, count FROM `my-project.my_dataset.stats` LIMIT 50")
140    ///     .with_project_id("my-project-id")
141    ///     .set_location("US")
142    ///     .until_done()
143    ///     .await?
144    ///     .read();
145    ///
146    /// while let Some(row) = rows.next().await.transpose()? {
147    ///     let name: String = row.get("name");
148    ///     let count: i64 = row.get("count");
149    ///     println!("{name}: {count}");
150    /// }
151    /// # Ok(())
152    /// # }
153    /// ```
154    pub fn query<S: Into<String>>(&self, sql: S) -> Query {
155        let builder = Query::new(self.job_service.clone(), sql.into());
156        self.project_id
157            .as_deref()
158            .into_iter()
159            .fold(builder, |builder, project_id| {
160                builder.with_project_id(project_id)
161            })
162    }
163
164    /// Binds an existing out-of-process query job reference to a high-level [`Query`](QueryHandle) handle.
165    ///
166    /// Fetches the job metadata via [`JobService::get_job`] and initializes a
167    /// [`Query`](QueryHandle) handle.
168    /// If `job_ref.project_id` is empty, it defaults to the client's billing project ID.
169    ///
170    /// # Arguments
171    /// * `job_ref` - A [`JobReference`] identifying the job to attach to.
172    ///
173    /// # Example
174    /// ```no_run
175    /// # use google_cloud_bigquery::client::BigQuery;
176    /// # use google_cloud_bigquery_v2::model::JobReference;
177    /// # async fn sample(client: &BigQuery) -> anyhow::Result<()> {
178    /// let job_ref = JobReference::new()
179    ///     .set_project_id("my-project")
180    ///     .set_job_id("my_job_id")
181    ///     .set_location("us-central1");
182    /// let query = client.attach_job(job_ref).await?;
183    /// let mut results = query.until_done().await?.read();
184    /// while let Some(row) = results.next().await {
185    ///     let row = row?;
186    ///     // process row
187    /// }
188    /// # Ok(())
189    /// # }
190    /// ```
191    pub async fn attach_job(&self, mut job_ref: JobReference) -> QueryResult<QueryHandle> {
192        if job_ref.project_id.is_empty()
193            && let Some(proj) = &self.project_id
194        {
195            job_ref.project_id = proj.clone();
196        }
197
198        let req = self
199            .job_service
200            .get_job()
201            .set_job_id(job_ref.job_id.clone())
202            .set_project_id(job_ref.project_id.clone());
203
204        let req = job_ref
205            .location
206            .clone()
207            .into_iter()
208            .fold(req, |req, location| req.set_location(location));
209
210        let job = req.send().await?;
211
212        let is_query = job
213            .configuration
214            .as_ref()
215            .and_then(|c| c.query.as_ref())
216            .is_some();
217        if !is_query {
218            return Err(QueryError::UnsupportedJobType);
219        }
220
221        Ok(QueryHandle::from_job(
222            self.job_service.clone(),
223            job,
224            None,
225            None,
226        ))
227    }
228}
229
230#[cfg(test)]
231mod tests {
232    use super::BigQuery;
233    use crate::error::QueryError;
234    use crate::query::tests::{MockJobService, create_job_service};
235    use google_cloud_auth::credentials::anonymous::Builder as Anonymous;
236    use google_cloud_bigquery_v2::client::JobService;
237    use google_cloud_bigquery_v2::model::{
238        Job, JobConfiguration, JobConfigurationQuery, JobReference,
239    };
240    use google_cloud_gax::response::Response;
241    use std::sync::Arc;
242
243    impl BigQuery {
244        fn from_job_service(job_service: Arc<JobService>, project_id: Option<String>) -> Self {
245            Self {
246                job_service,
247                project_id,
248            }
249        }
250    }
251
252    #[tokio::test]
253    async fn test_bigquery_builder() -> anyhow::Result<()> {
254        let client = BigQuery::builder()
255            .with_credentials(Anonymous::new().build())
256            .build()
257            .await?;
258        assert!(client.project_id.is_none());
259        Ok(())
260    }
261
262    #[tokio::test]
263    async fn test_bigquery_builder_with_project_id() -> anyhow::Result<()> {
264        let client = BigQuery::builder()
265            .with_project_id("test-proj")
266            .with_credentials(Anonymous::new().build())
267            .build()
268            .await?;
269        assert_eq!(client.project_id.as_deref(), Some("test-proj"));
270        Ok(())
271    }
272
273    #[tokio::test]
274    async fn test_bigquery_query_inherits_project_id() -> anyhow::Result<()> {
275        let client = BigQuery::builder()
276            .with_project_id("test-proj")
277            .with_credentials(Anonymous::new().build())
278            .build()
279            .await?;
280        let query_builder = client.query("SELECT 1");
281        assert_eq!(query_builder.project_id.as_deref(), Some("test-proj"));
282        Ok(())
283    }
284
285    #[tokio::test]
286    async fn test_bigquery_query_without_project_id() -> anyhow::Result<()> {
287        let client = BigQuery::builder()
288            .with_credentials(Anonymous::new().build())
289            .build()
290            .await?;
291        let query_builder = client.query("SELECT 1");
292        assert!(query_builder.project_id.is_none());
293        Ok(())
294    }
295
296    #[tokio::test]
297    async fn test_bigquery_attach_job() -> anyhow::Result<()> {
298        let mut mock = MockJobService::new();
299        mock.expect_get_job().returning(|req, _| {
300            assert_eq!(req.project_id, "test-proj");
301            assert_eq!(req.job_id, "job_123");
302            let job = Job::new()
303                .set_job_reference(
304                    JobReference::new()
305                        .set_project_id("test-proj")
306                        .set_job_id("job_123"),
307                )
308                .set_configuration(
309                    JobConfiguration::new()
310                        .set_query(JobConfigurationQuery::new().set_query("SELECT 1")),
311                );
312            Ok(Response::from(job))
313        });
314        let client = BigQuery::from_job_service(create_job_service(mock), None);
315        let job_ref = JobReference::new()
316            .set_project_id("test-proj")
317            .set_job_id("job_123");
318        let query = client.attach_job(job_ref).await?;
319        let job_ref = query
320            .metadata()
321            .job_reference
322            .as_ref()
323            .expect("job_reference should be set");
324        assert_eq!(job_ref.project_id, "test-proj");
325        assert_eq!(job_ref.job_id, "job_123");
326        Ok(())
327    }
328
329    #[tokio::test]
330    async fn test_bigquery_attach_job_inherits_project_id() -> anyhow::Result<()> {
331        let mut mock = MockJobService::new();
332        mock.expect_get_job().returning(|req, _| {
333            assert_eq!(req.project_id, "client-proj");
334            assert_eq!(req.job_id, "job_456");
335            let job = Job::new()
336                .set_job_reference(
337                    JobReference::new()
338                        .set_project_id("client-proj")
339                        .set_job_id("job_456"),
340                )
341                .set_configuration(
342                    JobConfiguration::new()
343                        .set_query(JobConfigurationQuery::new().set_query("SELECT 1")),
344                );
345            Ok(Response::from(job))
346        });
347        let client =
348            BigQuery::from_job_service(create_job_service(mock), Some("client-proj".to_string()));
349        let job_ref = JobReference::new().set_job_id("job_456");
350        let query = client.attach_job(job_ref).await?;
351        let job_ref = query
352            .metadata()
353            .job_reference
354            .as_ref()
355            .expect("job_reference should be set");
356        assert_eq!(job_ref.project_id, "client-proj");
357        assert_eq!(job_ref.job_id, "job_456");
358        Ok(())
359    }
360
361    #[tokio::test]
362    async fn test_bigquery_attach_job_missing_project_id() -> anyhow::Result<()> {
363        let client = BigQuery::builder()
364            .with_credentials(Anonymous::new().build())
365            .build()
366            .await?;
367        let job_ref = JobReference::new().set_job_id("job_789");
368        let err = client
369            .attach_job(job_ref)
370            .await
371            .expect_err("should return an error when project_id is missing");
372        assert!(
373            matches!(&err, QueryError::Rpc { source } if source.is_binding()),
374            "expected Binding error for missing project ID, got {err:?}"
375        );
376        Ok(())
377    }
378
379    #[tokio::test]
380    async fn test_bigquery_attach_job_empty_job_id() -> anyhow::Result<()> {
381        let client = BigQuery::builder()
382            .with_project_id("client-proj")
383            .with_credentials(Anonymous::new().build())
384            .build()
385            .await?;
386        let job_ref = JobReference::new();
387        let err = client
388            .attach_job(job_ref)
389            .await
390            .expect_err("should return an error when job_id is empty");
391        assert!(
392            matches!(&err, QueryError::Rpc { source } if source.is_binding()),
393            "expected Binding error for empty job ID, got {err:?}"
394        );
395        Ok(())
396    }
397
398    #[tokio::test]
399    async fn test_bigquery_attach_job_unsupported_job_type() -> anyhow::Result<()> {
400        let mut mock = MockJobService::new();
401        mock.expect_get_job().returning(|_, _| {
402            let job = Job::new().set_configuration(JobConfiguration::new());
403            Ok(Response::from(job))
404        });
405        let client =
406            BigQuery::from_job_service(create_job_service(mock), Some("client-proj".to_string()));
407        let job_ref = JobReference::new().set_job_id("job_extract");
408        let err = client
409            .attach_job(job_ref)
410            .await
411            .expect_err("should return an error for non-query job");
412        assert!(
413            matches!(&err, QueryError::UnsupportedJobType),
414            "expected UnsupportedJobType, got {err:?}"
415        );
416        Ok(())
417    }
418}
google_cloud_bigquery/query/client.rs

google_cloud_bigquery/query/
client.rs
