From 88f2cb91ee0ec2246d665c8a9eb182390f506e9b Mon Sep 17 00:00:00 2001 From: Will Eaton Date: Thu, 6 Jun 2024 13:51:06 -0400 Subject: [PATCH 1/6] first try --- snowflake-api/src/lib.rs | 53 ++++++++-------- snowflake-api/src/requests.rs | 10 +++ snowflake-api/src/session.rs | 112 +++++++++++++++------------------- 3 files changed, 85 insertions(+), 90 deletions(-) diff --git a/snowflake-api/src/lib.rs b/snowflake-api/src/lib.rs index 1fa7b36..2bd9677 100644 --- a/snowflake-api/src/lib.rs +++ b/snowflake-api/src/lib.rs @@ -198,6 +198,8 @@ impl AuthArgs { Ok(AuthType::Password(PasswordArgs { password })) } else if let Ok(private_key_pem) = std::env::var("SNOWFLAKE_PRIVATE_KEY") { Ok(AuthType::Certificate(CertificateArgs { private_key_pem })) + } else if let Ok(token) = std::env::var("SNOWFLAKE_OAUTH_TOKEN") { + Ok(AuthType::OAuth(OAuthArgs { token })) } else { Err(MissingEnvArgument( "SNOWFLAKE_PASSWORD or SNOWFLAKE_PRIVATE_KEY".to_owned(), @@ -221,6 +223,7 @@ impl AuthArgs { pub enum AuthType { Password(PasswordArgs), Certificate(CertificateArgs), + OAuth(OAuthArgs), } pub struct PasswordArgs { @@ -231,6 +234,10 @@ pub struct CertificateArgs { pub private_key_pem: String, } +pub struct OAuthArgs { + pub token: String, +} + #[must_use] pub struct SnowflakeApiBuilder { pub auth: AuthArgs, @@ -253,27 +260,20 @@ impl SnowflakeApiBuilder { None => Arc::new(Connection::new()?), }; + let session = Session::auth_builder( + Arc::clone(&connection), + &self.auth.account_identifier, + self.auth.warehouse.as_deref(), + self.auth.database.as_deref(), + self.auth.schema.as_deref(), + &self.auth.username, + self.auth.role.as_deref(), + ); + let session = match self.auth.auth_type { - AuthType::Password(args) => Session::password_auth( - Arc::clone(&connection), - &self.auth.account_identifier, - self.auth.warehouse.as_deref(), - self.auth.database.as_deref(), - self.auth.schema.as_deref(), - &self.auth.username, - self.auth.role.as_deref(), - &args.password, - ), - AuthType::Certificate(args) => Session::cert_auth( - Arc::clone(&connection), - &self.auth.account_identifier, - self.auth.warehouse.as_deref(), - self.auth.database.as_deref(), - self.auth.schema.as_deref(), - &self.auth.username, - self.auth.role.as_deref(), - &args.private_key_pem, - ), + AuthType::Password(args) => session.password(&args.password), + AuthType::Certificate(args) => session.cert(&args.private_key_pem), + AuthType::OAuth(args) => session.oauth(&args.token), }; let account_identifier = self.auth.account_identifier.to_uppercase(); @@ -313,8 +313,7 @@ impl SnowflakeApi { password: &str, ) -> Result { let connection = Arc::new(Connection::new()?); - - let session = Session::password_auth( + let session = Session::auth_builder( Arc::clone(&connection), account_identifier, warehouse, @@ -322,8 +321,8 @@ impl SnowflakeApi { schema, username, role, - password, - ); + ) + .password(password); let account_identifier = account_identifier.to_uppercase(); Ok(Self::new( @@ -345,7 +344,7 @@ impl SnowflakeApi { ) -> Result { let connection = Arc::new(Connection::new()?); - let session = Session::cert_auth( + let session = Session::auth_builder( Arc::clone(&connection), account_identifier, warehouse, @@ -353,8 +352,8 @@ impl SnowflakeApi { schema, username, role, - private_key_pem, - ); + ) + .cert(private_key_pem); let account_identifier = account_identifier.to_uppercase(); Ok(Self::new( diff --git a/snowflake-api/src/requests.rs b/snowflake-api/src/requests.rs index 77b0434..ace75d3 100644 --- a/snowflake-api/src/requests.rs +++ b/snowflake-api/src/requests.rs @@ -15,6 +15,7 @@ pub struct LoginRequest { } pub type PasswordLoginRequest = LoginRequest; +pub type OAuthLoginRequest = LoginRequest; #[cfg(feature = "cert-auth")] pub type CertLoginRequest = LoginRequest; @@ -62,6 +63,15 @@ pub struct CertRequestData { pub token: String, } +#[derive(Serialize, Debug)] +#[serde(rename_all = "SCREAMING_SNAKE_CASE")] +pub struct OAuthRequestData { + #[serde(flatten)] + pub login_request_common: LoginRequestCommon, + pub authenticator: String, + pub token: String, +} + #[derive(Serialize, Debug)] #[serde(rename_all = "camelCase")] pub struct RenewSessionRequest { diff --git a/snowflake-api/src/session.rs b/snowflake-api/src/session.rs index 90acaaf..977745e 100644 --- a/snowflake-api/src/session.rs +++ b/snowflake-api/src/session.rs @@ -11,8 +11,8 @@ use crate::connection::{Connection, QueryType}; #[cfg(feature = "cert-auth")] use crate::requests::{CertLoginRequest, CertRequestData}; use crate::requests::{ - ClientEnvironment, LoginRequest, LoginRequestCommon, PasswordLoginRequest, PasswordRequestData, - RenewSessionRequest, SessionParameters, + ClientEnvironment, LoginRequest, LoginRequestCommon, OAuthLoginRequest, OAuthRequestData, + PasswordLoginRequest, PasswordRequestData, RenewSessionRequest, SessionParameters, }; use crate::responses::AuthResponse; @@ -50,6 +50,9 @@ pub enum AuthError { #[error("Enable the cert-auth feature to use certificate authentication")] CertAuthNotEnabled, + + #[error("The authentication type has not been set yet on the connection!")] + AuthTypeUnset, } #[derive(Debug)] @@ -103,8 +106,11 @@ impl AuthToken { } enum AuthType { - Certificate, - Password, + Certificate(String), + Password(String), + OAuth(String), + // Used for initialization only + None, } /// Requests, caches, and renews authentication tokens. @@ -125,18 +131,11 @@ pub struct Session { username: String, role: Option, - // This is not used with the certificate auth crate - #[allow(dead_code)] - private_key_pem: Option, - password: Option, } // todo: make builder impl Session { - /// Authenticate using private certificate and JWT - // fixme: add builder or introduce structs - #[allow(clippy::too_many_arguments)] - pub fn cert_auth( + pub fn auth_builder( connection: Arc, account_identifier: &str, warehouse: Option<&str>, @@ -144,7 +143,6 @@ impl Session { schema: Option<&str>, username: &str, role: Option<&str>, - private_key_pem: &str, ) -> Self { // uppercase everything as this is the convention let account_identifier = account_identifier.to_uppercase(); @@ -154,58 +152,39 @@ impl Session { let username = username.to_uppercase(); let role = role.map(str::to_uppercase); - let private_key_pem = Some(private_key_pem.to_string()); Self { connection, auth_tokens: Mutex::new(None), - auth_type: AuthType::Certificate, - private_key_pem, + auth_type: AuthType::None, account_identifier, warehouse: warehouse.map(str::to_uppercase), database, username, role, schema, - password: None, } } - /// Authenticate using password + /// Authenticate using private certificate and JWT // fixme: add builder or introduce structs #[allow(clippy::too_many_arguments)] - pub fn password_auth( - connection: Arc, - account_identifier: &str, - warehouse: Option<&str>, - database: Option<&str>, - schema: Option<&str>, - username: &str, - role: Option<&str>, - password: &str, - ) -> Self { - let account_identifier = account_identifier.to_uppercase(); - - let database = database.map(str::to_uppercase); - let schema = schema.map(str::to_uppercase); + pub fn cert(mut self, private_key_pem: &str) -> Self { + self.auth_type = AuthType::Certificate(private_key_pem.to_string()); + self + } - let username = username.to_uppercase(); - let password = Some(password.to_string()); - let role = role.map(str::to_uppercase); + pub fn password(mut self, password: &str) -> Self { + self.auth_type = AuthType::Password(password.to_string()); + self + } - Self { - connection, - auth_tokens: Mutex::new(None), - auth_type: AuthType::Password, - account_identifier, - warehouse: warehouse.map(str::to_uppercase), - database, - username, - role, - password, - schema, - private_key_pem: None, - } + /// Authenticate using oauth + // fixme: add builder or introduce structs + #[allow(clippy::too_many_arguments)] + pub fn oauth(mut self, oauth_access_token: &str) -> Self { + self.auth_type = AuthType::OAuth(oauth_access_token.to_string()); + self } /// Get cached token or request a new one if old one has expired. @@ -217,19 +196,21 @@ impl Session { .is_some_and(|at| at.master_token.is_expired()) { // Create new session if tokens are absent or can not be exchange - let tokens = match self.auth_type { - AuthType::Certificate => { + let tokens = match &self.auth_type { + AuthType::Certificate(pem) => { log::info!("Starting session with certificate authentication"); if cfg!(feature = "cert-auth") { - self.create(self.cert_request_body()?).await + self.create(self.cert_request_body(pem)?).await } else { Err(AuthError::MissingCertificate)? } } - AuthType::Password => { + AuthType::Password(pwd) => { log::info!("Starting session with password authentication"); - self.create(self.passwd_request_body()?).await + self.create(self.passwd_request_body(pwd)).await } + AuthType::OAuth(token) => self.create(self.oauth_request_body(token)).await, + AuthType::None => Err(AuthError::AuthTypeUnset)?, }?; *auth_tokens = Some(tokens); } else if auth_tokens @@ -277,12 +258,9 @@ impl Session { } #[cfg(feature = "cert-auth")] - fn cert_request_body(&self) -> Result { + fn cert_request_body(&self, private_key_pem: &str) -> Result { let full_identifier = format!("{}.{}", &self.account_identifier, &self.username); - let private_key_pem = self - .private_key_pem - .as_ref() - .ok_or(AuthError::MissingCertificate)?; + let jwt_token = generate_jwt_token(private_key_pem, &full_identifier)?; Ok(CertLoginRequest { @@ -294,15 +272,23 @@ impl Session { }) } - fn passwd_request_body(&self) -> Result { - let password = self.password.as_ref().ok_or(AuthError::MissingPassword)?; - - Ok(PasswordLoginRequest { + fn passwd_request_body(&self, password: &str) -> PasswordLoginRequest { + PasswordLoginRequest { data: PasswordRequestData { login_request_common: self.login_request_common(), password: password.to_string(), }, - }) + } + } + + fn oauth_request_body(&self, oauth_access_token: &str) -> OAuthLoginRequest { + OAuthLoginRequest { + data: OAuthRequestData { + login_request_common: self.login_request_common(), + authenticator: "OAUTH".to_string(), + token: oauth_access_token.to_string(), + }, + } } /// Start new session, all the Snowflake temporary objects will be scoped towards it, From a6b6218cf00df1235a93bef21739d96d21cbafd5 Mon Sep 17 00:00:00 2001 From: Will Eaton Date: Thu, 6 Jun 2024 14:24:54 -0400 Subject: [PATCH 2/6] working checkpoint --- snowflake-api/src/session.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/snowflake-api/src/session.rs b/snowflake-api/src/session.rs index 977745e..4849b4a 100644 --- a/snowflake-api/src/session.rs +++ b/snowflake-api/src/session.rs @@ -209,7 +209,10 @@ impl Session { log::info!("Starting session with password authentication"); self.create(self.passwd_request_body(pwd)).await } - AuthType::OAuth(token) => self.create(self.oauth_request_body(token)).await, + AuthType::OAuth(token) => { + log::info!("Starting session with oauth authentication"); + self.create(self.oauth_request_body(token)).await + } AuthType::None => Err(AuthError::AuthTypeUnset)?, }?; *auth_tokens = Some(tokens); From 68f6fb772b6df3d6b509ae3284066c6aaf6f7282 Mon Sep 17 00:00:00 2001 From: Will Eaton Date: Thu, 6 Jun 2024 16:27:19 -0400 Subject: [PATCH 3/6] more radical changes --- snowflake-api/src/lib.rs | 86 +++----------------- snowflake-api/src/session.rs | 149 ++++++++++++++++++++++++----------- 2 files changed, 114 insertions(+), 121 deletions(-) diff --git a/snowflake-api/src/lib.rs b/snowflake-api/src/lib.rs index 2bd9677..39f344c 100644 --- a/snowflake-api/src/lib.rs +++ b/snowflake-api/src/lib.rs @@ -28,7 +28,7 @@ use reqwest_middleware::ClientWithMiddleware; use thiserror::Error; use responses::ExecResponse; -use session::{AuthError, Session}; +use session::{AuthError, Session, SessionBuilder}; use crate::connection::QueryType; use crate::connection::{Connection, ConnectionError}; @@ -260,20 +260,20 @@ impl SnowflakeApiBuilder { None => Arc::new(Connection::new()?), }; - let session = Session::auth_builder( - Arc::clone(&connection), - &self.auth.account_identifier, - self.auth.warehouse.as_deref(), - self.auth.database.as_deref(), - self.auth.schema.as_deref(), - &self.auth.username, - self.auth.role.as_deref(), - ); + let session = SessionBuilder::new(&self.auth.account_identifier, &self.auth.username) + .warehouse(self.auth.warehouse.as_deref()) + .database(self.auth.database.as_deref()) + .schema(self.auth.schema.as_deref()) + .role(self.auth.role.as_deref()); let session = match self.auth.auth_type { - AuthType::Password(args) => session.password(&args.password), - AuthType::Certificate(args) => session.cert(&args.private_key_pem), - AuthType::OAuth(args) => session.oauth(&args.token), + AuthType::Password(args) => { + session.build_password(Arc::clone(&connection), &args.password) + } + AuthType::Certificate(args) => { + session.build_cert(Arc::clone(&connection), &args.private_key_pem) + } + AuthType::OAuth(args) => session.build_oauth(Arc::clone(&connection), &args.token), }; let account_identifier = self.auth.account_identifier.to_uppercase(); @@ -302,66 +302,6 @@ impl SnowflakeApi { account_identifier, } } - /// Initialize object with password auth. Authentication happens on the first request. - pub fn with_password_auth( - account_identifier: &str, - warehouse: Option<&str>, - database: Option<&str>, - schema: Option<&str>, - username: &str, - role: Option<&str>, - password: &str, - ) -> Result { - let connection = Arc::new(Connection::new()?); - let session = Session::auth_builder( - Arc::clone(&connection), - account_identifier, - warehouse, - database, - schema, - username, - role, - ) - .password(password); - - let account_identifier = account_identifier.to_uppercase(); - Ok(Self::new( - Arc::clone(&connection), - session, - account_identifier, - )) - } - - /// Initialize object with private certificate auth. Authentication happens on the first request. - pub fn with_certificate_auth( - account_identifier: &str, - warehouse: Option<&str>, - database: Option<&str>, - schema: Option<&str>, - username: &str, - role: Option<&str>, - private_key_pem: &str, - ) -> Result { - let connection = Arc::new(Connection::new()?); - - let session = Session::auth_builder( - Arc::clone(&connection), - account_identifier, - warehouse, - database, - schema, - username, - role, - ) - .cert(private_key_pem); - - let account_identifier = account_identifier.to_uppercase(); - Ok(Self::new( - Arc::clone(&connection), - session, - account_identifier, - )) - } pub fn from_env() -> Result { SnowflakeApiBuilder::new(AuthArgs::from_env()?).build() diff --git a/snowflake-api/src/session.rs b/snowflake-api/src/session.rs index 4849b4a..bedab67 100644 --- a/snowflake-api/src/session.rs +++ b/snowflake-api/src/session.rs @@ -109,14 +109,12 @@ enum AuthType { Certificate(String), Password(String), OAuth(String), - // Used for initialization only - None, } /// Requests, caches, and renews authentication tokens. /// Tokens are given as response to creating new session in Snowflake. Session persists /// the configuration state and temporary objects (tables, procedures, etc). -// todo: split warehouse-database-schema and username-role-key into its own structs + // todo: close session after object is dropped pub struct Session { connection: Arc, @@ -125,68 +123,124 @@ pub struct Session { auth_type: AuthType, account_identifier: String, - warehouse: Option, - database: Option, - schema: Option, - username: String, role: Option, -} -// todo: make builder -impl Session { - pub fn auth_builder( - connection: Arc, - account_identifier: &str, - warehouse: Option<&str>, - database: Option<&str>, - schema: Option<&str>, - username: &str, - role: Option<&str>, - ) -> Self { - // uppercase everything as this is the convention - let account_identifier = account_identifier.to_uppercase(); + object_details: SessionObjectDetails, +} - let database = database.map(str::to_uppercase); - let schema = schema.map(str::to_uppercase); +#[must_use] +pub struct SessionBuilder { + account_identifier: String, + object_details: SessionObjectDetails, + username: String, + role: Option, +} +impl SessionBuilder { + pub fn new(account_identifier: &str, username: &str) -> Self { let username = username.to_uppercase(); - let role = role.map(str::to_uppercase); Self { - connection, - auth_tokens: Mutex::new(None), - auth_type: AuthType::None, - account_identifier, - warehouse: warehouse.map(str::to_uppercase), - database, + account_identifier: account_identifier.to_string(), + object_details: SessionObjectDetails::default(), username, - role, - schema, + role: None, } } - /// Authenticate using private certificate and JWT - // fixme: add builder or introduce structs - #[allow(clippy::too_many_arguments)] - pub fn cert(mut self, private_key_pem: &str) -> Self { - self.auth_type = AuthType::Certificate(private_key_pem.to_string()); + pub fn warehouse(mut self, warehouse: Option<&str>) -> Self { + self.object_details.warehouse = warehouse.map(str::to_string); + self + } + + pub fn database(mut self, database: Option<&str>) -> Self { + self.object_details.database = database.map(str::to_string); self } - pub fn password(mut self, password: &str) -> Self { - self.auth_type = AuthType::Password(password.to_string()); + pub fn schema(mut self, schema: Option<&str>) -> Self { + self.object_details.schema = schema.map(str::to_string); self } - /// Authenticate using oauth - // fixme: add builder or introduce structs - #[allow(clippy::too_many_arguments)] - pub fn oauth(mut self, oauth_access_token: &str) -> Self { - self.auth_type = AuthType::OAuth(oauth_access_token.to_string()); + pub fn role(mut self, role: Option<&str>) -> Self { + self.role = role.map(str::to_string); self } + pub fn build_oauth(&self, connection: Arc, oauth_access_token: &str) -> Session { + Session::new( + connection, + &self.account_identifier, + AuthType::OAuth(oauth_access_token.to_string()), + self.object_details.clone(), + &self.username, + self.role.as_deref(), + ) + } + + pub fn build_password(&self, connection: Arc, password: &str) -> Session { + Session::new( + connection, + &self.account_identifier, + AuthType::Password(password.to_string()), + self.object_details.clone(), + &self.username, + self.role.as_deref(), + ) + } + + pub fn build_cert(&self, connection: Arc, private_key_pem: &str) -> Session { + Session::new( + connection, + &self.account_identifier, + AuthType::Certificate(private_key_pem.to_string()), + self.object_details.clone(), + &self.username, + self.role.as_deref(), + ) + } +} + +#[derive(Debug, Clone)] +struct SessionObjectDetails { + warehouse: Option, + database: Option, + schema: Option, +} + +impl Default for SessionObjectDetails { + fn default() -> Self { + Self { + warehouse: None, + database: None, + schema: None, + } + } +} + +// todo: make builder +impl Session { + fn new( + connection: Arc, + account_identifier: &str, + auth_type: AuthType, + object_details: SessionObjectDetails, + username: &str, + role: Option<&str>, + ) -> Self { + Self { + connection, + auth_tokens: Mutex::new(None), + auth_type, + account_identifier: account_identifier.to_string(), + username: username.to_string(), + role: role.map(str::to_string), + object_details, + } + } + /// Get cached token or request a new one if old one has expired. pub async fn get_token(&self) -> Result { let mut auth_tokens = self.auth_tokens.lock().await; @@ -213,7 +267,6 @@ impl Session { log::info!("Starting session with oauth authentication"); self.create(self.oauth_request_body(token)).await } - AuthType::None => Err(AuthError::AuthTypeUnset)?, }?; *auth_tokens = Some(tokens); } else if auth_tokens @@ -301,15 +354,15 @@ impl Session { body: LoginRequest, ) -> Result { let mut get_params = Vec::new(); - if let Some(warehouse) = &self.warehouse { + if let Some(warehouse) = &self.object_details.warehouse { get_params.push(("warehouse", warehouse.as_str())); } - if let Some(database) = &self.database { + if let Some(database) = &self.object_details.database { get_params.push(("databaseName", database.as_str())); } - if let Some(schema) = &self.schema { + if let Some(schema) = &self.object_details.schema { get_params.push(("schemaName", schema.as_str())); } From bbcdba2f1925d84d3a2fad04daa0e94ffb38b6ba Mon Sep 17 00:00:00 2001 From: Will Eaton Date: Thu, 6 Jun 2024 16:51:37 -0400 Subject: [PATCH 4/6] fixup example --- snowflake-api/examples/run_sql.rs | 46 ++++++++++++++++++------------- 1 file changed, 27 insertions(+), 19 deletions(-) diff --git a/snowflake-api/examples/run_sql.rs b/snowflake-api/examples/run_sql.rs index 18ec8a9..e9edf38 100644 --- a/snowflake-api/examples/run_sql.rs +++ b/snowflake-api/examples/run_sql.rs @@ -5,7 +5,9 @@ use arrow::util::pretty::pretty_format_batches; use clap::Parser; use std::fs; -use snowflake_api::{QueryResult, SnowflakeApi}; +use snowflake_api::{ + AuthArgs, CertificateArgs, PasswordArgs, QueryResult, SnowflakeApi, SnowflakeApiBuilder, +}; #[derive(clap::ValueEnum, Clone, Debug)] enum Output { @@ -67,25 +69,31 @@ async fn main() -> Result<()> { let mut api = match (&args.private_key, &args.password) { (Some(pkey), None) => { let pem = fs::read_to_string(pkey)?; - SnowflakeApi::with_certificate_auth( - &args.account_identifier, - args.warehouse.as_deref(), - args.database.as_deref(), - args.schema.as_deref(), - &args.username, - args.role.as_deref(), - &pem, - )? + SnowflakeApiBuilder::new(AuthArgs { + account_identifier: args.account_identifier, + warehouse: args.warehouse, + database: args.database, + schema: args.schema, + username: args.username, + role: args.role, + auth_type: snowflake_api::AuthType::Certificate(CertificateArgs { + private_key_pem: pem, + }), + }) + .build()? } - (None, Some(pwd)) => SnowflakeApi::with_password_auth( - &args.account_identifier, - args.warehouse.as_deref(), - args.database.as_deref(), - args.schema.as_deref(), - &args.username, - args.role.as_deref(), - pwd, - )?, + (None, Some(pwd)) => SnowflakeApiBuilder::new(AuthArgs { + account_identifier: args.account_identifier, + warehouse: args.warehouse, + database: args.database, + schema: args.schema, + username: args.username, + role: args.role, + auth_type: snowflake_api::AuthType::Password(PasswordArgs { + password: pwd.to_string(), + }), + }) + .build()?, _ => { panic!("Either private key path or password must be set") } From 09b78cc6bf67b5f64302fe0aa68a868d09b6fb57 Mon Sep 17 00:00:00 2001 From: Will Eaton Date: Thu, 6 Jun 2024 16:56:03 -0400 Subject: [PATCH 5/6] clippy; fix unnessecary derive --- snowflake-api/src/session.rs | 13 +------------ 1 file changed, 1 insertion(+), 12 deletions(-) diff --git a/snowflake-api/src/session.rs b/snowflake-api/src/session.rs index bedab67..2549069 100644 --- a/snowflake-api/src/session.rs +++ b/snowflake-api/src/session.rs @@ -203,24 +203,13 @@ impl SessionBuilder { } } -#[derive(Debug, Clone)] +#[derive(Debug, Default, Clone)] struct SessionObjectDetails { warehouse: Option, database: Option, schema: Option, } -impl Default for SessionObjectDetails { - fn default() -> Self { - Self { - warehouse: None, - database: None, - schema: None, - } - } -} - -// todo: make builder impl Session { fn new( connection: Arc, From 39996b27acc9a966b8744513438448305064ff2f Mon Sep 17 00:00:00 2001 From: Will Eaton Date: Thu, 6 Jun 2024 18:12:12 -0400 Subject: [PATCH 6/6] fix doctests --- snowflake-api/README.md | 15 +++++++++------ snowflake-api/src/lib.rs | 20 ++++++++++++++++++++ 2 files changed, 29 insertions(+), 6 deletions(-) diff --git a/snowflake-api/README.md b/snowflake-api/README.md index b08f911..7d6a22a 100644 --- a/snowflake-api/README.md +++ b/snowflake-api/README.md @@ -44,21 +44,24 @@ snowflake-api = "0.7.0" Check [examples](./examples) for working programs using the library. - ```rust use anyhow::Result; -use snowflake_api::{QueryResult, SnowflakeApi}; +use snowflake_api::{QueryResult, AuthArgs, PasswordArgs, AuthType, SnowflakeApi, SnowflakeApiBuilder}; async fn run_query(sql: &str) -> Result { - let mut api = SnowflakeApi::with_password_auth( + + let auth = AuthArgs::new( "ACCOUNT_IDENTIFIER", Some("WAREHOUSE"), Some("DATABASE"), Some("SCHEMA"), "USERNAME", Some("ROLE"), - "PASSWORD", - )?; + AuthType::Password(PasswordArgs { password: "password".to_string() }) + ); + + let mut api: SnowflakeApi = SnowflakeApiBuilder::new(auth) + .build()?; let res = api.exec(sql).await?; Ok(res) @@ -68,7 +71,7 @@ async fn run_query(sql: &str) -> Result { Or using environment variables: ```rust - use anyhow::Result; +use anyhow::Result; use snowflake_api::{QueryResult, SnowflakeApi}; async fn run_query(sql: &str) -> Result { diff --git a/snowflake-api/src/lib.rs b/snowflake-api/src/lib.rs index 39f344c..db7947d 100644 --- a/snowflake-api/src/lib.rs +++ b/snowflake-api/src/lib.rs @@ -193,6 +193,26 @@ pub struct AuthArgs { } impl AuthArgs { + pub fn new( + account_identifier: &str, + warehouse: Option<&str>, + database: Option<&str>, + schema: Option<&str>, + username: &str, + role: Option<&str>, + auth_type: AuthType, + ) -> Self { + Self { + account_identifier: account_identifier.to_string(), + warehouse: warehouse.map(str::to_string), + database: database.map(str::to_string), + schema: schema.map(str::to_string), + username: username.to_string(), + role: role.map(str::to_string), + auth_type, + } + } + pub fn from_env() -> Result { let auth_type = if let Ok(password) = std::env::var("SNOWFLAKE_PASSWORD") { Ok(AuthType::Password(PasswordArgs { password }))