wip: add workload identity auth

py · py · commit 7abb2353d2b9 · 2024-12-19T14:11:19.000+05:30
diff --git a/crates/cli/src/lib.rs b/crates/cli/src/lib.rs
@@ -108,6 +108,11 @@ async fn initialize(with_metadata: bool, context: Context<impl Environment>) ->
                     description: "The BigQuery service key".to_string(),
                     default_value: None,
                 },
+                metadata::EnvironmentVariableDefinition {
+                    name: "HASURA_BIGQUERY_WORKLOAD_IDENTITY_AUTH".to_string(),
+                    description: "The BigQuery workload identity auth URL".to_string(),
+                    default_value: None,
+                },
                 metadata::EnvironmentVariableDefinition {
                     name: "HASURA_BIGQUERY_PROJECT_ID".to_string(),
                     description: "The BigQuery project ID/name".to_string(),
diff --git a/crates/configuration/src/configuration.rs b/crates/configuration/src/configuration.rs
@@ -19,7 +19,7 @@ use crate::values::PoolSettings;
 pub struct Configuration {
     pub metadata: metadata::Metadata,
     pub pool_settings: PoolSettings,
-    pub service_key: String,
+    pub auth: (bool, String),
     pub project_id: String,
     pub dataset_id: String,
     // pub mutations_version: Option<metadata::mutations::MutationsVersion>,
diff --git a/crates/configuration/src/connection_settings.rs b/crates/configuration/src/connection_settings.rs
@@ -1,19 +1,22 @@
 //! Database connection settings.
 
-use crate::values::{DatasetId, ProjectId, Secret, ServiceKey};
+use crate::values::{DatasetId, ProjectId, Secret, ServiceKey, WorkloadIdentityAuth};
 use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
 
 pub const DEFAULT_SERVICE_KEY_VARIABLE: &str = "HASURA_BIGQUERY_SERVICE_KEY";
+pub const DEFAULT_WORKLOAD_IDENTITY_AUTH_VARIABLE: &str = "HASURA_BIGQUERY_WORKLOAD_IDENTITY_AUTH";
 pub const DEFAULT_PROJECT_ID_VARIABLE: &str = "HASURA_BIGQUERY_PROJECT_ID";
 pub const DEFAULT_DATASET_ID_VARIABLE: &str = "HASURA_BIGQUERY_DATASET_ID";
 
 /// Database connection settings.
 #[derive(Clone, PartialEq, Eq, Debug, Deserialize, Serialize, JsonSchema)]
 #[serde(rename_all = "camelCase")]
 pub struct DatabaseConnectionSettings {
-    /// Connection string for a Postgres-compatible database.
-    pub service_key: ServiceKey,
+    /// Connection string for a BigQuery-compatible database.
+    pub service_key: Option<ServiceKey>,
+    /// Connection string for a BigQuery-compatible database.
+    pub workload_identity_auth: Option<WorkloadIdentityAuth>,
     /// Project ID for a BigQuery database.
     pub project_id: ProjectId,
     /// Dataset ID for a BigQuery database.
@@ -23,9 +26,12 @@ pub struct DatabaseConnectionSettings {
 impl DatabaseConnectionSettings {
     pub fn empty() -> Self {
         Self {
-            service_key: ServiceKey(Secret::FromEnvironment {
+            service_key: Some(ServiceKey(Secret::FromEnvironment {
                 variable: DEFAULT_SERVICE_KEY_VARIABLE.into(),
-            }),
+            })),
+            workload_identity_auth: Some(WorkloadIdentityAuth(Secret::FromEnvironment {
+                variable: DEFAULT_WORKLOAD_IDENTITY_AUTH_VARIABLE.into(),
+            })),
             project_id: ProjectId(Secret::FromEnvironment {
                 variable: DEFAULT_PROJECT_ID_VARIABLE.into(),
             }),
diff --git a/crates/configuration/src/to_runtime_configuration.rs b/crates/configuration/src/to_runtime_configuration.rs
@@ -6,25 +6,49 @@ use std::collections::BTreeMap;
 use super::version1::ParsedConfiguration;
 use crate::environment::Environment;
 use crate::error::MakeRuntimeConfigurationError;
-use crate::values::{DatasetId, ProjectId, Secret, ServiceKey};
+use crate::values::{DatasetId, ProjectId, Secret, ServiceKey, WorkloadIdentityAuth};
 use query_engine_metadata::{self, metadata};
-// use crate::VersionTag;
 
 /// Convert the parsed configuration metadata to internal engine metadata
 /// That can be used by the connector at runtime.
 pub fn make_runtime_configuration(
     parsed_config: ParsedConfiguration,
     environment: impl Environment,
 ) -> Result<crate::Configuration, MakeRuntimeConfigurationError> {
-    let service_key = match parsed_config.connection_settings.service_key {
-        ServiceKey(Secret::Plain(key)) => Ok(key),
-        ServiceKey(Secret::FromEnvironment { variable }) => {
-            environment.read(&variable).map_err(|error| {
+    let auth: (bool, String) = match parsed_config.connection_settings.workload_identity_auth {
+        Some(WorkloadIdentityAuth(Secret::Plain(key))) => {
+            Ok((true, key))
+        },
+        Some(WorkloadIdentityAuth(Secret::FromEnvironment { variable })) => {
+            let variable_value = environment.read(&variable).map_err(|error| {
                 MakeRuntimeConfigurationError::MissingEnvironmentVariable {
                     file_path: super::version1::CONFIGURATION_FILENAME.into(),
                     message: error.to_string(),
                 }
-            })
+            })?;
+            Ok((true, variable_value))
+        },
+        None => {
+            match parsed_config.connection_settings.service_key {
+                Some(ServiceKey(Secret::Plain(key))) => {
+                    Ok((false, key))
+                },
+                Some(ServiceKey(Secret::FromEnvironment { variable })) => {
+                    let variable_value = environment.read(&variable).map_err(|error| {
+                        MakeRuntimeConfigurationError::MissingEnvironmentVariable {
+                            file_path: super::version1::CONFIGURATION_FILENAME.into(),
+                            message: error.to_string(),
+                        }
+                    })?;
+                    Ok((false, variable_value))
+                },
+                None => {
+                    return Err(MakeRuntimeConfigurationError::MissingEnvironmentVariable {
+                        file_path: super::version1::CONFIGURATION_FILENAME.into(),
+                        message: "Neither Workload Identity Auth URL or Service key is provided".into(),
+                    });
+                },
+            }
         }
     }?;
 
@@ -53,10 +77,9 @@ pub fn make_runtime_configuration(
     Ok(crate::Configuration {
         metadata: convert_metadata(parsed_config.metadata),
         pool_settings: parsed_config.pool_settings,
-        service_key,
+        auth,
         project_id,
         dataset_id,
-        // mutations_version: convert_mutations_version(parsed_config.mutations_version),
     })
 }
 
diff --git a/crates/configuration/src/values/connection_info.rs b/crates/configuration/src/values/connection_info.rs
@@ -18,6 +18,21 @@ impl From<&str> for ServiceKey {
     }
 }
 
+#[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize, JsonSchema)]
+pub struct WorkloadIdentityAuth(pub Secret);
+
+impl From<String> for WorkloadIdentityAuth {
+    fn from(value: String) -> Self {
+        Self(value.into())
+    }
+}
+
+impl From<&str> for WorkloadIdentityAuth {
+    fn from(value: &str) -> Self {
+        Self::from(value.to_string())
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize, JsonSchema)]
 pub struct ProjectId(pub Secret);
 
diff --git a/crates/configuration/src/values/mod.rs b/crates/configuration/src/values/mod.rs
@@ -2,6 +2,6 @@ pub mod connection_info;
 mod pool_settings;
 mod secret;
 
-pub use connection_info::{DatasetId, ProjectId, ServiceKey};
+pub use connection_info::{DatasetId, ProjectId, ServiceKey, WorkloadIdentityAuth};
 pub use pool_settings::PoolSettings;
 pub use secret::Secret;
diff --git a/crates/configuration/src/version1.rs b/crates/configuration/src/version1.rs
@@ -3,7 +3,7 @@
 use crate::connection_settings;
 use crate::environment::Environment;
 use crate::error::WriteParsedConfigurationError;
-use crate::values::{DatasetId, PoolSettings, ProjectId, Secret, ServiceKey};
+use crate::values::{DatasetId, PoolSettings, ProjectId, Secret, ServiceKey, WorkloadIdentityAuth};
 
 use super::error::ParseConfigurationError;
 use gcp_bigquery_client::model::query_request::QueryRequest;
@@ -94,10 +94,44 @@ pub async fn configure(
     args: &ParsedConfiguration,
     environment: impl Environment,
 ) -> anyhow::Result<ParsedConfiguration> {
-    let service_key = match &args.connection_settings.service_key {
-        ServiceKey(Secret::Plain(value)) => Cow::Borrowed(value),
-        ServiceKey(Secret::FromEnvironment { variable }) => Cow::Owned(environment.read(variable)?),
+    let bigquery_client = match &args.connection_settings.workload_identity_auth {
+        Some(WorkloadIdentityAuth(Secret::Plain(value))) => {
+            let url = Cow::Borrowed(value);
+            std::env::set_var("BIG_QUERY_AUTH_URL", url.as_ref());
+            gcp_bigquery_client::Client::with_workload_identity(false).await.unwrap()
+        },
+        Some(WorkloadIdentityAuth(Secret::FromEnvironment { variable })) => {
+            let url: Cow<'_, String> = Cow::Owned(environment.read(variable)?);
+            std::env::set_var("BIG_QUERY_AUTH_URL", url.as_ref());
+            gcp_bigquery_client::Client::with_workload_identity(false).await.unwrap()
+        },
+        None => {
+            match &args.connection_settings.service_key {
+                Some(ServiceKey(Secret::Plain(value))) => {
+                    let service_key = Cow::Borrowed(value);
+                    let service_account_key = yup_oauth2::parse_service_account_key(service_key.as_str()).unwrap();
+                    gcp_bigquery_client::Client::from_service_account_key(service_account_key, false).await.unwrap()
+                },
+                Some(ServiceKey(Secret::FromEnvironment { variable })) => {
+                    let service_key: Cow<'_, String> = Cow::Owned(environment.read(variable)?);
+                    let service_account_key = yup_oauth2::parse_service_account_key(service_key.as_str()).unwrap();
+                    gcp_bigquery_client::Client::from_service_account_key(service_account_key, false).await.unwrap()
+                },
+                None => {
+                    return Err(anyhow::anyhow!("Neither Workload Identity Auth URL or Service key is provided"));
+                },
+            }
+        }
     };
+    // let service_key = match &args.connection_settings.service_key {
+    //     Some(ServiceKey(Secret::Plain(value))) => Cow::Borrowed(value),
+    //     Some(ServiceKey(Secret::FromEnvironment { variable })) => Cow::Owned(environment.read(variable)?),
+    // };
+
+    // let workload_identity_auth = match &args.connection_settings.workload_identity_auth {
+    //     WorkloadIdentityAuth(Secret::Plain(value)) => Cow::Borrowed(value),
+    //     WorkloadIdentityAuth(Secret::FromEnvironment { variable }) => Cow::Owned(environment.read(variable)?),
+    // };
 
     let project_id_ = match &args.connection_settings.project_id {
         ProjectId(Secret::Plain(value)) => Cow::Borrowed(value),
@@ -109,19 +143,37 @@ pub async fn configure(
         DatasetId(Secret::FromEnvironment { variable }) => Cow::Owned(environment.read(variable)?),
     };
 
-    let service_account_key = yup_oauth2::parse_service_account_key(service_key.as_str()).unwrap();
+    // let service_account_key = yup_oauth2::parse_service_account_key(service_key.as_str()).unwrap();
 
     let project_id = project_id_.as_str();
     let dataset_id = dataset_id_.as_str();
 
     let schema_name = format!("{project_id}.{dataset_id}");
     let database_name = schema_name.clone();
 
-    // Init BigQuery client
-    let bigquery_client =
-        gcp_bigquery_client::Client::from_service_account_key(service_account_key, false)
-            .await
-            .unwrap();
+    // let bigquery_client = match
+    // BIG_QUERY_AUTH_URL
+
+    // let big_query_auth_url = match std::env::var("HASURA_BIGQUERY_WORKLOAD_IDENTITY_AUTH") {
+    //     Ok(val) => val,
+    //     Err(_) => {
+    //         return Err(anyhow::anyhow!(
+    //             "Environment variable HASURA_BIGQUERY_WORKLOAD_IDENTITY_AUTH not set"
+    //         ))
+    //     }
+    // };
+
+    // std::env::set_var("BIG_QUERY_AUTH_URL", workload_identity_auth.as_ref());
+
+    // if 
+
+    // let bigquery_client_auth = gcp_bigquery_client::Client::with_workload_identity(true).await.unwrap();
+
+    // // Init BigQuery client
+    // let bigquery_client =
+    //     gcp_bigquery_client::Client::from_service_account_key(service_account_key, false)
+    //         .await
+    //         .unwrap();
 
     // get scalar_types
 
@@ -221,6 +273,7 @@ pub async fn configure(
         version: 1,
         connection_settings: connection_settings::DatabaseConnectionSettings {
             service_key: args.connection_settings.service_key.clone(),
+            workload_identity_auth: args.connection_settings.workload_identity_auth.clone(),
             project_id: args.connection_settings.project_id.clone(),
             dataset_id: args.connection_settings.dataset_id.clone(),
         },
diff --git a/crates/connectors/ndc-bigquery/src/state.rs b/crates/connectors/ndc-bigquery/src/state.rs
@@ -5,10 +5,6 @@
 use thiserror::Error;
 use tracing::{info_span, Instrument};
 
-// use ndc_bigquery_configuration::PoolSettings;
-// use ndc_bigquery_configuration::ConfigurationError;
-// use super::configuration::{Configuration, ConfigurationError};
-// use query_engine_execution::database_info::{self, DatabaseInfo, DatabaseVersion};
 use query_engine_execution::metrics;
 
 /// State for our connector.
@@ -33,14 +29,19 @@ pub async fn create_state(
     .instrument(info_span!("Setup metrics"))
     .await?;
 
-    let service_account_key =
-        yup_oauth2::parse_service_account_key(configuration.service_key.clone()).unwrap();
-
-    // Init BigQuery client
-    let bigquery_client =
-        gcp_bigquery_client::Client::from_service_account_key(service_account_key, false)
-            .await
-            .unwrap();
+    let bigquery_client = match &configuration.auth {
+        (true, workload_identity_auth) => {
+            std::env::set_var("BIG_QUERY_AUTH_URL", workload_identity_auth);
+            gcp_bigquery_client::Client::with_workload_identity(false).await.unwrap()
+        }
+        (false, service_key) => {
+            let service_account_key =
+                yup_oauth2::parse_service_account_key(service_key.clone()).unwrap();
+            gcp_bigquery_client::Client::from_service_account_key(service_account_key, false)
+                .await
+                .unwrap()
+        }
+    };
 
     Ok(State {
         metrics,
