feat: add google drive auth url

src/client_builder.rs

@@ -7,18 +7,18 @@ use crate::auth::{
     service_account_authenticator, ServiceAccountAuthenticator,
 };
 use crate::error::BQError;
-use crate::{Client, BIG_QUERY_AUTH_URL, BIG_QUERY_V2_URL};
+use crate::{Client, BIG_QUERY_AUTH_URL, BIG_QUERY_V2_URL, DRIVE_AUTH_URL};
 
 pub struct ClientBuilder {
     v2_base_url: String,
-    auth_base_url: String,
+    auth_base_urls: Vec<String>,
 }
 
 impl ClientBuilder {
     pub fn new() -> Self {
         Self {
             v2_base_url: BIG_QUERY_V2_URL.to_string(),
-            auth_base_url: BIG_QUERY_AUTH_URL.to_string(),
+            auth_base_urls: vec![BIG_QUERY_AUTH_URL.to_string(), DRIVE_AUTH_URL.to_string()],
         }
     }
 
@@ -27,8 +27,8 @@ impl ClientBuilder {
         self
     }
 
-    pub fn with_auth_base_url(&mut self, base_url: String) -> &mut Self {
-        self.auth_base_url = base_url;
+    pub fn with_auth_base_url(&mut self, base_urls: Vec<String>) -> &mut Self {
+        self.auth_base_urls = base_urls;
         self
     }
 
@@ -37,20 +37,29 @@ impl ClientBuilder {
         sa_key: ServiceAccountKey,
         readonly: bool,
     ) -> Result<Client, BQError> {
-        let scope = if readonly {
-            format!("{}.readonly", self.auth_base_url)
-        } else {
-            self.auth_base_url.clone()
-        };
-        let sa_auth = ServiceAccountAuthenticator::from_service_account_key(sa_key, &[&scope]).await?;
+        let scopes: Vec<String> = self
+            .auth_base_urls
+            .iter()
+            .map(|auth_base_url| {
+                if readonly {
+                    format!("{}.readonly", auth_base_url)
+                } else {
+                    auth_base_url.to_owned()
+                }
+            })
+            .collect();
+
+        let scopes_as_str: Vec<&str> = scopes.iter().map(String::as_str).collect();
+
+        let sa_auth = ServiceAccountAuthenticator::from_service_account_key(sa_key, &scopes_as_str).await?;
 
         let mut client = Client::from_authenticator(sa_auth);
         client.v2_base_url(self.v2_base_url.clone());
         Ok(client)
     }
 
     pub async fn build_from_service_account_key_file(&self, sa_key_file: &str) -> Result<Client, BQError> {
-        let scopes = vec![self.auth_base_url.as_str()];
+        let scopes = self.auth_base_as_str_urls();
         let sa_auth = service_account_authenticator(scopes, sa_key_file).await?;
 
         let mut client = Client::from_authenticator(sa_auth);
@@ -59,13 +68,20 @@ impl ClientBuilder {
     }
 
     pub async fn build_with_workload_identity(&self, readonly: bool) -> Result<Client, BQError> {
-        let scope = if readonly {
-            format!("{BIG_QUERY_AUTH_URL}.readonly")
-        } else {
-            BIG_QUERY_AUTH_URL.to_string()
-        };
+        let scopes: Vec<String> = vec![BIG_QUERY_AUTH_URL, DRIVE_AUTH_URL]
+            .iter()
+            .map(|auth_base_url| {
+                if readonly {
+                    format!("{auth_base_url}.readonly")
+                } else {
+                    auth_base_url.to_string()
+                }
+            })
+            .collect();
 
-        let sa_auth = ServiceAccountAuthenticator::with_workload_identity(&[&scope]).await?;
+        let scopes_as_str: Vec<&str> = scopes.iter().map(String::as_str).collect();
+
+        let sa_auth = ServiceAccountAuthenticator::with_workload_identity(&scopes_as_str).await?;
 
         let mut client = Client::from_authenticator(sa_auth);
         client.v2_base_url(self.v2_base_url.clone());
@@ -77,8 +93,7 @@ impl ClientBuilder {
         secret: S,
         persistant_file_path: P,
     ) -> Result<Client, BQError> {
-        let scopes = vec![self.auth_base_url.as_str()];
-        let auth = installed_flow_authenticator(secret, &scopes, persistant_file_path).await?;
+        let auth = installed_flow_authenticator(secret, &self.auth_base_as_str_urls(), persistant_file_path).await?;
 
         let mut client = Client::from_authenticator(auth);
         client.v2_base_url(self.v2_base_url.clone());
@@ -100,8 +115,7 @@ impl ClientBuilder {
     }
 
     pub async fn build_from_application_default_credentials(&self) -> Result<Client, BQError> {
-        let scopes = vec![self.auth_base_url.as_str()];
-        let auth = application_default_credentials_authenticator(&scopes).await?;
+        let auth = application_default_credentials_authenticator(&self.auth_base_as_str_urls()).await?;
 
         let mut client = Client::from_authenticator(auth);
         client.v2_base_url(self.v2_base_url.clone());
@@ -112,13 +126,16 @@ impl ClientBuilder {
         &self,
         authorized_user_secret_path: P,
     ) -> Result<Client, BQError> {
-        let scopes = vec![self.auth_base_url.as_str()];
-        let auth = authorized_user_authenticator(authorized_user_secret_path, &scopes).await?;
+        let auth = authorized_user_authenticator(authorized_user_secret_path, &self.auth_base_as_str_urls()).await?;
 
         let mut client = Client::from_authenticator(auth);
         client.v2_base_url(self.v2_base_url.clone());
         Ok(client)
     }
+
+    pub fn auth_base_as_str_urls(&self) -> Vec<&str> {
+        self.auth_base_urls.iter().map(String::as_str).collect()
+    }
 }
 
 impl Default for ClientBuilder {

src/error.rs

@@ -53,6 +53,9 @@ pub enum BQError {
 
     #[error("Json serialization error (error: {0})")]
     SerializationError(#[from] serde_json::Error),
+
+    #[error("Timeout Error")]
+    TimeoutError,
 }
 
 #[derive(Debug, Deserialize)]

src/job.rs

@@ -66,6 +66,13 @@ impl JobApi {
         let resp = self.client.execute(request).await?;
 
         let query_response: QueryResponse = process_response(resp).await?;
+
+        if let Some(job_complete) = query_response.job_complete {
+            if !job_complete {
+                return Err(BQError::TimeoutError);
+            }
+        }
+
         Ok(ResultSet::new(query_response))
     }
 

src/lib.rs

@@ -51,6 +51,7 @@ 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";
+const DRIVE_AUTH_URL: &str = "https://www.googleapis.com/auth/drive";
 
 /// An asynchronous BigQuery client.
 #[derive(Clone)]