1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
//! [<img alt="github" src="https://img.shields.io/badge/github-lquerel/gcp_bigquery_client-8da0cb?style=for-the-badge&labelColor=555555&logo=github" height="20">](https://github.com/lquerel/gcp-bigquery-client)
//! [<img alt="crates.io" src="https://img.shields.io/crates/v/gcp_bigquery_client.svg?style=for-the-badge&color=fc8d62&logo=rust" height="20">](https://crates.io/crates/gcp-bigquery-client)
//! [<img alt="build status" src="https://img.shields.io/github/workflow/status/lquerel/gcp-bigquery-client/Rust/main?style=for-the-badge" height="20">](https://github.com/lquerel/gcp-bigquery-client/actions?query=branch%3Amain)
//!
//! An ergonomic async client library for GCP BigQuery.
//! * Support for dataset, table, streaming API and query (see [status section](#status) for an exhaustive list of supported API endpoints)
//! * Support Service Account Key authentication (other OAuth flows will be added later)
//! * Create tables and rows via builder patterns
//! * Persist complex Rust structs in structured BigQuery tables
//! * Async API
//!
//! <br>
//!
//! Other OAuth flows will be added later.
//!
//! For a detailed tutorial on the different ways to use GCP BigQuery Client please check out the [GitHub repository](https://github.com/lquerel/gcp-bigquery-client).
#[macro_use]
extern crate serde;
extern crate serde_json;

use std::env;
use std::path::PathBuf;
use std::sync::Arc;

use client_builder::ClientBuilder;
use reqwest::Response;
use serde::Deserialize;
use yup_oauth2::ServiceAccountKey;

use crate::auth::Authenticator;
use crate::dataset::DatasetApi;
use crate::error::BQError;
use crate::job::JobApi;
use crate::model_api::ModelApi;
use crate::project::ProjectApi;
use crate::routine::RoutineApi;
use crate::table::TableApi;
use crate::tabledata::TableDataApi;

/// Since yup_oauth2 structs are used as parameters in public functions there is already semver
/// coupling, as it is an error if consumer uses different version of yup_oauth than gcp-bigquery-client
/// Export yup_oauth2 so consumers don't need to carefully keep their dependency versions in sync.
/// (see https://github.com/lquerel/gcp-bigquery-client/pull/86)
pub use yup_oauth2;

pub mod auth;
pub mod client_builder;
pub mod dataset;
pub mod error;
pub mod job;
pub mod model;
pub mod model_api;
pub mod project;
pub mod routine;
pub mod table;
pub mod tabledata;

const BIG_QUERY_V2_URL: &str = "https://bigquery.googleapis.com/bigquery/v2";
const BIG_QUERY_AUTH_URL: &str = "https://www.googleapis.com/auth/bigquery";

/// An asynchronous BigQuery client.
#[derive(Clone)]
pub struct Client {
    dataset_api: DatasetApi,
    table_api: TableApi,
    job_api: JobApi,
    tabledata_api: TableDataApi,
    routine_api: RoutineApi,
    model_api: ModelApi,
    project_api: ProjectApi,
}

impl Client {
    pub fn from_authenticator(auth: Arc<dyn Authenticator>) -> Self {
        let client = reqwest::Client::new();
        Self {
            dataset_api: DatasetApi::new(client.clone(), Arc::clone(&auth)),
            table_api: TableApi::new(client.clone(), Arc::clone(&auth)),
            job_api: JobApi::new(client.clone(), Arc::clone(&auth)),
            tabledata_api: TableDataApi::new(client.clone(), Arc::clone(&auth)),
            routine_api: RoutineApi::new(client.clone(), Arc::clone(&auth)),
            model_api: ModelApi::new(client.clone(), Arc::clone(&auth)),
            project_api: ProjectApi::new(client, auth),
        }
    }

    /// Constructs a new BigQuery client.
    /// # Argument
    /// * `sa_key_file` - A GCP Service Account Key file.
    pub async fn from_service_account_key_file(sa_key_file: &str) -> Result<Self, BQError> {
        ClientBuilder::new()
            .build_from_service_account_key_file(sa_key_file)
            .await
    }

    /// Constructs a new BigQuery client from a [`ServiceAccountKey`].
    /// # Argument
    /// * `sa_key` - A GCP Service Account Key `yup-oauth2` object.
    /// * `readonly` - A boolean setting whether the acquired token scope should be readonly.
    ///
    /// [`ServiceAccountKey`]: https://docs.rs/yup-oauth2/*/yup_oauth2/struct.ServiceAccountKey.html
    pub async fn from_service_account_key(sa_key: ServiceAccountKey, readonly: bool) -> Result<Self, BQError> {
        ClientBuilder::new()
            .build_from_service_account_key(sa_key, readonly)
            .await
    }

    pub async fn with_workload_identity(readonly: bool) -> Result<Self, BQError> {
        ClientBuilder::new().build_with_workload_identity(readonly).await
    }

    pub(crate) fn v2_base_url(&mut self, base_url: String) -> &mut Self {
        self.dataset_api.with_base_url(base_url.clone());
        self.table_api.with_base_url(base_url.clone());
        self.job_api.with_base_url(base_url.clone());
        self.tabledata_api.with_base_url(base_url.clone());
        self.routine_api.with_base_url(base_url.clone());
        self.model_api.with_base_url(base_url.clone());
        self.project_api.with_base_url(base_url);
        self
    }

    pub async fn from_installed_flow_authenticator<S: AsRef<[u8]>, P: Into<PathBuf>>(
        secret: S,
        persistant_file_path: P,
    ) -> Result<Self, BQError> {
        ClientBuilder::new()
            .build_from_installed_flow_authenticator(secret, persistant_file_path)
            .await
    }

    pub async fn from_installed_flow_authenticator_from_secret_file<P: Into<PathBuf>>(
        secret_file: &str,
        persistant_file_path: P,
    ) -> Result<Self, BQError> {
        Self::from_installed_flow_authenticator(
            tokio::fs::read(secret_file)
                .await
                .expect("expecting a valid secret file."),
            persistant_file_path,
        )
        .await
    }

    pub async fn from_application_default_credentials() -> Result<Self, BQError> {
        ClientBuilder::new().build_from_application_default_credentials().await
    }

    pub async fn from_authorized_user_secret(secret: &str) -> Result<Self, BQError> {
        ClientBuilder::new()
            .build_from_authorized_user_authenticator(secret)
            .await
    }

    /// Returns a dataset API handler.
    pub fn dataset(&self) -> &DatasetApi {
        &self.dataset_api
    }

    /// Returns a table API handler.
    pub fn table(&self) -> &TableApi {
        &self.table_api
    }

    /// Returns a job API handler.
    pub fn job(&self) -> &JobApi {
        &self.job_api
    }

    /// Returns a table data API handler.
    pub fn tabledata(&self) -> &TableDataApi {
        &self.tabledata_api
    }

    /// Returns a routine API handler.
    pub fn routine(&self) -> &RoutineApi {
        &self.routine_api
    }

    /// Returns a model API handler.
    pub fn model(&self) -> &ModelApi {
        &self.model_api
    }

    /// Returns a project API handler.
    pub fn project(&self) -> &ProjectApi {
        &self.project_api
    }
}

pub(crate) fn urlencode<T: AsRef<str>>(s: T) -> String {
    url::form_urlencoded::byte_serialize(s.as_ref().as_bytes()).collect()
}

async fn process_response<T: for<'de> Deserialize<'de>>(resp: Response) -> Result<T, BQError> {
    if resp.status().is_success() {
        Ok(resp.json().await?)
    } else {
        Err(BQError::ResponseError {
            error: resp.json().await?,
        })
    }
}

pub fn env_vars() -> (String, String, String, String) {
    let project_id = env::var("PROJECT_ID").expect("Environment variable PROJECT_ID");
    let dataset_id = env::var("DATASET_ID").expect("Environment variable DATASET_ID");
    let table_id = env::var("TABLE_ID").expect("Environment variable TABLE_ID");
    let gcp_sa_key =
        env::var("GOOGLE_APPLICATION_CREDENTIALS").expect("Environment variable GOOGLE_APPLICATION_CREDENTIALS");

    (project_id, dataset_id, table_id, gcp_sa_key)
}