Refinery migration support

Cargo.toml

@@ -5,4 +5,5 @@ members = [
     "snowflake-api",
     "snowflake-api/examples/tracing",
     "snowflake-api/examples/polars",
+    "snowflake-api/examples/migrate",
 ]

snowflake-api/Cargo.toml

@@ -13,10 +13,12 @@ version = "0.7.0"
 
 [features]
 default = ["cert-auth"]
-all = ["cert-auth", "polars"]
+all = ["cert-auth", "polars", "refinery"]
 cert-auth = ["dep:snowflake-jwt"]
 # support for conversion of arrow and json payloads to dataframes
 polars = ["dep:polars-core", "dep:polars-io"]
+refinery = ["dep:refinery-core", "dep:time", "dep:sqlparser"]
+
 
 [dependencies]
 arrow = "50"
@@ -40,9 +42,20 @@ snowflake-jwt = { version = "0.3.0", optional = true }
 thiserror = "1"
 url = "2"
 uuid = { version = "1", features = ["v4"] }
-polars-io = { version = ">=0.32", features = ["json", "ipc_streaming"], optional = true}
-polars-core = { version = ">=0.32", optional = true}
+polars-io = { version = ">=0.32", features = [
+    "json",
+    "ipc_streaming",
+], optional = true }
+polars-core = { version = ">=0.32", optional = true }
+tap = "1"
+
 
+refinery-core = { version = "=0.8.12", optional = true }
+time = { version = "0.3.5", features = [
+    "parsing",
+    "formatting",
+], optional = true }
+sqlparser = { version = "0.44", optional = true }
 
 [dev-dependencies]
 anyhow = "1"

snowflake-api/examples/migrate/Cargo.toml

@@ -0,0 +1,16 @@
+[package]
+name = "migrate-example"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+anyhow = "1.0.79"
+dotenv = "0.15.0"
+tokio = { version = "1.35.1", features = ["full"] }
+tracing = "0.1.40"
+tracing-subscriber = "0.3"
+
+snowflake-api = { path = "../../../snowflake-api", features = ["refinery"] }
+refinery = { version = "0.8.12" }

snowflake-api/examples/migrate/src/main.rs

@@ -0,0 +1,20 @@
+use anyhow::Result;
+use snowflake_api::{AuthArgs, SnowflakeApiBuilder};
+
+mod embedded {
+    use refinery::embed_migrations;
+    embed_migrations!("./tests/sql_migrations");
+}
+
+#[tokio::main]
+async fn main() -> Result<()> {
+    dotenv::dotenv().ok();
+    tracing_subscriber::fmt::init();
+
+    let auth_args = AuthArgs::from_env()?;
+    let mut conn = SnowflakeApiBuilder::new(auth_args).build()?;
+
+    embedded::migrations::runner().run_async(&mut conn).await?;
+
+    Ok(())
+}

snowflake-api/examples/migrate/tests/sql_migrations/V1__select.sql

@@ -0,0 +1,2 @@
+SELECT
+    1;

snowflake-api/examples/migrate/tests/sql_migrations/V2__create_table.sql

@@ -0,0 +1,36 @@
+-- create a dummy snowflake table
+CREATE
+OR REPLACE TABLE snowflake (
+    id INT,
+    NAME STRING,
+    age INT
+);
+-- insert some data into the table
+INSERT INTO
+    snowflake
+VALUES
+    (
+        1,
+        'John',
+        25
+    ),
+    (
+        2,
+        'Jane',
+        30
+    ),
+    (
+        3,
+        'Jim',
+        35
+    ),
+    (
+        4,
+        'Jill',
+        40
+    ),
+    (
+        5,
+        'Jack',
+        45
+    );

snowflake-api/examples/polars/src/main.rs

@@ -1,24 +1,13 @@
 use anyhow::Result;
 use polars::frame::DataFrame;
-use snowflake_api::{AuthArgs, AuthType, PasswordArgs, SnowflakeApiBuilder};
+use snowflake_api::{AuthArgs, SnowflakeApiBuilder};
 
 #[tokio::main]
 async fn main() -> Result<()> {
     dotenv::dotenv().ok();
     tracing_subscriber::fmt::init();
 
-    let auth_args = AuthArgs {
-        account_identifier: std::env::var("SNOWFLAKE_ACCOUNT").expect("SNOWFLAKE_ACCOUNT not set"),
-        warehouse: std::env::var("SNOWLFLAKE_WAREHOUSE").ok(),
-        database: std::env::var("SNOWFLAKE_DATABASE").ok(),
-        schema: std::env::var("SNOWFLAKE_SCHEMA").ok(),
-        username: std::env::var("SNOWFLAKE_USER").expect("SNOWFLAKE_USER not set"),
-        role: std::env::var("SNOWFLAKE_ROLE").ok(),
-        auth_type: AuthType::Password(PasswordArgs {
-            password: std::env::var("SNOWFLAKE_PASSWORD").expect("SNOWFLAKE_PASSWORD not set"),
-        }),
-    };
-
+    let auth_args = AuthArgs::from_env()?;
     let api = SnowflakeApiBuilder::new(auth_args).build()?;
 
     // run a query that returns a tabular arrow response

snowflake-api/examples/tracing/src/main.rs

@@ -4,7 +4,7 @@ use opentelemetry::global;
 use opentelemetry_otlp::WithExportConfig;
 
 use snowflake_api::connection::Connection;
-use snowflake_api::{AuthArgs, AuthType, PasswordArgs, QueryResult, SnowflakeApiBuilder};
+use snowflake_api::{AuthArgs, QueryResult, SnowflakeApiBuilder};
 use tracing_subscriber::layer::SubscriberExt;
 
 use opentelemetry::KeyValue;
@@ -35,18 +35,7 @@ async fn main() -> Result<()> {
 
     dotenv::dotenv().ok();
 
-    let auth_args = AuthArgs {
-        account_identifier: std::env::var("SNOWFLAKE_ACCOUNT").expect("SNOWFLAKE_ACCOUNT not set"),
-        warehouse: std::env::var("SNOWLFLAKE_WAREHOUSE").ok(),
-        database: std::env::var("SNOWFLAKE_DATABASE").ok(),
-        schema: std::env::var("SNOWFLAKE_SCHEMA").ok(),
-        username: std::env::var("SNOWFLAKE_USER").expect("SNOWFLAKE_USER not set"),
-        role: std::env::var("SNOWFLAKE_ROLE").ok(),
-        auth_type: AuthType::Password(PasswordArgs {
-            password: std::env::var("SNOWFLAKE_PASSWORD").expect("SNOWFLAKE_PASSWORD not set"),
-        }),
-    };
-
+    let auth_args = AuthArgs::from_env()?;
     let mut client = Connection::default_client_builder()?;
     client = client
         .with_init(Extension(OtelName(std::borrow::Cow::Borrowed(

snowflake-api/src/lib.rs

@@ -42,6 +42,8 @@ use crate::responses::{
 };
 
 pub mod connection;
+#[cfg(feature = "refinery")]
+mod migration;
 #[cfg(feature = "polars")]
 mod polars;
 mod requests;
@@ -91,6 +93,9 @@ pub enum SnowflakeApiError {
 
     #[error("Unexpected API response")]
     UnexpectedResponse,
+
+    #[error("Error parsing query: {0}")]
+    QueryParserError(sqlparser::parser::ParserError),
 }
 
 /// Even if Arrow is specified as a return type non-select queries
@@ -200,6 +205,34 @@ pub struct CertificateArgs {
     pub private_key_pem: String,
 }
 
+impl AuthArgs {
+    /// Construct `AuthArgs` from environment variables
+    pub fn from_env() -> Result<Self, SnowflakeApiError> {
+        let auth_type = match std::env::var("SNOWFLAKE_PASSWORD") {
+            Ok(_) => AuthType::Password(PasswordArgs {
+                password: std::env::var("SNOWFLAKE_PASSWORD")
+                    .map_err(|_| SnowflakeApiError::AuthError(AuthError::MissingPassword))?,
+            }),
+            Err(_) => AuthType::Certificate(CertificateArgs {
+                private_key_pem: std::env::var("SNOWFLAKE_PRIVATE_KEY_PEM")
+                    .map_err(|_| SnowflakeApiError::AuthError(AuthError::MissingPrivateKey))?,
+            }),
+        };
+
+        Ok(Self {
+            account_identifier: std::env::var("SNOWFLAKE_ACCOUNT")
+                .map_err(|_| SnowflakeApiError::AuthError(AuthError::MissingAccount))?,
+            warehouse: std::env::var("SNOWLFLAKE_WAREHOUSE").ok(),
+            database: std::env::var("SNOWFLAKE_DATABASE").ok(),
+            schema: std::env::var("SNOWFLAKE_SCHEMA").ok(),
+            username: std::env::var("SNOWFLAKE_USER")
+                .map_err(|_| SnowflakeApiError::AuthError(AuthError::MissingUsername))?,
+            role: std::env::var("SNOWFLAKE_ROLE").ok(),
+            auth_type,
+        })
+    }
+}
+
 #[must_use]
 pub struct SnowflakeApiBuilder {
     pub auth: AuthArgs,