collect all env var as part of connection_settings

pranshi06 · pranshi06 · commit b78d7297d53c · 2024-09-18T14:31:20.000+05:30
diff --git a/crates/configuration/src/connection_settings.rs b/crates/configuration/src/connection_settings.rs
@@ -0,0 +1,37 @@
+//! Database connection settings.
+
+use crate::values::{DatasetId, ProjectId, Secret, ServiceKey};
+use schemars::JsonSchema;
+use serde::{Deserialize, Serialize};
+
+pub const DEFAULT_SERVICE_KEY_VARIABLE: &str = "HASURA_BIGQUERY_SERVICE_KEY";
+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,
+    /// Project ID for a BigQuery database.
+    pub project_id: ProjectId,
+    /// Dataset ID for a BigQuery database.
+    pub dataset_id: DatasetId,
+}
+
+impl DatabaseConnectionSettings {
+    pub fn empty() -> Self {
+        Self {
+            service_key: ServiceKey(Secret::FromEnvironment {
+                variable: DEFAULT_SERVICE_KEY_VARIABLE.into(),
+            }),
+            project_id: ProjectId(Secret::FromEnvironment {
+                variable: DEFAULT_PROJECT_ID_VARIABLE.into(),
+            }),
+            dataset_id: DatasetId(Secret::FromEnvironment {
+                variable: DEFAULT_DATASET_ID_VARIABLE.into(),
+            }),
+        }
+    }
+}
diff --git a/crates/configuration/src/error.rs b/crates/configuration/src/error.rs
@@ -16,8 +16,14 @@ pub enum ParseConfigurationError {
         column: usize,
         message: String,
     },
-    #[error("empty connection URI")]
-    EmptyConnectionUri { file_path: std::path::PathBuf },
+    #[error("empty service account key")]
+    EmptyServiceKey { file_path: std::path::PathBuf },
+
+    #[error("empty project ID")]
+    EmptyProjectId { file_path: std::path::PathBuf },
+
+    #[error("empty dataset ID")]
+    EmptyDatasetId { file_path: std::path::PathBuf },
 
     #[error("I/O error: {0}")]
     IoErrorButStringified(String),
diff --git a/crates/configuration/src/lib.rs b/crates/configuration/src/lib.rs
@@ -1,12 +1,13 @@
 pub mod configuration;
+pub mod connection_settings;
 pub mod environment;
 pub mod error;
 pub mod to_runtime_configuration;
 pub mod values;
 pub mod version1;
 
 pub use configuration::Configuration;
-pub use values::uri::ConnectionUri;
+pub use values::connection_info::ServiceKey;
 pub use version1::{
     configure,
     parse_configuration,
diff --git a/crates/configuration/src/to_runtime_configuration.rs b/crates/configuration/src/to_runtime_configuration.rs
@@ -6,7 +6,7 @@ use std::collections::BTreeMap;
 use super::version1::ParsedConfiguration;
 use crate::environment::Environment;
 use crate::error::MakeRuntimeConfigurationError;
-use crate::values::{ConnectionUri, Secret};
+use crate::values::{Secret, ServiceKey};
 use query_engine_metadata::{self, metadata};
 // use crate::VersionTag;
 
@@ -16,9 +16,9 @@ pub fn make_runtime_configuration(
     parsed_config: ParsedConfiguration,
     environment: impl Environment,
 ) -> Result<crate::Configuration, MakeRuntimeConfigurationError> {
-    let connection_uri = match parsed_config.service_key {
-        ConnectionUri(Secret::Plain(uri)) => Ok(uri),
-        ConnectionUri(Secret::FromEnvironment { variable }) => {
+    let connection_uri = match parsed_config.connection_settings.service_key {
+        ServiceKey(Secret::Plain(uri)) => Ok(uri),
+        ServiceKey(Secret::FromEnvironment { variable }) => {
             environment.read(&variable).map_err(|error| {
                 MakeRuntimeConfigurationError::MissingEnvironmentVariable {
                     file_path: super::version1::CONFIGURATION_FILENAME.into(),
diff --git a/crates/configuration/src/values/connection_info.rs b/crates/configuration/src/values/connection_info.rs
@@ -3,6 +3,21 @@ use serde::{Deserialize, Serialize};
 
 use super::Secret;
 
+#[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize, JsonSchema)]
+pub struct ServiceKey(pub Secret);
+
+impl From<String> for ServiceKey {
+    fn from(value: String) -> Self {
+        Self(value.into())
+    }
+}
+
+impl From<&str> for ServiceKey {
+    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
@@ -1,9 +1,7 @@
-pub mod database_info;
+pub mod connection_info;
 mod pool_settings;
 mod secret;
-pub mod uri;
 
-pub use database_info::{DatasetId, ProjectId};
+pub use connection_info::{DatasetId, ProjectId, ServiceKey};
 pub use pool_settings::PoolSettings;
 pub use secret::Secret;
-pub use uri::ConnectionUri;
diff --git a/crates/configuration/src/values/uri.rs b/crates/configuration/src/values/uri.rs
diff --git a/crates/configuration/src/version1.rs b/crates/configuration/src/version1.rs
@@ -1,8 +1,9 @@
 //! Internal Configuration and state for our connector.
 
+use crate::connection_settings;
 use crate::environment::Environment;
 use crate::error::WriteParsedConfigurationError;
-use crate::values::{ConnectionUri, DatasetId, PoolSettings, ProjectId, Secret};
+use crate::values::{DatasetId, PoolSettings, ProjectId, Secret, ServiceKey};
 
 use super::error::ParseConfigurationError;
 use gcp_bigquery_client::model::query_request::QueryRequest;
@@ -51,10 +52,7 @@ const NOT_APPROX_COUNTABLE: [&str; 4] = ["image", "sql_variant", "ntext", "text"
 pub struct ParsedConfiguration {
     // Which version of the configuration format are we using
     pub version: u32,
-    // Connection string for a Postgres-compatible database
-    pub service_key: ConnectionUri,
-    pub project_id: ProjectId,
-    pub dataset_id: DatasetId,
+    pub connection_settings: connection_settings::DatabaseConnectionSettings,
     #[serde(skip_serializing_if = "PoolSettings::is_default")]
     #[serde(default)]
     pub pool_settings: PoolSettings,
@@ -80,15 +78,7 @@ impl ParsedConfiguration {
     pub fn empty() -> Self {
         Self {
             version: CURRENT_VERSION,
-            service_key: ConnectionUri(Secret::FromEnvironment {
-                variable: DEFAULT_SERVICE_KEY_VARIABLE.into(),
-            }),
-            project_id: ProjectId(Secret::FromEnvironment {
-                variable: DEFAULT_PROJECT_ID_VARIABLE.into(),
-            }),
-            dataset_id: DatasetId(Secret::FromEnvironment {
-                variable: DEFAULT_DATASET_ID_VARIABLE.into(),
-            }),
+            connection_settings: connection_settings::DatabaseConnectionSettings::empty(),
             pool_settings: PoolSettings::default(),
             metadata: metadata::Metadata::default(),
             // aggregate_functions: metadata::AggregateFunctions::default(),
@@ -101,19 +91,17 @@ pub async fn configure(
     args: &ParsedConfiguration,
     environment: impl Environment,
 ) -> anyhow::Result<ParsedConfiguration> {
-    let service_key = match &args.service_key {
-        ConnectionUri(Secret::Plain(value)) => Cow::Borrowed(value),
-        ConnectionUri(Secret::FromEnvironment { variable }) => {
-            Cow::Owned(environment.read(variable)?)
-        }
+    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 project_id_ = match &args.project_id {
+    let project_id_ = match &args.connection_settings.project_id {
         ProjectId(Secret::Plain(value)) => Cow::Borrowed(value),
         ProjectId(Secret::FromEnvironment { variable }) => Cow::Owned(environment.read(variable)?),
     };
 
-    let dataset_id_ = match &args.dataset_id {
+    let dataset_id_ = match &args.connection_settings.dataset_id {
         DatasetId(Secret::Plain(value)) => Cow::Borrowed(value),
         DatasetId(Secret::FromEnvironment { variable }) => Cow::Owned(environment.read(variable)?),
     };
@@ -220,9 +208,11 @@ pub async fn configure(
 
     Ok(ParsedConfiguration {
         version: 1,
-        service_key: args.service_key.clone(),
-        project_id: args.project_id.clone(),
-        dataset_id: args.dataset_id.clone(),
+        connection_settings: connection_settings::DatabaseConnectionSettings {
+            service_key: args.connection_settings.service_key.clone(),
+            project_id: args.connection_settings.project_id.clone(),
+            dataset_id: args.connection_settings.dataset_id.clone(),
+        },
         pool_settings: args.pool_settings.clone(),
         metadata: metadata::Metadata {
             tables: tables_info,
diff --git a/crates/connectors/ndc-bigquery/src/connector.rs b/crates/connectors/ndc-bigquery/src/connector.rs
@@ -178,12 +178,32 @@ impl<Env: Environment + Send + Sync> ConnectorSetup for BigQuerySetup<Env> {
                     message,
                 })
                 .into(),
-                configuration::error::ParseConfigurationError::EmptyConnectionUri { file_path } => {
+                configuration::error::ParseConfigurationError::EmptyServiceKey { file_path } => {
                     connector::ParseError::ValidateError(connector::InvalidNodes(vec![
                         connector::InvalidNode {
                             file_path,
-                            node_path: vec![connector::KeyOrIndex::Key("connectionUri".into())],
-                            message: "database connection URI must be specified".to_string(),
+                            node_path: vec![connector::KeyOrIndex::Key("serviceKey".into())],
+                            message: "Service account key must be specified".to_string(),
+                        },
+                    ]))
+                    .into()
+                }
+                configuration::error::ParseConfigurationError::EmptyProjectId { file_path } => {
+                    connector::ParseError::ValidateError(connector::InvalidNodes(vec![
+                        connector::InvalidNode {
+                            file_path,
+                            node_path: vec![connector::KeyOrIndex::Key("projectId".into())],
+                            message: "BigQuery project ID must be specified".to_string(),
+                        },
+                    ]))
+                    .into()
+                }
+                configuration::error::ParseConfigurationError::EmptyDatasetId { file_path } => {
+                    connector::ParseError::ValidateError(connector::InvalidNodes(vec![
+                        connector::InvalidNode {
+                            file_path,
+                            node_path: vec![connector::KeyOrIndex::Key("datasetId".into())],
+                            message: "BigQuery dataset ID must be specified".to_string(),
                         },
                     ]))
                     .into()
diff --git a/crates/connectors/ndc-bigquery/tests/common/mod.rs b/crates/connectors/ndc-bigquery/tests/common/mod.rs
@@ -4,10 +4,13 @@ use std::collections::HashMap;
 
 use ndc_postgres::connector;
 
-const POSTGRESQL_CONNECTION_STRING: &str = "postgresql://postgres:password@localhost:64002";
-
 pub const CHINOOK_DEPLOYMENT_PATH: &str = "static/chinook-deployment.json";
 
+pub const BIGQUERY_SERVICE_KEY: &str = "{\"type\": \"service_account\",\"project_id\": \"hasura-development\",\"private_key_id\": \"222dd3f9e98b6743bb8d74d7a126fe89e6ac221d\",\"private_key\": \"-----BEGIN PRIVATE KEY-----\nMIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQDZuxyxWk6bOxHr\nht+MPZ7Q+F4D7AzTqCTZOmcldod+KHMlUCKwIOLQabAO8TEPhvcYyzBQ4gCwoN2i\n7VoQbHmlTpQu1s43K25oIoEIicFTbHcL4MALiFnT44XYl+PxL+e//GibJYGjwqI+\n3o2Go//o8BOfQEO/PPiQdub8/4VQjXGE0+xLeUjYURiJPu8ojrL2FdIqvzMaDWKq\n2r6KuukIeWkqmt6fHnrGiseavg3g7pBPjqmRtX6ekY74XbkTQk1kmCKf9MLjZ1UI\n+8QNp1C4pO4eDbp1Zkqz3uHhzccUvStkSCmppjKfD64Tp+6ExbUnMfq1UJ0GJBDM\nVPeJF6+PAgMBAAECggEAFFSf88GKeH02CaeQ5S/1ze57HOOkOPlI443MBtgAA9w0\nEEZgztBrTWmo+mQ0IA6KsSJ78vl/df63Y1jFYaY3X6OsO4lsPQONriSWptzyE9+b\naB0G4azMMnhazaQ1MRa3jZo8jEwexFNOwg8W6P0UTsRoGKUwDkHbteWcYQBdCu3W\nFa/CX3Tw0n/DdAVNi8Ai9K0d+Okmcv+ZRopeNuLENR28/VGSXj+Li1V7A0s+nX9E\nyxuGrDY4WMxSXHkW2yjrDnPUs6dXLFk1HBQPaHrs3i6gGyNXfTNWUJ3nGQwZIqJI\na1b4TMiGVapq33qCo/3Yi6jQ+I6KnpmWgQ7y5LXhoQKBgQDuA80oWCXQv7MERg91\nFwammtXrMjoD234u3RGNtnU67yH87kvL+p18EiNlbmy+CWyoc1mOjLtTHvMBfMGh\nfKt3BSuzrZZArA1GJF6J2Rew5dkJGzwPogLSnXMgrVwknAejKJw97wTJzzIZuuSc\nb7P57+mFoSdR+eSb44WFcuMyoQKBgQDqLu9LWz+LcljDWDeMQ4kl8gkNZMe5//Qd\nOpa6mN6T2nfRgxasaLo7WO8TqT4X28eBJKuru4BOeHVx0Y8GvWYyaW0uEEycdXVl\n6man+YUhZezTjjB/nCeaz7E7LCcUao1JP2Y9xlnpO5jdyi2tYkCqu7vOxmnLArN/\nl3zuXgrkLwKBgEzCzReF1ixMpt9p+PI6StrQdM01laBI2ZkjktWxUn1/Qebgs3FF\nkiTBdMjxpABl6bUp/mgK2x8jjBuesJP0MRhhgoagJSUWV/GXKSYr7YgPmL9nGSex\niFeEj+yp/F2SNKRaJImU3GZ5fB7wN2p8W/7vcNC3+IZnoWLlLdqsAroBAoGAdzZh\nVoki9gfFq9uym1Kd9JUbipftHIBxcpeqt16un7GtIRiMaEP/2cpSGj4jf92/17wl\nMA0JKekkUEaPeqzb43nLvJFLjrI0iyciDwx0eyX5w1A03CFP//0OicLWOgxr1AfU\nMkpQ5uwRy4XqbsL/jGp5Fq/mlxPO8HrbfDSfcr0CgYEAxN/RMCYODz+p9xZ6tbiS\nfHFrCgvPpYR9hEWhb/DyT4Q/OSzk0TItuSXGc3uicYeIycHIndyWej/a1HGg0IRK\nqjGbqGvRJIrzhLvLog1oOGADFSE2IJrxV2m9lQG8IUow4QUFcoZaCXZAQEvWeo+D\nq+4Pe2w4aMZeyqpt/mOSGzQ=\n-----END PRIVATE KEY-----\n\",\"client_email\": \"skm-bq-test@hasura-development.iam.gserviceaccount.com\",\"client_id\": \"116460406056940511807\",\"auth_uri\": \"https://accounts.google.com/o/oauth2/auth\",\"token_uri\": \"https://oauth2.googleapis.com/token\",\"auth_provider_x509_cert_url\": \"https://www.googleapis.com/oauth2/v1/certs\",\"client_x509_cert_url\": \"https://www.googleapis.com/robot/v1/metadata/x509/skm-bq-test%40hasura-development.iam.gserviceaccount.com\",\"universe_domain\": \"googleapis.com\"}";
+
+pub const BIGQUERY_PROJECT_ID: &str = "hasura-development";
+pub const BIGQUERY_DATASET_ID: &str = "chinook_sample";
+
 /// Creates a router with a fresh state from the test deployment.
 pub async fn create_router() -> axum::Router {
     create_router_from_deployment(CHINOOK_DEPLOYMENT_PATH).await
@@ -17,10 +20,20 @@ pub async fn create_router() -> axum::Router {
 pub async fn create_router_from_deployment(deployment_path: &str) -> axum::Router {
     let _ = env_logger::builder().is_test(true).try_init();
 
-    let environment = HashMap::from([(
-        ndc_bigquery_configuration::version1::DEFAULT_SERVICE_KEY_VARIABLE.into(),
-        POSTGRESQL_CONNECTION_STRING.to_string(),
-    )]);
+    let environment = HashMap::from([
+        (
+            ndc_bigquery_configuration::version1::DEFAULT_SERVICE_KEY_VARIABLE.into(),
+            BIGQUERY_SERVICE_KEY.to_string(),
+        ),
+        (
+            ndc_bigquery_configuration::version1::DEFAULT_PROJECT_ID_VARIABLE.into(),
+            BIGQUERY_PROJECT_ID.to_string(),
+        ),
+        (
+            ndc_bigquery_configuration::version1::DEFAULT_DATASET_ID_VARIABLE.into(),
+            BIGQUERY_DATASET_ID.to_string(),
+        ),
+    ]);
 
     let setup = connector::BigQuerySetup::new(environment);
 
diff --git a/crates/connectors/ndc-bigquery/tests/configuration_tests.rs b/crates/connectors/ndc-bigquery/tests/configuration_tests.rs
@@ -11,7 +11,7 @@ use similar_asserts::assert_eq;
 use ndc_bigquery_configuration::{
     values::Secret,
     version1::{self, DEFAULT_SERVICE_KEY_VARIABLE},
-    ConnectionUri,
+    ServiceKey,
 };
 
 use tests_common::deployment::helpers::get_path_from_project_root;
@@ -40,7 +40,8 @@ async fn test_configure() {
         POSTGRESQL_CONNECTION_STRING.into(),
     )]);
 
-    args.service_key = ConnectionUri(Secret::Plain(DEFAULT_SERVICE_KEY_VARIABLE.to_string()));
+    args.connection_settings.service_key =
+        ServiceKey(Secret::Plain(DEFAULT_SERVICE_KEY_VARIABLE.to_string()));
 
     let actual = version1::configure(&args, environment)
         .await
diff --git a/crates/query-engine/sql/src/sql/helpers.rs b/crates/query-engine/sql/src/sql/helpers.rs
@@ -4,6 +4,7 @@ use super::ast::*;
 use std::collections::BTreeMap;
 
 /// Used as input to helpers to construct SELECTs which return 'rows' and/or 'aggregates' results.
+#[derive(Debug, Clone, PartialEq)]
 pub enum SelectSet {
     Rows(Select),
     Aggregates(Select),
diff --git a/crates/query-engine/translation/src/translation/query/mod.rs b/crates/query-engine/translation/src/translation/query/mod.rs
@@ -42,6 +42,8 @@ pub fn translate(
         &query_request.query,
     )?;
 
+    dbg!(&select_set);
+
     // form a single JSON item shaped `{ rows: [], aggregates: {} }`
     // that matches the models::RowSet type
     let json_select = sql::helpers::select_rowset(
diff --git a/crates/query-engine/translation/src/translation/query/values.rs b/crates/query-engine/translation/src/translation/query/values.rs
@@ -82,10 +82,17 @@ fn type_to_ast_scalar_type_name(
                 None => Ok(sql::ast::ScalarTypeName::Unqualified(
                     scalar_type.type_name.to_string(),
                 )),
-                Some(schema_name) => Ok(sql::ast::ScalarTypeName::Qualified {
-                    schema_name: sql::ast::SchemaName(schema_name),
-                    type_name: scalar_type.type_name.to_string(),
-                }),
+                Some(_schema_name) =>
+                // FIXME(PY): How to use Qualified types. In cast it gives something like Cast(@param1 as hasura-development.chinook_sample.string) when using _in operator for strings
+                // Ok(sql::ast::ScalarTypeName::Qualified {
+                //     schema_name: sql::ast::SchemaName(schema_name),
+                //     type_name: scalar_type.type_name.to_string(),
+                // }),
+                {
+                    Ok(sql::ast::ScalarTypeName::Unqualified(
+                        scalar_type.type_name.to_string(),
+                    ))
+                }
             }
         }
         query_engine_metadata::metadata::Type::CompositeType(t) => {
