diff --git a/Cargo.lock b/Cargo.lock index f1e318a8ef7..df2768e4996 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4140,6 +4140,7 @@ dependencies = [ "feldera-observability", "feldera-samply", "feldera-size-of", + "feldera-snowflake", "feldera-sqllib", "feldera-storage", "feldera-types", @@ -4617,6 +4618,15 @@ dependencies = [ "unicode-xid", ] +[[package]] +name = "des" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ffdd80ce8ce993de27e9f063a444a4d53ce8e8db4c1f00cc03af5ad5a9867a1e" +dependencies = [ + "cipher", +] + [[package]] name = "deunicode" version = "1.6.2" @@ -5480,6 +5490,36 @@ dependencies = [ "xxhash-rust", ] +[[package]] +name = "feldera-snowflake" +version = "0.324.0" +dependencies = [ + "anyhow", + "arrow", + "async-compression", + "async-stream", + "base64 0.22.1", + "bytes", + "chrono", + "dbsp", + "feldera-adapterlib", + "feldera-types", + "flate2", + "futures-util", + "jsonwebtoken", + "log", + "pkcs8 0.10.2", + "rand 0.8.6", + "reqwest 0.12.24", + "rsa", + "serde", + "serde_json", + "sha2 0.10.9", + "tokio", + "tokio-util", + "uuid", +] + [[package]] name = "feldera-sqllib" version = "0.324.0" @@ -8950,6 +8990,7 @@ dependencies = [ "aes", "cbc", "der 0.7.10", + "des", "pbkdf2", "scrypt", "sha2 0.10.9", diff --git a/Cargo.toml b/Cargo.toml index 22f32397157..3433d9be7cc 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -25,6 +25,7 @@ members = [ "sql-to-dbsp-compiler/lib/readers", "crates/datagen", "crates/iceberg", + "crates/snowflake", "crates/storage", "crates/rest-api", "crates/ir", @@ -61,6 +62,7 @@ arrow-json = "58" arrow-schema = "58" ascii_table = "=4.0.2" async-channel = "2.3.1" +async-compression = { version = "0.4.42", features = ["gzip", "tokio"] } async-nats = "0.47.0" async-std = "1.12.0" async-stream = "0.3.5" @@ -143,6 +145,7 @@ feldera-cloud1-client = "0.1.2" feldera-datagen = { path = "crates/datagen" } feldera-fxp = { version = "0.324.0", path = "crates/fxp", features = ["dbsp"] } feldera-iceberg = { path = "crates/iceberg" } +feldera-snowflake = { path = "crates/snowflake" } feldera-observability = { version = "0.324.0", path = "crates/feldera-observability" } feldera-macros = { version = "0.324.0", path = "crates/feldera-macros" } feldera-sqllib = { version = "0.324.0", path = "crates/sqllib" } @@ -213,6 +216,7 @@ paste = "1.0.12" petgraph = "0.6.0" pg-client-config = "0.1.2" pin-project-lite = "0.2.16" +pkcs8 = { version = "0.10.2", features = ["3des"] } postgres = "0.19.10" postgres-openssl = "0.5.1" postgresql_embedded = { version = "0.20.0", features = ["bundled"] } @@ -244,6 +248,7 @@ rmpv = "1.3.0" # Pin `roaring` exactly because the crate only supports the current stable # Rust and may start requiring newer compiler features in patch releases. roaring = "=0.11.3" +rsa = { version = "0.9.10", features = ["pkcs5"] } rstest = "0.15" # Make sure this is the same rustls version used by the `tonic` crate. # See the `ensure_default_crypto_provider` function. diff --git a/crates/adapters/Cargo.toml b/crates/adapters/Cargo.toml index 576af7b0a4a..f8e4b79e249 100644 --- a/crates/adapters/Cargo.toml +++ b/crates/adapters/Cargo.toml @@ -19,6 +19,7 @@ default = [ "with-kafka", "with-deltalake", "with-iceberg", + "with-snowflake", "with-avro", "with-nexmark", "with-pubsub", @@ -39,6 +40,7 @@ with-deltalake = [ "parquet/object_store", ] with-iceberg = ["feldera-iceberg"] +with-snowflake = ["feldera-snowflake"] with-pubsub = ["google-cloud-pubsub", "google-cloud-gax"] with-avro = [ "apache-avro", @@ -82,6 +84,7 @@ feldera-observability = { workspace = true } feldera-samply = { workspace = true } feldera-storage = { workspace = true } feldera-iceberg = { workspace = true, optional = true } +feldera-snowflake = { workspace = true, optional = true } awc = { workspace = true, features = [ "compress-gzip", "compress-brotli", diff --git a/crates/adapters/src/controller.rs b/crates/adapters/src/controller.rs index b9cd2600227..34904a1caae 100644 --- a/crates/adapters/src/controller.rs +++ b/crates/adapters/src/controller.rs @@ -8415,7 +8415,7 @@ impl ControllerInner { } TransactionState::Started { tid, - start, + start: _, processed_records, } if next.is_none() || next != open => { // The next step belongs to a different transaction, or the diff --git a/crates/adapters/src/integrated.rs b/crates/adapters/src/integrated.rs index fd84911875a..904125fd6c0 100644 --- a/crates/adapters/src/integrated.rs +++ b/crates/adapters/src/integrated.rs @@ -134,6 +134,10 @@ pub fn create_integrated_input_endpoint( consumer, )) } + #[cfg(feature = "with-snowflake")] + TransportConfig::SnowflakeInput(config) => Box::new( + feldera_snowflake::SnowflakeInputEndpoint::new(endpoint_name, config, consumer), + ), TransportConfig::PostgresInput(config) => { Box::new(PostgresInputEndpoint::new(endpoint_name, config, consumer)) } diff --git a/crates/adapters/src/transport.rs b/crates/adapters/src/transport.rs index f634e86c0e4..7c8871dff95 100644 --- a/crates/adapters/src/transport.rs +++ b/crates/adapters/src/transport.rs @@ -128,6 +128,7 @@ pub fn input_transport_config_to_endpoint( | TransportConfig::HttpOutput(_) | TransportConfig::RedisOutput(_) | TransportConfig::IcebergInput(_) + | TransportConfig::SnowflakeInput(_) | TransportConfig::NullOutput => return Ok(None), }; Ok(Some(endpoint)) diff --git a/crates/feldera-types/src/config.rs b/crates/feldera-types/src/config.rs index 44f3ddface3..83fef0ced04 100644 --- a/crates/feldera-types/src/config.rs +++ b/crates/feldera-types/src/config.rs @@ -25,6 +25,7 @@ use crate::transport::postgres::{ use crate::transport::pubsub::PubSubInputConfig; use crate::transport::redis::RedisOutputConfig; use crate::transport::s3::S3InputConfig; +use crate::transport::snowflake::SnowflakeReaderConfig; use crate::transport::url::UrlInputConfig; use core::fmt; use feldera_ir::{MirNode, MirNodeId}; @@ -1879,6 +1880,7 @@ pub enum TransportConfig { RedisOutput(RedisOutputConfig), // Prevent rust from complaining about large size difference between enum variants. IcebergInput(Box), + SnowflakeInput(Box), PostgresInput(PostgresReaderConfig), PostgresCdcInput(PostgresCdcReaderConfig), PostgresOutput(PostgresWriterConfig), @@ -1912,6 +1914,7 @@ impl TransportConfig { TransportConfig::DeltaTableOutput(_) => "delta_table_output".to_string(), TransportConfig::DynamoDBOutput(_) => "dynamodb_output".to_string(), TransportConfig::IcebergInput(_) => "iceberg_input".to_string(), + TransportConfig::SnowflakeInput(_) => "snowflake_input".to_string(), TransportConfig::PostgresInput(_) => "postgres_input".to_string(), TransportConfig::PostgresCdcInput(_) => "postgres_cdc_input".to_string(), TransportConfig::PostgresOutput(_) => "postgres_output".to_string(), diff --git a/crates/feldera-types/src/transport.rs b/crates/feldera-types/src/transport.rs index 4fc44d779a4..e182c3272e4 100644 --- a/crates/feldera-types/src/transport.rs +++ b/crates/feldera-types/src/transport.rs @@ -13,4 +13,5 @@ pub mod postgres; pub mod pubsub; pub mod redis; pub mod s3; +pub mod snowflake; pub mod url; diff --git a/crates/feldera-types/src/transport/snowflake.rs b/crates/feldera-types/src/transport/snowflake.rs new file mode 100644 index 00000000000..ee51e26ec8f --- /dev/null +++ b/crates/feldera-types/src/transport/snowflake.rs @@ -0,0 +1,399 @@ +use serde::{Deserialize, Serialize}; +use std::collections::BTreeMap; +use utoipa::ToSchema; + +/// Snowflake table read mode. +#[derive(Debug, Clone, Copy, Default, Eq, PartialEq, Deserialize, Serialize, ToSchema)] +pub enum SnowflakeIngestMode { + /// Read a snapshot of the table and stop. + #[default] + #[serde(rename = "snapshot")] + Snapshot, +} + +/// Snowflake authentication method. +#[derive(Debug, Clone, Copy, Default, Eq, PartialEq, Deserialize, Serialize, ToSchema)] +pub enum SnowflakeAuthenticator { + /// Key-pair authentication using a JWT signed with an RSA private key. + #[default] + #[serde(rename = "SNOWFLAKE_JWT")] + SnowflakeJwt, +} + +/// Snowflake input connector transaction mode. +/// +/// Determines whether the connector wraps snapshot ingestion in a Feldera +/// transaction. +/// +/// * `none` - the connector does not start or commit transactions. +/// * `snapshot` - the connector ingests the snapshot in one Feldera +/// transaction. +#[derive(Debug, Clone, Copy, Eq, PartialEq, Deserialize, Serialize, ToSchema, Default)] +pub enum SnowflakeTransactionMode { + /// Do not request Feldera transaction boundaries. + #[default] + #[serde(rename = "none")] + None, + + /// Ingest the snapshot in one Feldera transaction. + /// + /// If snapshot ingestion fails after records have been queued, the partial + /// transaction is committed because connector transactions do not support + /// rollback. Atomic rollback on failure is not currently supported. + #[serde(rename = "snapshot")] + Snapshot, +} + +/// Representation of scaled Snowflake `NUMBER` values. +#[derive(Debug, Clone, Copy, Default, Eq, PartialEq, Deserialize, Serialize, ToSchema)] +pub enum SnowflakeNumberMode { + /// Convert scaled `NUMBER` values to Arrow Decimal128, preserving precision. + #[default] + #[serde(rename = "decimal")] + Decimal, + + /// Convert scaled `NUMBER` values to Arrow Float64. + /// + /// This can lose precision. Scale-zero `NUMBER` values remain integral. + #[serde(rename = "double")] + Double, +} + +const DEFAULT_MAX_CONCURRENT_READERS: u32 = 4; + +fn default_num_parsers() -> u32 { + 4 +} + +/// Snowflake input connector configuration. +#[derive(Debug, Clone, Eq, PartialEq, Deserialize, Serialize, ToSchema)] +pub struct SnowflakeReaderConfig { + /// Snowflake account identifier. + /// + /// Uses the same account identifier accepted by Snowflake drivers, for + /// example `"org-account"` or `"xy12345.us-east-1"`. + pub account: String, + + /// Snowflake login name. + pub user: String, + + /// Snowflake authenticator. + /// + /// Defaults to `"SNOWFLAKE_JWT"`. + #[serde(default)] + pub authenticator: SnowflakeAuthenticator, + + /// Snowflake role to use for the session. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub role: Option, + + /// Snowflake warehouse to use for the session. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub warehouse: Option, + + /// Snowflake database to use for the session. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub database: Option, + + /// Snowflake schema to use for the session. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub schema: Option, + + /// RSA private key file for `SNOWFLAKE_JWT` authentication. + /// + /// The file must contain a PEM-encoded PKCS#8 RSA private key. + pub private_key_file: String, + + /// RSA private key passphrase for `SNOWFLAKE_JWT` authentication. + /// + /// Required when `private_key_file` contains an encrypted PKCS#8 private + /// key. + #[serde( + default, + alias = "private_key_passphrase", + alias = "private_key_file_password", + skip_serializing_if = "Option::is_none" + )] + pub private_key_file_pwd: Option, + + /// Source table name. + /// + /// The value can be a fully-qualified Snowflake table name, for example + /// `"MY_DATABASE.MY_SCHEMA.MY_TABLE"`, or an unqualified table name when + /// `database` and `schema` are configured on the Snowflake session. + pub table: String, + + /// Map Feldera table columns to Snowflake source columns. + /// + /// Each key is the bare name of a column in the Feldera SQL table this + /// connector is attached to (without SQL quote delimiters). Each value is + /// the corresponding Snowflake SQL column identifier. The connector aliases + /// mapped columns back to their Feldera names in the generated snapshot + /// query. Quote a value as a Snowflake SQL identifier when its spelling is + /// case-sensitive. + /// + /// For example, `{"uuid": "UUID"}` reads the Snowflake column `UUID` + /// into the Feldera column `"uuid"`. + #[serde(default, skip_serializing_if = "BTreeMap::is_empty")] + pub column_mapping: BTreeMap, + + /// Representation of scaled Snowflake `NUMBER` values. + /// + /// Defaults to `"decimal"`, which preserves precision. Set to `"double"` + /// to convert these values to double-precision floating point. + #[serde(default)] + pub number_mode: SnowflakeNumberMode, + + /// Table read mode. + /// + /// Only `"snapshot"` is supported. + #[serde(default)] + pub mode: SnowflakeIngestMode, + + /// Transaction mode. + /// + /// Determines whether the connector wraps snapshot ingestion in a Feldera + /// transaction. Defaults to `"none"`. + #[serde(default)] + pub transaction_mode: SnowflakeTransactionMode, + + /// Optional row filter. + /// + /// When specified, only rows that satisfy this predicate are included in + /// the snapshot. The predicate is appended to the generated Snowflake query + /// as the body of a `WHERE` clause. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub snapshot_filter: Option, + + /// Number of parallel tasks used to deserialize Arrow record batches. + /// + /// Recommended range: 1–10. Default: 4. + #[serde(default = "default_num_parsers")] + #[schema(minimum = 1)] + pub num_parsers: u32, + + /// Maximum number of concurrent reads of Snowflake result chunks. + /// + /// Snowflake returns large result sets as multiple Arrow IPC chunks stored + /// behind pre-signed URLs. Increasing this value can improve snapshot + /// throughput when chunk download/decompression is the bottleneck. + /// + /// Default: 4. + #[serde(default, skip_serializing_if = "Option::is_none")] + #[schema(minimum = 1)] + pub max_concurrent_readers: Option, +} + +impl SnowflakeReaderConfig { + pub fn validate(&self) -> Result<(), String> { + for (name, value) in [ + ("account", &self.account), + ("user", &self.user), + ("private_key_file", &self.private_key_file), + ("table", &self.table), + ] { + if value.trim().is_empty() { + return Err(format!("'{name}' must not be empty")); + } + } + + for (feldera_column, snowflake_column) in &self.column_mapping { + if feldera_column.trim().is_empty() { + return Err("'column_mapping' contains an empty Feldera column name".to_string()); + } + if snowflake_column.trim().is_empty() { + return Err(format!( + "'column_mapping' contains an empty Snowflake column name for Feldera column '{feldera_column}'" + )); + } + } + + for (name, value) in [ + ("role", &self.role), + ("warehouse", &self.warehouse), + ("database", &self.database), + ("schema", &self.schema), + ("private_key_file_pwd", &self.private_key_file_pwd), + ] { + if value.as_ref().is_some_and(|value| value.trim().is_empty()) { + return Err(format!("'{name}' must not be empty when specified")); + } + } + + if self + .snapshot_filter + .as_ref() + .is_some_and(|filter| filter.trim().is_empty()) + { + return Err("'snapshot_filter' must not be empty when specified".to_string()); + } + + if self + .max_concurrent_readers + .is_some_and(|readers| readers == 0) + { + return Err("'max_concurrent_readers' must be greater than 0".to_string()); + } + + if self.num_parsers == 0 { + return Err("'num_parsers' must be greater than 0".to_string()); + } + + Ok(()) + } + + pub fn max_concurrent_readers(&self) -> usize { + self.max_concurrent_readers + .unwrap_or(DEFAULT_MAX_CONCURRENT_READERS) as usize + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn minimal_config() -> SnowflakeReaderConfig { + SnowflakeReaderConfig { + account: "org-account".to_string(), + user: "svc_user".to_string(), + authenticator: SnowflakeAuthenticator::SnowflakeJwt, + role: None, + warehouse: None, + database: None, + schema: None, + private_key_file: "/secrets/key.p8".to_string(), + private_key_file_pwd: None, + table: "DB.SCHEMA.TABLE".to_string(), + column_mapping: BTreeMap::new(), + number_mode: SnowflakeNumberMode::Decimal, + mode: SnowflakeIngestMode::Snapshot, + transaction_mode: SnowflakeTransactionMode::None, + snapshot_filter: None, + num_parsers: default_num_parsers(), + max_concurrent_readers: None, + } + } + + #[test] + fn validates_minimal_config() { + let config = minimal_config(); + config.validate().unwrap(); + assert_eq!(config.max_concurrent_readers(), 4); + assert_eq!(config.num_parsers, 4); + } + + #[test] + fn rejects_zero_concurrent_readers() { + let mut config = minimal_config(); + config.max_concurrent_readers = Some(0); + assert_eq!( + config.validate().unwrap_err(), + "'max_concurrent_readers' must be greater than 0" + ); + } + + #[test] + fn rejects_zero_parsers() { + let mut config = minimal_config(); + config.num_parsers = 0; + assert_eq!( + config.validate().unwrap_err(), + "'num_parsers' must be greater than 0" + ); + } + + #[test] + fn parses_snapshot_transaction_mode() { + let config: SnowflakeReaderConfig = serde_json::from_str( + r#"{ + "account": "org-account", + "user": "svc_user", + "private_key_file": "/secrets/key.p8", + "table": "DB.SCHEMA.TABLE", + "transaction_mode": "snapshot" + }"#, + ) + .unwrap(); + + assert_eq!(config.transaction_mode, SnowflakeTransactionMode::Snapshot); + assert_eq!(config.number_mode, SnowflakeNumberMode::Decimal); + config.validate().unwrap(); + } + + #[test] + fn accepts_private_key_passphrase_aliases() { + let config: SnowflakeReaderConfig = serde_json::from_str( + r#"{ + "account": "org-account", + "user": "svc_user", + "private_key_file": "/secrets/key.p8", + "private_key_passphrase": "secret", + "table": "DB.SCHEMA.TABLE" + }"#, + ) + .unwrap(); + + assert_eq!(config.private_key_file_pwd.as_deref(), Some("secret")); + config.validate().unwrap(); + } + + #[test] + fn parses_column_mapping() { + let config: SnowflakeReaderConfig = serde_json::from_str( + r#"{ + "account": "org-account", + "user": "svc_user", + "private_key_file": "/secrets/key.p8", + "table": "DB.SCHEMA.TABLE", + "column_mapping": { + "uuid": "UUID" + } + }"#, + ) + .unwrap(); + + assert_eq!(config.column_mapping.get("uuid").unwrap(), "UUID"); + config.validate().unwrap(); + } + + #[test] + fn parses_number_mode() { + let config: SnowflakeReaderConfig = serde_json::from_str( + r#"{ + "account": "org-account", + "user": "svc_user", + "private_key_file": "/secrets/key.p8", + "table": "DB.SCHEMA.TABLE", + "number_mode": "double" + }"#, + ) + .unwrap(); + + assert_eq!(config.number_mode, SnowflakeNumberMode::Double); + } + + #[test] + fn rejects_empty_column_mapping_names() { + let mut config = minimal_config(); + config + .column_mapping + .insert("uuid".to_string(), String::new()); + assert!( + config + .validate() + .unwrap_err() + .contains("Snowflake column name") + ); + + let mut config = minimal_config(); + config + .column_mapping + .insert(" ".to_string(), "UUID".to_string()); + assert!( + config + .validate() + .unwrap_err() + .contains("Feldera column name") + ); + } +} diff --git a/crates/pipeline-manager/src/db/types/program.rs b/crates/pipeline-manager/src/db/types/program.rs index 4057794b639..96882633bc4 100644 --- a/crates/pipeline-manager/src/db/types/program.rs +++ b/crates/pipeline-manager/src/db/types/program.rs @@ -729,6 +729,7 @@ pub fn generate_program_info( | TransportConfig::DeltaTableInput(_) | TransportConfig::PostgresInput(_) | TransportConfig::IcebergInput(_) + | TransportConfig::SnowflakeInput(_) | TransportConfig::Datagen(_) | TransportConfig::Nexmark(_) | TransportConfig::EmptyInput => {} diff --git a/crates/rest-api/build.rs b/crates/rest-api/build.rs index a2fe10ce9f4..2a668b30861 100644 --- a/crates/rest-api/build.rs +++ b/crates/rest-api/build.rs @@ -132,6 +132,22 @@ fn type_replacement() -> Vec<(&'static str, &'static str)> { "DeltaTableWriterConfig", "feldera_types::transport::delta_table::DeltaTableWriterConfig", ), + ( + "SnowflakeIngestMode", + "feldera_types::transport::snowflake::SnowflakeIngestMode", + ), + ( + "SnowflakeAuthenticator", + "feldera_types::transport::snowflake::SnowflakeAuthenticator", + ), + ( + "SnowflakeTransactionMode", + "feldera_types::transport::snowflake::SnowflakeTransactionMode", + ), + ( + "SnowflakeReaderConfig", + "feldera_types::transport::snowflake::SnowflakeReaderConfig", + ), ( "DynamoDBWriterConfig", "feldera_types::transport::dynamodb::DynamoDBWriterConfig", diff --git a/crates/snowflake/Cargo.toml b/crates/snowflake/Cargo.toml new file mode 100644 index 00000000000..25849937ed8 --- /dev/null +++ b/crates/snowflake/Cargo.toml @@ -0,0 +1,39 @@ +[package] +name = "feldera-snowflake" +edition = { workspace = true } +version = { workspace = true } +homepage = { workspace = true } +repository = { workspace = true } +license = { workspace = true } +authors = { workspace = true } +rust-version = { workspace = true } +readme = { workspace = true } +publish = false + +[dependencies] +feldera-types = { workspace = true } +feldera-adapterlib = { workspace = true } +dbsp = { workspace = true } +anyhow = { workspace = true, features = ["backtrace"] } +async-compression = { workspace = true } +async-stream = { workspace = true } +bytes = { workspace = true } +tokio = { workspace = true, features = ["io-util", "macros", "rt", "sync", "time"] } +tokio-util = { workspace = true, features = ["io"] } +chrono = { workspace = true, features = ["clock"] } +log = { workspace = true } +serde = { workspace = true, features = ["derive"] } +serde_json = { workspace = true } +reqwest = { workspace = true, features = ["json", "stream", "rustls-tls-native-roots"] } +arrow = { workspace = true, features = ["ipc"] } +futures-util = { workspace = true } +base64 = { workspace = true } +jsonwebtoken = { workspace = true } +pkcs8 = { workspace = true } +rand = { workspace = true } +rsa = { workspace = true } +sha2 = { workspace = true } +uuid = { workspace = true, features = ["v4"] } + +[dev-dependencies] +flate2 = { workspace = true } diff --git a/crates/snowflake/src/arrow.rs b/crates/snowflake/src/arrow.rs new file mode 100644 index 00000000000..66838b23159 --- /dev/null +++ b/crates/snowflake/src/arrow.rs @@ -0,0 +1,538 @@ +use std::io::Cursor; + +use anyhow::{anyhow, bail, Context, Result as AnyResult}; +use arrow::{ + array::{ + Array, ArrayRef, Decimal128Array, Float64Array, Int16Array, Int32Array, Int64Array, + Int8Array, StructArray, Time64NanosecondArray, TimestampMicrosecondArray, + }, + datatypes::{DataType, Field, Schema, TimeUnit}, + ipc::reader::StreamReader, + record_batch::RecordBatch, +}; +use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine as _}; +use feldera_types::transport::snowflake::SnowflakeNumberMode; + +/// Decode an inline Snowflake `rowsetBase64` payload into Arrow record batches. +pub(crate) fn decode_base64_ipc_stream( + rowset_base64: &str, + number_mode: SnowflakeNumberMode, +) -> AnyResult> { + let bytes = BASE64_STANDARD + .decode(rowset_base64) + .context("invalid base64 in Snowflake Arrow rowset")?; + decode_ipc_stream(&bytes, number_mode) +} + +/// Decode a raw Snowflake Arrow chunk into Arrow record batches. +pub(crate) fn decode_ipc_stream( + bytes: &[u8], + number_mode: SnowflakeNumberMode, +) -> AnyResult> { + if bytes.is_empty() { + return Ok(Vec::new()); + } + + let reader = StreamReader::try_new(Cursor::new(bytes), None) + .context("error opening Snowflake Arrow IPC stream")?; + + reader + .map(|batch| batch.context("error reading Snowflake Arrow IPC record batch")) + .map(|batch| batch.and_then(|batch| normalize_snowflake_batch(batch, number_mode))) + .collect() +} + +/// Convert Snowflake-specific Arrow encodings into standard Arrow types understood by Feldera. +pub(crate) fn normalize_snowflake_batch( + batch: RecordBatch, + number_mode: SnowflakeNumberMode, +) -> AnyResult { + let schema = batch.schema(); + let mut fields = Vec::with_capacity(schema.fields().len()); + let mut columns = Vec::with_capacity(batch.num_columns()); + + for (field, column) in schema.fields().iter().zip(batch.columns()) { + let logical_type = field.metadata().get("logicalType").map(String::as_str); + + // Snowflake encodes scaled FIXED values as unscaled signed integers. Its C/C++ client + // applies the scale during conversion as well: + // https://github.com/snowflakedb/libsnowflakeclient/blob/master/cpp/lib/ArrowChunkIterator.cpp + if logical_type == Some("FIXED") && field_scale(field)? != 0 { + let normalized: ArrayRef = match number_mode { + SnowflakeNumberMode::Decimal => { + std::sync::Arc::new(normalize_fixed(field.as_ref(), column.as_ref())?) + } + SnowflakeNumberMode::Double => { + std::sync::Arc::new(normalize_fixed_to_double(field.as_ref(), column.as_ref())?) + } + }; + fields.push( + Field::new( + field.name(), + normalized.data_type().clone(), + field.is_nullable(), + ) + .with_metadata(field.metadata().clone()), + ); + columns.push(normalized); + } else if matches!( + logical_type, + Some("TIMESTAMP_NTZ" | "TIMESTAMP_LTZ" | "TIMESTAMP_TZ") + ) { + let timezone = match logical_type { + Some("TIMESTAMP_LTZ" | "TIMESTAMP_TZ") => Some("UTC".into()), + _ => None, + }; + let timestamps = normalize_timestamp(field.as_ref(), column.as_ref(), logical_type)? + .with_timezone_opt(timezone.clone()); + fields.push( + Field::new( + field.name(), + DataType::Timestamp(TimeUnit::Microsecond, timezone), + field.is_nullable(), + ) + .with_metadata(field.metadata().clone()), + ); + columns.push(std::sync::Arc::new(timestamps) as _); + } else if logical_type == Some("TIME") { + let times = normalize_time(field.as_ref(), column.as_ref())?; + fields.push( + Field::new( + field.name(), + DataType::Time64(TimeUnit::Nanosecond), + field.is_nullable(), + ) + .with_metadata(field.metadata().clone()), + ); + columns.push(std::sync::Arc::new(times) as _); + } else { + fields.push(field.as_ref().clone()); + columns.push(column.clone()); + } + } + + let schema = Schema::new(fields).with_metadata(schema.metadata().clone()); + RecordBatch::try_new(std::sync::Arc::new(schema), columns) + .context("error constructing normalized Snowflake Arrow record batch") +} + +fn normalize_fixed(field: &Field, column: &dyn Array) -> AnyResult { + let precision = field_metadata::(field, "precision")?; + let scale = field_metadata::(field, "scale")?; + if let Some(values) = column.as_any().downcast_ref::() { + return values + .clone() + .with_precision_and_scale(precision, scale) + .with_context(|| { + format!( + "invalid precision or scale metadata for Snowflake FIXED column '{}'", + field.name() + ) + }); + } + + let values = (0..column.len()) + .map(|row| fixed_integer_at(column, row, field.name())) + .collect::>>()?; + + Decimal128Array::from(values) + .with_precision_and_scale(precision, scale) + .with_context(|| { + format!( + "invalid precision or scale metadata for Snowflake FIXED column '{}'", + field.name() + ) + }) +} + +fn normalize_fixed_to_double(field: &Field, column: &dyn Array) -> AnyResult { + let scale = field_metadata::(field, "scale")?; + let divisor = 10_f64.powi(scale); + + (0..column.len()) + .map(|row| { + fixed_integer_at(column, row, field.name()) + .map(|value| value.map(|value| value as f64 / divisor)) + }) + .collect() +} + +fn fixed_integer_at(column: &dyn Array, row: usize, field_name: &str) -> AnyResult> { + if let Some(values) = column.as_any().downcast_ref::() { + return Ok((!values.is_null(row)).then(|| values.value(row))); + } + + signed_integer_at(column, row, field_name).map(|value| value.map(i128::from)) +} + +fn normalize_timestamp( + field: &Field, + column: &dyn Array, + logical_type: Option<&str>, +) -> AnyResult { + // Snowflake's reference implementation documents both the current structured timestamp + // encoding and the legacy TIMESTAMP_TZ layout handled below: + // https://github.com/snowflakedb/libsnowflakeclient/blob/master/cpp/lib/ArrowChunkIterator.cpp + if let Some(values) = column.as_any().downcast_ref::() { + let epoch = struct_child::(values, field.name(), "epoch")?; + let fraction = struct_child::(values, field.name(), "fraction")?; + let legacy_timestamp_tz = + logical_type == Some("TIMESTAMP_TZ") && values.column_by_name("timezone").is_none(); + + return (0..values.len()) + .map(|row| { + if values.is_null(row) || epoch.is_null(row) || fraction.is_null(row) { + return Ok(None); + } + + // Modern structured timestamps store seconds since the Unix epoch plus a + // nanosecond fraction. In Snowflake's legacy TIMESTAMP_TZ struct, `epoch` + // instead contains the complete scaled timestamp and `fraction` contains + // the display timezone. + let micros = if legacy_timestamp_tz { + scaled_integer_to_micros( + epoch.value(row), + field_scale(field)?, + field.name(), + row, + )? + } else { + epoch + .value(row) + .checked_mul(1_000_000) + .and_then(|value| value.checked_add(i64::from(fraction.value(row)) / 1_000)) + .ok_or_else(|| timestamp_out_of_range(field.name(), row))? + }; + Ok(Some(micros)) + }) + .collect::>(); + } + + let scale = field_scale(field)?; + (0..column.len()) + .map(|row| { + signed_integer_at(column, row, field.name())? + .map(|value| scaled_integer_to_micros(value, scale, field.name(), row)) + .transpose() + }) + .collect::>() +} + +fn normalize_time(field: &Field, column: &dyn Array) -> AnyResult { + let scale = field_scale(field)?; + if scale > 9 { + bail!( + "Snowflake TIME column '{}' has unsupported scale {scale}", + field.name() + ); + } + let multiplier = 10_i64.pow(9 - scale); + + (0..column.len()) + .map(|row| { + signed_integer_at(column, row, field.name())? + .map(|value| { + value.checked_mul(multiplier).ok_or_else(|| { + anyhow!( + "Snowflake TIME value in column '{}' at row {row} is out of range", + field.name() + ) + }) + }) + .transpose() + }) + .collect::>() +} + +fn signed_integer_at(column: &dyn Array, row: usize, field_name: &str) -> AnyResult> { + macro_rules! value { + ($array:ty) => { + if let Some(values) = column.as_any().downcast_ref::<$array>() { + return Ok((!values.is_null(row)).then(|| i64::from(values.value(row)))); + } + }; + } + + value!(Int8Array); + value!(Int16Array); + value!(Int32Array); + value!(Int64Array); + bail!( + "Snowflake column '{field_name}' has unsupported physical type {}", + column.data_type() + ) +} + +fn field_scale(field: &Field) -> AnyResult { + field_metadata(field, "scale") +} + +fn field_metadata(field: &Field, key: &str) -> AnyResult +where + T: std::str::FromStr, + T::Err: std::error::Error + Send + Sync + 'static, +{ + field + .metadata() + .get(key) + .ok_or_else(|| anyhow!("Snowflake column '{}' has no {key} metadata", field.name()))? + .parse() + .with_context(|| { + format!( + "invalid {key} metadata for Snowflake column '{}'", + field.name() + ) + }) +} + +fn scaled_integer_to_micros( + value: i64, + scale: u32, + field_name: &str, + row: usize, +) -> AnyResult { + if scale > 9 { + bail!("Snowflake timestamp column '{field_name}' has unsupported scale {scale}"); + } + let micros = if scale <= 6 { + value.checked_mul(10_i64.pow(6 - scale)) + } else { + Some(value / 10_i64.pow(scale - 6)) + }; + micros.ok_or_else(|| timestamp_out_of_range(field_name, row)) +} + +fn timestamp_out_of_range(field_name: &str, row: usize) -> anyhow::Error { + anyhow!("Snowflake timestamp value in column '{field_name}' at row {row} is out of range") +} + +fn struct_child<'a, T: Array + 'static>( + values: &'a StructArray, + field_name: &str, + child_name: &str, +) -> AnyResult<&'a T> { + let child = values.column_by_name(child_name).ok_or_else(|| { + anyhow!("Snowflake timestamp column '{field_name}' has no '{child_name}' child") + })?; + child.as_any().downcast_ref::().ok_or_else(|| { + anyhow!( + "Snowflake timestamp column '{field_name}' has an invalid '{child_name}' child type" + ) + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::{ + array::{Int32Array, Int64Array, StringArray, StructArray}, + datatypes::{DataType, Field, Schema}, + ipc::writer::StreamWriter, + }; + use std::{collections::HashMap, sync::Arc}; + + #[test] + fn decodes_ipc_stream() { + let schema = Arc::new(Schema::new(vec![ + Field::new("ID", DataType::Int64, false), + Field::new("NAME", DataType::Utf8, true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int64Array::from(vec![1, 2])), + Arc::new(StringArray::from(vec![Some("a"), None])), + ], + ) + .unwrap(); + + let mut bytes = Vec::new(); + { + let mut writer = StreamWriter::try_new(&mut bytes, &schema).unwrap(); + writer.write(&batch).unwrap(); + writer.finish().unwrap(); + } + + let batches = decode_ipc_stream(&bytes, SnowflakeNumberMode::Decimal).unwrap(); + assert_eq!(batches.len(), 1); + assert_eq!(batches[0].num_rows(), 2); + assert_eq!(batches[0].schema().field(0).name(), "ID"); + } + + #[test] + fn normalizes_snowflake_timestamp_structs() { + let timestamp_fields = vec![ + Arc::new(Field::new("epoch", DataType::Int64, true)), + Arc::new(Field::new("fraction", DataType::Int32, true)), + Arc::new(Field::new("timezone", DataType::Int32, true)), + ]; + let timestamps = StructArray::from(vec![ + ( + timestamp_fields[0].clone(), + Arc::new(Int64Array::from(vec![Some(1), None])) as ArrayRef, + ), + ( + timestamp_fields[1].clone(), + Arc::new(Int32Array::from(vec![Some(234_567_890), None])) as ArrayRef, + ), + ( + timestamp_fields[2].clone(), + Arc::new(Int32Array::from(vec![Some(1_440), None])) as ArrayRef, + ), + ]); + let field = Field::new( + "CREATED_AT", + DataType::Struct(timestamp_fields.into()), + true, + ) + .with_metadata(HashMap::from([( + "logicalType".to_string(), + "TIMESTAMP_TZ".to_string(), + )])); + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![field])), + vec![Arc::new(timestamps)], + ) + .unwrap(); + + let normalized = normalize_snowflake_batch(batch, SnowflakeNumberMode::Decimal).unwrap(); + assert_eq!( + normalized.schema().field(0).data_type(), + &DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())) + ); + let values = normalized + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(values.value(0), 1_234_567); + assert!(values.is_null(1)); + } + + #[test] + fn normalizes_scaled_timestamp_and_time_integers() { + let timestamp = Field::new("TS", DataType::Int64, false).with_metadata(HashMap::from([ + ("logicalType".to_string(), "TIMESTAMP_NTZ".to_string()), + ("scale".to_string(), "3".to_string()), + ])); + let time = Field::new("T", DataType::Int64, false).with_metadata(HashMap::from([ + ("logicalType".to_string(), "TIME".to_string()), + ("scale".to_string(), "3".to_string()), + ])); + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![timestamp, time])), + vec![ + Arc::new(Int64Array::from(vec![1_234])), + Arc::new(Int64Array::from(vec![3_723_004])), + ], + ) + .unwrap(); + + let normalized = normalize_snowflake_batch(batch, SnowflakeNumberMode::Decimal).unwrap(); + let timestamp = normalized + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let time = normalized + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(timestamp.value(0), 1_234_000); + assert_eq!(time.value(0), 3_723_004_000_000); + } + + #[test] + fn normalizes_scaled_fixed_integers() { + let field = Field::new("AMOUNT", DataType::Int8, true).with_metadata(HashMap::from([ + ("logicalType".to_string(), "FIXED".to_string()), + ("precision".to_string(), "10".to_string()), + ("scale".to_string(), "2".to_string()), + ])); + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![field])), + vec![Arc::new(Int8Array::from(vec![Some(123), None]))], + ) + .unwrap(); + + let normalized = normalize_snowflake_batch(batch, SnowflakeNumberMode::Decimal).unwrap(); + assert_eq!( + normalized.schema().field(0).data_type(), + &DataType::Decimal128(10, 2) + ); + let values = normalized + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(values.value(0), 123); + assert!(values.is_null(1)); + } + + #[test] + fn converts_scaled_fixed_integers_to_double() { + let field = Field::new("AMOUNT", DataType::Int8, true).with_metadata(HashMap::from([ + ("logicalType".to_string(), "FIXED".to_string()), + ("precision".to_string(), "10".to_string()), + ("scale".to_string(), "2".to_string()), + ])); + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![field])), + vec![Arc::new(Int8Array::from(vec![Some(123), None]))], + ) + .unwrap(); + + let normalized = normalize_snowflake_batch(batch, SnowflakeNumberMode::Double).unwrap(); + assert_eq!(normalized.schema().field(0).data_type(), &DataType::Float64); + let values = normalized + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(values.value(0), 1.23); + assert!(values.is_null(1)); + } + + #[test] + fn converts_scaled_decimal128_to_double() { + let field = + Field::new("AMOUNT", DataType::Decimal128(38, 2), true).with_metadata(HashMap::from([ + ("logicalType".to_string(), "FIXED".to_string()), + ("precision".to_string(), "38".to_string()), + ("scale".to_string(), "2".to_string()), + ])); + let values = Decimal128Array::from(vec![Some(-123), None]) + .with_precision_and_scale(38, 2) + .unwrap(); + let batch = + RecordBatch::try_new(Arc::new(Schema::new(vec![field])), vec![Arc::new(values)]) + .unwrap(); + + let normalized = normalize_snowflake_batch(batch, SnowflakeNumberMode::Double).unwrap(); + let values = normalized + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(values.value(0), -1.23); + assert!(values.is_null(1)); + } + + #[test] + fn double_mode_leaves_scale_zero_numbers_integral() { + let field = Field::new("ID", DataType::Int8, false).with_metadata(HashMap::from([ + ("logicalType".to_string(), "FIXED".to_string()), + ("precision".to_string(), "10".to_string()), + ("scale".to_string(), "0".to_string()), + ])); + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![field])), + vec![Arc::new(Int8Array::from(vec![123]))], + ) + .unwrap(); + + let normalized = normalize_snowflake_batch(batch, SnowflakeNumberMode::Double).unwrap(); + assert_eq!(normalized.schema().field(0).data_type(), &DataType::Int8); + } +} diff --git a/crates/snowflake/src/auth.rs b/crates/snowflake/src/auth.rs new file mode 100644 index 00000000000..6e43bb0b3b2 --- /dev/null +++ b/crates/snowflake/src/auth.rs @@ -0,0 +1,139 @@ +use anyhow::{bail, Context, Result as AnyResult}; +use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine as _}; +use feldera_types::transport::snowflake::{SnowflakeAuthenticator, SnowflakeReaderConfig}; +use jsonwebtoken::{encode, Algorithm, EncodingKey, Header}; +use pkcs8::{DecodePrivateKey, EncodePrivateKey, EncodePublicKey, LineEnding}; +use rsa::{RsaPrivateKey, RsaPublicKey}; +use serde::Serialize; +use sha2::{Digest, Sha256}; +use std::{ + fs, + time::{SystemTime, UNIX_EPOCH}, +}; + +const JWT_LIFETIME_SECONDS: u64 = 60; + +#[derive(Debug, Serialize)] +struct JwtClaims { + iss: String, + sub: String, + iat: u64, + exp: u64, +} + +pub(crate) fn jwt_token(config: &SnowflakeReaderConfig) -> AnyResult { + match config.authenticator { + SnowflakeAuthenticator::SnowflakeJwt => {} + } + + let pem = fs::read(&config.private_key_file).with_context(|| { + format!( + "error reading Snowflake private key {}", + config.private_key_file + ) + })?; + let key = parse_private_key(&pem, config.private_key_file_pwd.as_deref())?; + let public_key_fingerprint = public_key_fingerprint(&key)?; + + let account = login_account_name(&config.account).to_ascii_uppercase(); + let user = config.user.to_ascii_uppercase(); + let subject = format!("{account}.{user}"); + let now = epoch_seconds(); + let claims = JwtClaims { + iss: format!("{subject}.SHA256:{public_key_fingerprint}"), + sub: subject, + iat: now, + exp: now + JWT_LIFETIME_SECONDS, + }; + let mut header = Header::new(Algorithm::RS256); + header.typ = Some("JWT".to_string()); + let signing_key = key + .to_pkcs8_pem(LineEnding::LF) + .context("error serializing Snowflake RSA private key")?; + let signing_key = EncodingKey::from_rsa_pem(signing_key.as_bytes()) + .context("error parsing Snowflake RSA private key")?; + encode(&header, &claims, &signing_key).context("error signing Snowflake JWT") +} + +pub(crate) fn login_account_name(account: &str) -> &str { + if account.to_ascii_lowercase().contains(".global") { + let account = account.split('.').next().unwrap_or(account); + account + .rsplit_once('-') + .map_or(account, |(account, _)| account) + } else { + account.split('.').next().unwrap_or(account) + } +} + +fn parse_private_key(pem: &[u8], passphrase: Option<&str>) -> AnyResult { + let pem = + std::str::from_utf8(pem).context("Snowflake RSA private key is not valid UTF-8 PEM")?; + + if let Some(passphrase) = passphrase { + if let Ok(key) = RsaPrivateKey::from_pkcs8_encrypted_pem(pem, passphrase) { + return Ok(key); + } + } + if let Ok(key) = RsaPrivateKey::from_pkcs8_pem(pem) { + return Ok(key); + } + + bail!( + "error parsing Snowflake PKCS#8 RSA private key; provide 'private_key_file_pwd' if the key is encrypted" + ) +} + +fn public_key_fingerprint(key: &RsaPrivateKey) -> AnyResult { + let public_key_der = RsaPublicKey::from(key) + .to_public_key_der() + .context("error extracting Snowflake public key DER")?; + Ok(BASE64_STANDARD.encode(Sha256::digest(public_key_der.as_bytes()))) +} + +fn epoch_seconds() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs() +} + +#[cfg(test)] +mod tests { + use super::*; + use pkcs8::{ + der::Decode, + pkcs5::pbes2::{EncryptionScheme, Parameters, Pbkdf2Params}, + PrivateKeyInfo, + }; + use rand::rngs::OsRng; + + #[test] + fn normalizes_account_name_for_login_and_jwt() { + assert_eq!(login_account_name("org-account"), "org-account"); + assert_eq!(login_account_name("xy12345.us-east-1"), "xy12345"); + assert_eq!(login_account_name("xy12345-external.global"), "xy12345"); + } + + #[test] + fn parses_encrypted_private_key() { + let key = RsaPrivateKey::new(&mut OsRng, 512).unwrap(); + let key_der = key.to_pkcs8_der().unwrap(); + let key_info = PrivateKeyInfo::from_der(key_der.as_bytes()).unwrap(); + let salt = b"test salt"; + let iv = b"test iv!"; + let encryption = Parameters { + kdf: Pbkdf2Params::hmac_with_sha256(10, salt).unwrap().into(), + encryption: EncryptionScheme::DesEde3Cbc { iv }, + }; + let encrypted = key_info + .encrypt_with_params(encryption, "secret") + .unwrap() + .to_pem("ENCRYPTED PRIVATE KEY", LineEnding::LF) + .unwrap(); + + assert!(parse_private_key(encrypted.as_bytes(), None).is_err()); + assert!(parse_private_key(encrypted.as_bytes(), Some("wrong password")).is_err()); + assert!(parse_private_key(encrypted.as_bytes(), Some("secret")).is_ok()); + } +} diff --git a/crates/snowflake/src/client.rs b/crates/snowflake/src/client.rs new file mode 100644 index 00000000000..851c0786118 --- /dev/null +++ b/crates/snowflake/src/client.rs @@ -0,0 +1,327 @@ +use crate::auth::{jwt_token, login_account_name}; +use anyhow::{anyhow, bail, Context, Result as AnyResult}; +use feldera_types::transport::snowflake::SnowflakeReaderConfig; +use reqwest::{ + header::{HeaderMap, HeaderValue, ACCEPT, AUTHORIZATION, USER_AGENT}, + Client, Url, +}; +use serde::Deserialize; +use serde_json::{json, Value as JsonValue}; +use std::{env, sync::atomic::AtomicU64, time::Duration}; +use uuid::Uuid; + +// Match Snowflake's C/C++ connector, which defines its client id as "C API": +// https://github.com/snowflakedb/libsnowflakeclient/blob/master/include/snowflake/client.h +const CLIENT_APP_ID: &str = "C API"; +const CLIENT_APP_VERSION: &str = env!("CARGO_PKG_VERSION"); +const QUERY_RESULT_FORMAT: &str = "ARROW_FORCE"; + +pub(crate) struct SnowflakeClient { + pub(super) http: Client, + pub(super) base_url: Url, + session_token: String, + pub(super) sequence_id: AtomicU64, +} + +#[derive(Deserialize)] +pub(super) struct SnowflakeApiResponse { + #[serde(default)] + success: Option, + #[serde(default)] + pub(super) code: Option, + #[serde(default)] + message: Option, + pub(super) data: Option, +} + +#[derive(Deserialize)] +struct LoginData { + #[serde(default)] + token: Option, + #[serde(default, rename = "sessionToken")] + session_token: Option, +} + +impl SnowflakeClient { + pub(crate) async fn from_reader_config(config: &SnowflakeReaderConfig) -> AnyResult { + config.validate().map_err(|error| anyhow!(error))?; + let http = Client::builder() + .user_agent(user_agent()) + .connect_timeout(Duration::from_secs(30)) + .no_gzip() + .build() + .context("error building Snowflake HTTP client")?; + let base_url = account_base_url(&config.account)?; + let jwt = jwt_token(config)?; + + let login_url = login_url(&base_url, config)?; + let body = login_body(config, &jwt); + let response = http + .post(login_url) + .header(ACCEPT, "application/snowflake") + .json(&body) + .send() + .await + .context("error sending Snowflake login request")? + .error_for_status() + .context("Snowflake login request failed")? + .json::>() + .await + .context("error parsing Snowflake login response")?; + + let data = response_data(response, "Snowflake login")?; + + let client = Self { + http, + base_url, + session_token: data.token.or(data.session_token).ok_or_else(|| { + anyhow!("Snowflake login response did not include a session token") + })?, + sequence_id: AtomicU64::new(0), + }; + + // Snowflake's C/C++ connector enables Arrow with this session parameter: + // https://github.com/snowflakedb/libsnowflakeclient/blob/master/cpp/lib/ArrowChunkIterator.hpp + client + .execute_statement(&format!( + "alter session set C_API_QUERY_RESULT_FORMAT={QUERY_RESULT_FORMAT}" + )) + .await + .context("error enabling Snowflake Arrow result format")?; + + Ok(client) + } + + pub(super) fn endpoint(&self, path: &str) -> AnyResult { + self.base_url + .join(path.trim_start_matches('/')) + .with_context(|| format!("invalid Snowflake endpoint path: {path}")) + } + + pub(super) fn session_headers(&self) -> AnyResult { + let mut headers = HeaderMap::new(); + headers.insert(ACCEPT, HeaderValue::from_static("application/snowflake")); + headers.insert(USER_AGENT, HeaderValue::from_str(&user_agent())?); + headers.insert( + AUTHORIZATION, + HeaderValue::from_str(&format!("Snowflake Token=\"{}\"", self.session_token)) + .context("invalid Snowflake session token header")?, + ); + Ok(headers) + } +} + +fn account_base_url(account: &str) -> AnyResult { + let account = account.trim(); + validate_account_identifier(account)?; + let top_level_domain = if account + .split('.') + .nth(1) + .is_some_and(|region| region.to_ascii_lowercase().starts_with("cn-")) + { + "cn" + } else { + "com" + }; + let url = format!("https://{account}.snowflakecomputing.{top_level_domain}"); + Url::parse(&url) + .with_context(|| format!("invalid Snowflake account URL derived from '{account}'")) +} + +fn validate_account_identifier(account: &str) -> AnyResult<()> { + if account.is_empty() + || account + .to_ascii_lowercase() + .contains(".snowflakecomputing.") + || account.split('.').any(|part| { + part.is_empty() + || !part + .bytes() + .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-')) + }) + { + bail!( + "invalid Snowflake account identifier '{account}'; expected dot-separated letters, digits, underscores, or hyphens" + ); + } + Ok(()) +} + +fn login_url(base_url: &Url, config: &SnowflakeReaderConfig) -> AnyResult { + let mut url = base_url.join("session/v1/login-request")?; + { + let mut query = url.query_pairs_mut(); + query.append_pair("request_id", &Uuid::new_v4().to_string()); + if let Some(database) = nonempty(config.database.as_deref()) { + query.append_pair("databaseName", database); + } + if let Some(schema) = nonempty(config.schema.as_deref()) { + query.append_pair("schemaName", schema); + } + if let Some(warehouse) = nonempty(config.warehouse.as_deref()) { + query.append_pair("warehouse", warehouse); + } + if let Some(role) = nonempty(config.role.as_deref()) { + query.append_pair("roleName", role); + } + } + Ok(url) +} + +fn nonempty(value: Option<&str>) -> Option<&str> { + value.and_then(|value| { + let value = value.trim(); + (!value.is_empty()).then_some(value) + }) +} + +fn login_body(config: &SnowflakeReaderConfig, jwt: &str) -> JsonValue { + let mut session_parameters = serde_json::Map::new(); + session_parameters.insert( + "AUTOCOMMIT".to_string(), + JsonValue::String("true".to_string()), + ); + session_parameters.insert("TIMEZONE".to_string(), JsonValue::String("UTC".to_string())); + + json!({ + "data": { + "CLIENT_APP_ID": CLIENT_APP_ID, + "CLIENT_APP_VERSION": CLIENT_APP_VERSION, + "ACCOUNT_NAME": login_account_name(&config.account), + "LOGIN_NAME": config.user, + "AUTHENTICATOR": "SNOWFLAKE_JWT", + "TOKEN": jwt, + "CLIENT_ENVIRONMENT": { + "APPLICATION": CLIENT_APP_ID, + "OS": env::consts::OS, + "ARCH": env::consts::ARCH, + }, + "SESSION_PARAMETERS": session_parameters, + } + }) +} + +pub(super) fn response_data(response: SnowflakeApiResponse, operation: &str) -> AnyResult { + if response.success == Some(false) { + bail!( + "{operation} failed: code={} message={}", + response.code.unwrap_or_default(), + response.message.unwrap_or_default() + ); + } + + response.data.ok_or_else(|| { + anyhow!( + "{operation} response did not include a data object: code={} message={}", + response.code.unwrap_or_default(), + response.message.unwrap_or_default() + ) + }) +} + +fn user_agent() -> String { + format!("{CLIENT_APP_ID}/{CLIENT_APP_VERSION}") +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::{ + array::{Array, Decimal128Array, Float64Array}, + datatypes::DataType, + }; + use feldera_types::transport::snowflake::{SnowflakeAuthenticator, SnowflakeNumberMode}; + use futures_util::StreamExt; + + #[test] + fn derives_account_urls() { + assert_eq!( + account_base_url("org-account").unwrap().as_str(), + "https://org-account.snowflakecomputing.com/" + ); + assert_eq!( + account_base_url("xy12345.us-east-1").unwrap().as_str(), + "https://xy12345.us-east-1.snowflakecomputing.com/" + ); + assert!(account_base_url("http://localhost:8080").is_err()); + assert!(account_base_url("xy12345.snowflakecomputing.com").is_err()); + } + + #[ignore] + #[tokio::test] + async fn snowflake_arrow_smoke_test() -> AnyResult<()> { + let config = SnowflakeReaderConfig { + account: env::var("SNOWFLAKE_ACCOUNT").context("SNOWFLAKE_ACCOUNT is not set")?, + user: env::var("SNOWFLAKE_USER").context("SNOWFLAKE_USER is not set")?, + authenticator: SnowflakeAuthenticator::SnowflakeJwt, + role: env::var("SNOWFLAKE_ROLE").ok(), + warehouse: env::var("SNOWFLAKE_WAREHOUSE").ok(), + database: env::var("SNOWFLAKE_DATABASE").ok(), + schema: env::var("SNOWFLAKE_SCHEMA").ok(), + private_key_file: env::var("SNOWFLAKE_PRIVATE_KEY_FILE") + .context("SNOWFLAKE_PRIVATE_KEY_FILE is not set")?, + private_key_file_pwd: env::var("SNOWFLAKE_PRIVATE_KEY_FILE_PWD").ok(), + table: "unused".to_string(), + column_mapping: Default::default(), + number_mode: Default::default(), + mode: Default::default(), + transaction_mode: Default::default(), + snapshot_filter: None, + num_parsers: 4, + max_concurrent_readers: None, + }; + let client = SnowflakeClient::from_reader_config(&config).await?; + let custom_query = env::var("SNOWFLAKE_TEST_QUERY").ok(); + let using_default_query = custom_query.is_none(); + let query = custom_query.unwrap_or_else(|| { + "select -1.23::number(10,2) as AMOUNT, 12::number(10,0) as ID, null::number(10,2) as OPTIONAL_AMOUNT" + .to_string() + }); + for (number_mode, expected_type) in [ + (SnowflakeNumberMode::Decimal, DataType::Decimal128(10, 2)), + (SnowflakeNumberMode::Double, DataType::Float64), + ] { + let result = client + .query_arrow_batch_stream(&query, 1, number_mode) + .await?; + let batches = result.batches.collect::>().await; + let batches = batches.into_iter().collect::>>()?; + assert_eq!( + batches.iter().map(|batch| batch.num_rows()).sum::(), + 1 + ); + if using_default_query { + assert_eq!(batches[0].schema().field(0).data_type(), &expected_type); + assert!(matches!( + batches[0].schema().field(1).data_type(), + DataType::Int8 + | DataType::Int16 + | DataType::Int32 + | DataType::Int64 + | DataType::Decimal128(_, 0) + )); + match number_mode { + SnowflakeNumberMode::Decimal => { + let amounts = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(amounts.value(0), -123); + assert!(batches[0].column(2).is_null(0)); + } + SnowflakeNumberMode::Double => { + let amounts = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(amounts.value(0), -1.23); + assert!(batches[0].column(2).is_null(0)); + } + } + } + } + Ok(()) + } +} diff --git a/crates/snowflake/src/input.rs b/crates/snowflake/src/input.rs new file mode 100644 index 00000000000..ee615178a1f --- /dev/null +++ b/crates/snowflake/src/input.rs @@ -0,0 +1,468 @@ +use crate::{client::SnowflakeClient, query::build_snapshot_query, snowflake_input_serde_config}; +use anyhow::{Context, Result as AnyResult, anyhow}; +use arrow::record_batch::RecordBatch; +use chrono::Utc; +use dbsp::circuit::tokio::TOKIO; +use feldera_adapterlib::{ + PipelineState, + catalog::{ArrowStream, InputCollectionHandle}, + errors::journal::ControllerError, + format::{InputBuffer, ParseError}, + transport::{ + InputConsumer, InputEndpoint, InputQueue, InputQueueEntry, InputReader, InputReaderCommand, + IntegratedInputEndpoint, Resume, Watermark, parse_resume_info, + }, +}; +use feldera_types::{ + adapter_stats::ConnectorHealth, + config::FtModel, + program_schema::Relation, + transport::snowflake::{SnowflakeIngestMode, SnowflakeReaderConfig, SnowflakeTransactionMode}, +}; +use futures_util::StreamExt; +use log::{debug, info}; +use serde::{Deserialize, Serialize}; +use std::{ + sync::{Arc, Mutex}, + thread, +}; +use tokio::{ + select, + sync::{ + mpsc, + watch::{Receiver, Sender, channel}, + }, +}; + +/// Integrated input connector that reads from a Snowflake table. +pub struct SnowflakeInputEndpoint { + inner: Arc, +} + +impl SnowflakeInputEndpoint { + pub fn new( + endpoint_name: &str, + config: &SnowflakeReaderConfig, + consumer: Box, + ) -> Self { + Self { + inner: Arc::new(SnowflakeInputEndpointInner::new( + endpoint_name, + config.clone(), + consumer, + )), + } + } +} + +impl InputEndpoint for SnowflakeInputEndpoint { + fn fault_tolerance(&self) -> Option { + // This test-only fault tolerance supports full-refresh checkpoints, not recovery + // within a snapshot: before EOI, recovery reruns the entire snapshot. + Some(FtModel::AtLeastOnce) + } +} + +impl IntegratedInputEndpoint for SnowflakeInputEndpoint { + fn open( + self: Box, + input_handle: &InputCollectionHandle, + seek: Option, + ) -> AnyResult> { + Ok(Box::new(SnowflakeInputReader::new( + &self.inner, + input_handle, + seek, + )?)) + } +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +struct SnowflakeResumeInfo { + eoi: bool, +} + +impl SnowflakeResumeInfo { + fn initial() -> Self { + Self { eoi: false } + } + + fn eoi() -> Self { + Self { eoi: true } + } + + fn to_resume(&self) -> Resume { + Resume::Seek { + seek: serde_json::to_value(self).unwrap(), + } + } +} + +struct SnowflakeInputReader { + sender: Sender, + inner: Arc, +} + +impl SnowflakeInputReader { + fn new( + endpoint: &Arc, + input_handle: &InputCollectionHandle, + resume_info: Option, + ) -> AnyResult { + endpoint.config.validate().map_err(|e| { + ControllerError::invalid_transport_configuration(&endpoint.endpoint_name, &e) + })?; + + let resume_info = resume_info + .as_ref() + .map(parse_resume_info::) + .transpose()?; + let eoi = resume_info + .as_ref() + .is_some_and(|resume_info| resume_info.eoi); + + if let Some(resume_info) = &resume_info { + *endpoint.last_resume_status.lock().unwrap() = Some(resume_info.clone()); + endpoint + .queue + .push_with_aux((None, Vec::new()), Utc::now(), Some(resume_info.clone())); + } + + let (sender, receiver) = channel(PipelineState::Paused); + + let input_stream = input_handle + .handle + .configure_arrow_deserializer(snowflake_input_serde_config())?; + let schema = input_handle.schema.clone(); + let endpoint_name = endpoint.endpoint_name.clone(); + + if eoi { + info!( + "snowflake {endpoint_name}: skipping connector initialization because the snapshot is already complete" + ); + endpoint.consumer.eoi(); + } else { + let endpoint_clone = endpoint.clone(); + let receiver_clone = receiver.clone(); + let (init_status_sender, mut init_status_receiver) = + mpsc::channel::>(1); + + thread::Builder::new() + .name(format!("{endpoint_name}-snowflake-input-tokio-wrapper")) + .spawn(move || { + TOKIO.block_on(async { + endpoint_clone + .worker_task(input_stream, schema, receiver_clone, init_status_sender) + .await; + }) + }) + .expect("failed to spawn snowflake-input tokio wrapper thread"); + + init_status_receiver.blocking_recv().ok_or_else(|| { + anyhow!("worker thread terminated unexpectedly during initialization") + })??; + } + + Ok(Self { + sender, + inner: endpoint.clone(), + }) + } +} + +impl InputReader for SnowflakeInputReader { + fn as_any(self: Arc) -> Arc { + self + } + + fn request(&self, command: InputReaderCommand) { + match command { + InputReaderCommand::Replay { .. } => panic!( + "replay command is not supported by SnowflakeInputReader; this is a bug, please report it to developers" + ), + InputReaderCommand::Extend => { + let _ = self.sender.send_replace(PipelineState::Running); + } + InputReaderCommand::Pause => { + let _ = self.sender.send_replace(PipelineState::Paused); + } + InputReaderCommand::Queue { + checkpoint_requested, + } => { + let stop_at: &dyn Fn(&Option) -> bool = if checkpoint_requested + { + &Option::is_some + } else { + &|_| false + }; + let (total, _, resume_info) = self.inner.queue.flush_with_aux_until(stop_at); + let resume_status = match resume_info.last() { + None => self.inner.last_resume_status.lock().unwrap().clone(), + Some((_, resume_info)) => resume_info.clone(), + }; + *self.inner.last_resume_status.lock().unwrap() = resume_status.clone(); + + let resume = match resume_status { + Some(resume_info) => resume_info.to_resume(), + None => Resume::Barrier, + }; + + self.inner.consumer.extended( + total, + Some(resume), + resume_info + .into_iter() + .map(|(timestamp, resume_info)| { + Watermark::new( + timestamp, + resume_info + .map(|resume_info| serde_json::to_value(resume_info).unwrap()), + ) + }) + .collect(), + ); + } + InputReaderCommand::Disconnect => { + let _ = self.sender.send_replace(PipelineState::Terminated); + } + } + } + + fn is_closed(&self) -> bool { + self.inner.queue.is_empty() && self.sender.is_closed() + } +} + +impl Drop for SnowflakeInputReader { + fn drop(&mut self) { + self.disconnect(); + } +} + +struct SnowflakeInputEndpointInner { + endpoint_name: String, + config: SnowflakeReaderConfig, + consumer: Box, + queue: InputQueue>, + last_resume_status: Mutex>, +} + +impl SnowflakeInputEndpointInner { + fn new( + endpoint_name: &str, + config: SnowflakeReaderConfig, + consumer: Box, + ) -> Self { + let queue = InputQueue::new(consumer.clone()); + Self { + endpoint_name: endpoint_name.to_string(), + config, + consumer, + queue, + last_resume_status: Mutex::new(Some(SnowflakeResumeInfo::initial())), + } + } + + async fn worker_task( + self: Arc, + input_stream: Box, + schema: Relation, + receiver: Receiver, + init_status_sender: mpsc::Sender>, + ) { + let mut receiver_clone = receiver.clone(); + select! { + _ = Self::worker_task_inner(self.clone(), input_stream, schema, receiver, init_status_sender) => { + debug!("snowflake {}: worker task terminated", &self.endpoint_name); + } + _ = receiver_clone.wait_for(|state| state == &PipelineState::Terminated) => { + debug!("snowflake {}: received termination command; worker task canceled", &self.endpoint_name); + } + } + } + + async fn worker_task_inner( + self: Arc, + input_stream: Box, + schema: Relation, + mut receiver: Receiver, + init_status_sender: mpsc::Sender>, + ) { + let query = match self.config.mode { + SnowflakeIngestMode::Snapshot => build_snapshot_query( + &self.config.table, + self.config.snapshot_filter.as_deref(), + &self.config.column_mapping, + &schema, + ), + }; + let query = match query { + Ok(query) => query, + Err(e) => { + let _ = init_status_sender + .send(Err(ControllerError::invalid_transport_configuration( + &self.endpoint_name, + &e.to_string(), + ))) + .await; + return; + } + }; + + let _ = init_status_sender.send(Ok(())).await; + wait_running(&mut receiver).await; + + debug!( + "snowflake {}: prepared snapshot query for table '{}'", + &self.endpoint_name, &self.config.table + ); + + let client = match SnowflakeClient::from_reader_config(&self.config).await { + Ok(client) => client, + Err(e) => { + let message = format!("error connecting to Snowflake for snapshot input: {e:#}"); + self.consumer + .update_connector_health(ConnectorHealth::unhealthy(&message)); + self.consumer.error(true, anyhow!(message), None); + return; + } + }; + + let result_stream = match client + .query_arrow_batch_stream( + &query, + self.config.max_concurrent_readers(), + self.config.number_mode, + ) + .await + { + Ok(result) => result, + Err(e) => { + let message = format!("error executing Snowflake snapshot query: {e:#}"); + self.consumer + .update_connector_health(ConnectorHealth::unhealthy(&message)); + self.consumer.error(true, anyhow!(message), None); + return; + } + }; + let query_id = result_stream + .metadata + .query_id + .clone() + .unwrap_or_else(|| "".to_string()); + let total_rows = result_stream.metadata.total_rows; + self.consumer + .update_connector_health(ConnectorHealth::healthy()); + + info!( + "snowflake {}: reading snapshot with Snowflake query '{}' (rows: {})", + &self.endpoint_name, + query_id, + total_rows + .map(|rows| rows.to_string()) + .unwrap_or_else(|| "".to_string()) + ); + + let mut timestamp = Utc::now(); + let transaction_label = match self.config.transaction_mode { + SnowflakeTransactionMode::None => None, + SnowflakeTransactionMode::Snapshot => { + Some(Some(format!("snowflake-snapshot:{query_id}"))) + } + }; + let batches = result_stream.batches; + let parser_receiver = receiver.clone(); + let parser_query_id = query_id.clone(); + let mut parsed_batches = batches + .map(move |batch_result| { + let batch_stream = input_stream.fork(); + let batch_query_id = parser_query_id.clone(); + let mut receiver = parser_receiver.clone(); + async move { + let batch = batch_result?; + wait_running(&mut receiver).await; + tokio::task::spawn_blocking(move || { + parse_arrow_batch(batch_stream, batch, &batch_query_id) + }) + .await + .context("Snowflake Arrow batch parser task failed") + } + }) + .buffer_unordered(self.config.num_parsers as usize); + + while let Some(parsed_batch) = parsed_batches.next().await { + wait_running(&mut receiver).await; + let (buffer, errors) = match parsed_batch { + Ok(parsed) => parsed, + Err(error) => { + self.commit_transaction_on_error(transaction_label.is_some()); + let message = format!("error processing Snowflake snapshot query: {error:#}"); + self.consumer + .update_connector_health(ConnectorHealth::unhealthy(&message)); + self.consumer.error(true, anyhow!(message), None); + return; + } + }; + + let entry = InputQueueEntry::new_with_aux(timestamp, None) + .with_buffer(buffer) + .with_start_transaction(transaction_label.clone()); + + self.queue.push_entry(entry, errors); + timestamp = Utc::now(); + } + + let mut completion_entry = + InputQueueEntry::new_with_aux(timestamp, Some(SnowflakeResumeInfo::eoi())); + if transaction_label.is_some() { + completion_entry = completion_entry + .with_start_transaction(transaction_label) + .with_commit_transaction(true); + } + self.queue.push_entry(completion_entry, Vec::new()); + + info!( + "snowflake {}: snapshot load completed (query: '{}')", + &self.endpoint_name, query_id + ); + + self.consumer.eoi(); + } + + fn commit_transaction_on_error(&self, transaction_enabled: bool) { + if transaction_enabled { + // Feldera does not currently expose transaction rollback to connectors. Commit any + // already queued records so a failed connector cannot leave the pipeline blocked. + self.queue.push_entry( + InputQueueEntry::new_with_aux(Utc::now(), None).with_commit_transaction(true), + Vec::new(), + ); + } + } +} + +fn parse_arrow_batch( + mut input_stream: Box, + batch: RecordBatch, + query_id: &str, +) -> (Option>, Vec) { + let errors = input_stream.insert(&batch, &None).map_or_else( + |error| { + vec![ParseError::bin_envelope_error( + format!( + "error deserializing Snowflake Arrow batch from query '{query_id}': {error}" + ), + &[], + None, + )] + }, + |()| Vec::new(), + ); + (input_stream.take_all(), errors) +} + +async fn wait_running(receiver: &mut Receiver) { + let _ = receiver + .wait_for(|state| state == &PipelineState::Running) + .await; +} diff --git a/crates/snowflake/src/lib.rs b/crates/snowflake/src/lib.rs new file mode 100644 index 00000000000..c61a79002a3 --- /dev/null +++ b/crates/snowflake/src/lib.rs @@ -0,0 +1,24 @@ +mod arrow; +mod auth; +mod client; +mod input; +mod query; + +pub use input::SnowflakeInputEndpoint; + +use feldera_types::serde_with_context::{ + serde_config::{BinaryFormat, DecimalFormat, VariantFormat}, + DateFormat, SqlSerdeConfig, TimeFormat, TimestampFormat, +}; + +pub fn snowflake_input_serde_config() -> SqlSerdeConfig { + let mut config = SqlSerdeConfig::default() + .with_date_format(DateFormat::String("%Y-%m-%d")) + .with_time_format(TimeFormat::NanosSigned) + .with_timestamp_format(TimestampFormat::MicrosSinceEpoch) + .with_decimal_format(DecimalFormat::String) + .with_variant_format(VariantFormat::JsonString) + .with_binary_format(BinaryFormat::Array); + config.timestamp_tz_format = TimestampFormat::MicrosSinceEpoch; + config +} diff --git a/crates/snowflake/src/query.rs b/crates/snowflake/src/query.rs new file mode 100644 index 00000000000..aa9a148861c --- /dev/null +++ b/crates/snowflake/src/query.rs @@ -0,0 +1,756 @@ +use crate::{ + arrow::{decode_base64_ipc_stream, normalize_snowflake_batch}, + client::{response_data, SnowflakeApiResponse, SnowflakeClient}, +}; +use anyhow::{anyhow, bail, Context, Result as AnyResult}; +use arrow::{buffer::Buffer, ipc::reader::StreamDecoder, record_batch::RecordBatch}; +use async_compression::tokio::bufread::GzipDecoder; +use async_stream::try_stream; +use bytes::BytesMut; +use feldera_types::{program_schema::Relation, transport::snowflake::SnowflakeNumberMode}; +use futures_util::{stream, stream::BoxStream, StreamExt, TryStreamExt}; +use rand::Rng; +use reqwest::{ + header::{HeaderMap, HeaderName, HeaderValue, CONTENT_ENCODING}, + StatusCode, Url, +}; +use serde::Deserialize; +use serde_json::json; +use std::{ + collections::{BTreeMap, HashMap}, + io, + pin::Pin, + sync::atomic::Ordering, + time::{Duration, SystemTime, UNIX_EPOCH}, +}; +use tokio::io::{AsyncBufReadExt, AsyncRead, AsyncReadExt, BufReader}; +use tokio_util::io::StreamReader as AsyncStreamReader; +use uuid::Uuid; + +const QUERY_IN_PROGRESS_CODE: &str = "333333"; +const QUERY_IN_PROGRESS_ASYNC_CODE: &str = "333334"; +const CHUNK_DOWNLOAD_MAX_ATTEMPTS: u32 = 6; +const CHUNK_RETRY_INITIAL_BACKOFF_MS: u64 = 250; +const CHUNK_RETRY_MAX_BACKOFF_MS: u64 = 8_000; +const ARROW_DECODE_BUFFER_SIZE: usize = 2 * 1024 * 1024; + +pub(crate) struct SnowflakeArrowQueryMetadata { + pub(crate) query_id: Option, + pub(crate) total_rows: Option, +} + +pub(crate) struct SnowflakeArrowBatchStream<'a> { + pub(crate) metadata: SnowflakeArrowQueryMetadata, + pub(crate) batches: BoxStream<'a, AnyResult>, +} + +fn snowflake_column_identifier(identifier: &str) -> AnyResult<&str> { + let identifier = identifier.trim(); + + if identifier.starts_with('"') { + if identifier.len() < 2 || !identifier.ends_with('"') { + bail!("invalid quoted Snowflake column identifier '{identifier}'"); + } + + let mut characters = identifier[1..identifier.len() - 1].chars(); + while let Some(character) = characters.next() { + if character == '"' && characters.next() != Some('"') { + bail!("invalid quoted Snowflake column identifier '{identifier}'"); + } + } + + return Ok(identifier); + } + + let mut characters = identifier.chars(); + if !matches!(characters.next(), Some(character) if character.is_ascii_alphabetic() || character == '_') + || !characters + .all(|character| character.is_ascii_alphanumeric() || matches!(character, '_' | '$')) + { + bail!("invalid Snowflake column identifier '{identifier}'"); + } + + Ok(identifier) +} + +pub(crate) fn build_snapshot_query( + table: &str, + snapshot_filter: Option<&str>, + column_mapping: &BTreeMap, + relation: &Relation, +) -> AnyResult { + let table = table.trim(); + if table.is_empty() { + bail!("Snowflake table name must not be empty"); + } + + let skip_unused_columns = relation.get_property("skip_unused_columns") == Some("true"); + let mut source_columns = HashMap::new(); + for (target, source) in column_mapping { + let target_lowercase = target.to_lowercase(); + let field = relation + .fields + .iter() + .find(|field| field.name.case_sensitive && field.name.name() == *target) + .or_else(|| { + relation.fields.iter().find(|field| { + !field.name.case_sensitive && field.name.name() == target_lowercase + }) + }); + let Some(field) = field else { + bail!( + "Snowflake column mapping refers to unknown Feldera column '{target}' in input relation '{}'", + relation.name + ); + }; + + let source = snowflake_column_identifier(source)?; + let target = field.name.name(); + if source_columns.insert(target.clone(), source).is_some() { + bail!( + "Snowflake column mapping contains multiple entries for Feldera column '{target}'" + ); + } + } + + let columns = relation + .fields + .iter() + .filter(|field| { + !skip_unused_columns + || !field.unused + || (!field.columntype.nullable && field.default.is_none()) + }) + .map(|field| { + source_columns.get(&field.name.name()).map_or_else( + || field.name.sql_name(), + |source| format!("{source} AS {}", field.name.sql_name()), + ) + }) + .collect::>(); + + if columns.is_empty() { + bail!( + "Snowflake snapshot query for input relation '{}' has no columns to read", + relation.name + ); + } + + let mut query = format!("SELECT {} FROM {table}", columns.join(", ")); + + if let Some(filter) = snapshot_filter { + let filter = filter.trim(); + if filter.is_empty() { + bail!("Snowflake snapshot filter must not be empty when specified"); + } + query.push_str(" WHERE "); + query.push_str(filter); + } + + Ok(query) +} + +#[derive(Deserialize)] +pub(super) struct QueryData { + #[serde(default, rename = "queryId")] + query_id: Option, + #[serde(default, rename = "queryResultFormat")] + query_result_format: Option, + #[serde(default, rename = "rowsetBase64")] + rowset_base64: Option, + #[serde(default)] + chunks: Vec, + #[serde(default, rename = "chunkHeaders")] + chunk_headers: HashMap, + #[serde(default)] + qrmk: Option, + #[serde(default)] + total: Option, + #[serde(default, rename = "getResultUrl")] + get_result_url: Option, +} + +#[derive(Deserialize)] +struct SnowflakeChunk { + url: String, +} + +impl SnowflakeClient { + pub(crate) async fn query_arrow_batch_stream( + &self, + sql: &str, + max_concurrent_readers: usize, + number_mode: SnowflakeNumberMode, + ) -> AnyResult> { + let data = self.execute_statement(sql).await?; + let format = data.query_result_format.as_deref().unwrap_or_default(); + if !format.eq_ignore_ascii_case("arrow") && !format.eq_ignore_ascii_case("arrow_force") { + bail!("Snowflake returned queryResultFormat='{format}', expected Arrow"); + } + + let inline_batches = data + .rowset_base64 + .as_deref() + .map(|rowset| decode_base64_ipc_stream(rowset, number_mode)) + .transpose()?; + let chunk_headers = chunk_headers(&data)?; + let metadata = SnowflakeArrowQueryMetadata { + query_id: data.query_id, + total_rows: data.total, + }; + + let inline_stream = stream::iter( + inline_batches + .into_iter() + .flatten() + .map(Ok::<_, anyhow::Error>), + ); + let downloaded_batches = stream::iter( + data.chunks + .into_iter() + .map(move |chunk| self.stream_chunk(chunk, chunk_headers.clone(), number_mode)), + ) + .flatten_unordered(max_concurrent_readers.max(1)); + + Ok(SnowflakeArrowBatchStream { + metadata, + batches: inline_stream.chain(downloaded_batches).boxed(), + }) + } + + pub(super) async fn execute_statement(&self, sql: &str) -> AnyResult { + let sql = sql.trim(); + if sql.is_empty() { + bail!("Snowflake query must not be empty"); + } + + let mut url = self.endpoint("/queries/v1/query-request")?; + url.query_pairs_mut() + .append_pair("requestId", &Uuid::new_v4().to_string()); + + let body = json!({ + "sqlText": sql, + "asyncExec": false, + "sequenceId": self.sequence_id.fetch_add(1, Ordering::Relaxed) + 1, + "querySubmissionTime": epoch_millis(), + }); + + let mut response = self + .http + .post(url) + .headers(self.session_headers()?) + .json(&body) + .send() + .await + .context("error sending Snowflake query request")? + .error_for_status() + .context("Snowflake query request failed")? + .json::>() + .await + .context("error parsing Snowflake query response")?; + + while response.code.as_deref().is_some_and(|code| { + code == QUERY_IN_PROGRESS_CODE || code == QUERY_IN_PROGRESS_ASYNC_CODE + }) { + let result_url = response + .data + .as_ref() + .and_then(|data| data.get_result_url.as_deref()) + .ok_or_else(|| { + anyhow!("Snowflake query-in-progress response did not include getResultUrl") + })?; + let result_url = self + .base_url + .join(result_url.trim_start_matches('/')) + .with_context(|| format!("invalid Snowflake query result URL '{result_url}'"))?; + response = self + .http + .get(result_url) + .headers(self.session_headers()?) + .send() + .await + .context("error polling Snowflake query result")? + .error_for_status() + .context("Snowflake query result poll failed")? + .json::>() + .await + .context("error parsing Snowflake query result response")?; + } + + response_data(response, "Snowflake query") + } + + fn stream_chunk<'a>( + &'a self, + chunk: SnowflakeChunk, + headers: HeaderMap, + number_mode: SnowflakeNumberMode, + ) -> BoxStream<'a, AnyResult> { + try_stream! { + let chunk_url = redact_url(&chunk.url); + let response = self + .open_chunk_response(&chunk.url, &headers) + .await + .with_context(|| format!("error opening Snowflake result chunk {chunk_url}"))?; + let gzip_header = response_is_gzip_encoded(&response, &chunk.url); + let body_url = chunk_url.clone(); + let body = response.bytes_stream().map_err(move |error| { + io::Error::other(format!( + "error reading Snowflake result chunk {body_url}: {error}" + )) + }); + let mut body = BufReader::new(AsyncStreamReader::new(body)); + let gzip_magic = body + .fill_buf() + .await + .with_context(|| format!("error reading Snowflake result chunk {chunk_url}"))? + .starts_with(&[0x1f, 0x8b]); + + let reader: Pin> = if gzip_header || gzip_magic { + Box::pin(GzipDecoder::new(body)) + } else { + Box::pin(body) + }; + let mut batches = decode_arrow_reader(reader, chunk_url, number_mode); + while let Some(batch) = batches.next().await { + yield batch?; + } + } + .boxed() + } + + async fn open_chunk_response( + &self, + url: &str, + headers: &HeaderMap, + ) -> AnyResult { + for attempt in 1..=CHUNK_DOWNLOAD_MAX_ATTEMPTS { + let response = match self.http.get(url).headers(headers.clone()).send().await { + Ok(response) => response, + Err(_) if attempt < CHUNK_DOWNLOAD_MAX_ATTEMPTS => { + tokio::time::sleep(chunk_retry_backoff(attempt)).await; + continue; + } + Err(error) => { + return Err(error).with_context(|| { + format!("error sending Snowflake chunk request, attempt {attempt}") + }); + } + }; + + let status = response.status(); + if status.is_success() { + return Ok(response); + } + + if !is_retryable_chunk_status(status) || attempt == CHUNK_DOWNLOAD_MAX_ATTEMPTS { + response.error_for_status().with_context(|| { + format!("Snowflake chunk request failed with status {status}") + })?; + unreachable!("error_for_status returned Ok for unsuccessful status {status}"); + } + + tokio::time::sleep(chunk_retry_backoff(attempt)).await; + } + + unreachable!("chunk retry loop should return or error"); + } +} + +/// Incrementally decode Arrow IPC batches from an asynchronous byte source. +fn decode_arrow_reader( + mut reader: R, + source: String, + number_mode: SnowflakeNumberMode, +) -> BoxStream<'static, AnyResult> +where + R: AsyncRead + Send + Unpin + 'static, +{ + try_stream! { + let mut decoder = StreamDecoder::new(); + + loop { + let mut bytes = BytesMut::with_capacity(ARROW_DECODE_BUFFER_SIZE); + let count = reader + .read_buf(&mut bytes) + .await + .with_context(|| format!("error reading Arrow IPC data from {source}"))?; + if count == 0 { + break; + } + + let mut buffer = Buffer::from(bytes.freeze()); + while !buffer.is_empty() { + if let Some(batch) = decoder + .decode(&mut buffer) + .with_context(|| format!("error decoding Arrow IPC data from {source}"))? + { + yield normalize_snowflake_batch(batch, number_mode).with_context(|| { + format!("error normalizing Arrow data from {source}") + })?; + } + } + } + + decoder + .finish() + .with_context(|| format!("incomplete Arrow IPC stream from {source}"))?; + } + .boxed() +} + +fn chunk_retry_backoff(attempt: u32) -> Duration { + let base_delay_ms = CHUNK_RETRY_INITIAL_BACKOFF_MS + .checked_shl(attempt.saturating_sub(1)) + .unwrap_or(u64::MAX) + .min(CHUNK_RETRY_MAX_BACKOFF_MS); + let jitter_ms = rand::thread_rng().gen_range(0..=(base_delay_ms / 4)); + + Duration::from_millis((base_delay_ms + jitter_ms).min(CHUNK_RETRY_MAX_BACKOFF_MS)) +} + +fn chunk_headers(data: &QueryData) -> AnyResult { + let mut headers = HeaderMap::new(); + for (name, value) in &data.chunk_headers { + let name = HeaderName::from_bytes(name.as_bytes()) + .with_context(|| format!("invalid Snowflake chunk header name: {name}"))?; + let value = HeaderValue::from_str(value) + .with_context(|| format!("invalid Snowflake chunk header value for {name}"))?; + headers.insert(name, value); + } + + if headers.is_empty() { + if let Some(qrmk) = &data.qrmk { + headers.insert( + HeaderName::from_static("x-amz-server-side-encryption-customer-algorithm"), + HeaderValue::from_static("AES256"), + ); + headers.insert( + HeaderName::from_static("x-amz-server-side-encryption-customer-key"), + HeaderValue::from_str(qrmk).context("invalid Snowflake qrmk chunk header")?, + ); + } + } + + Ok(headers) +} + +fn response_is_gzip_encoded(response: &reqwest::Response, url: &str) -> bool { + response + .headers() + .get(CONTENT_ENCODING) + .and_then(|value| value.to_str().ok()) + .is_some_and(|value| value.eq_ignore_ascii_case("gzip")) + || url.contains("response-content-encoding=gzip") +} + +fn is_retryable_chunk_status(status: StatusCode) -> bool { + status == StatusCode::TOO_MANY_REQUESTS + || status == StatusCode::INTERNAL_SERVER_ERROR + || status == StatusCode::SERVICE_UNAVAILABLE + || status == StatusCode::BAD_GATEWAY + || status == StatusCode::GATEWAY_TIMEOUT +} + +fn redact_url(url: &str) -> String { + match Url::parse(url) { + Ok(mut parsed) => { + parsed.set_query(None); + parsed.set_fragment(None); + parsed.to_string() + } + Err(_) => "".to_string(), + } +} + +fn epoch_millis() -> u128 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_millis() +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::{ + array::Int64Array, + datatypes::{DataType, Field as ArrowField, Schema as ArrowSchema}, + ipc::writer::StreamWriter, + }; + use feldera_types::program_schema::{ + ColumnType, Field, PropertyValue, Relation, SourcePosition, SqlIdentifier, + }; + use flate2::{write::GzEncoder, Compression}; + use std::{io::Write, sync::Arc}; + use tokio::{ + io::{duplex, AsyncWriteExt}, + sync::oneshot, + time::timeout, + }; + + fn relation(fields: &[&str]) -> Relation { + Relation { + name: SqlIdentifier::from("T".to_string()), + fields: fields + .iter() + .map(|name| { + Field::new( + SqlIdentifier::from((*name).to_string()), + ColumnType::varchar(true), + ) + }) + .collect(), + materialized: false, + properties: Default::default(), + primary_key: None, + } + } + + fn enable_skip_unused_columns(relation: &mut Relation) { + let position = SourcePosition { + start_line_number: 0, + start_column: 0, + end_line_number: 0, + end_column: 0, + }; + relation.properties.insert( + "skip_unused_columns".to_string(), + PropertyValue { + value: "true".to_string(), + key_position: position, + value_position: position, + }, + ); + } + + #[test] + fn builds_snapshot_query() { + let query = build_snapshot_query( + "DB.SCHEMA.T", + Some("ID > 10"), + &BTreeMap::new(), + &relation(&["ID", "\"customer name\""]), + ) + .unwrap(); + + assert_eq!( + query, + "SELECT ID, \"customer name\" FROM DB.SCHEMA.T WHERE ID > 10" + ); + } + + #[test] + fn reads_declared_unused_columns_by_default() { + let mut relation = relation(&["ID", "UNUSED"]); + relation.fields[1].unused = true; + + assert_eq!( + build_snapshot_query("T", None, &BTreeMap::new(), &relation).unwrap(), + "SELECT ID, UNUSED FROM T" + ); + } + + #[test] + fn skips_unused_columns_when_requested_by_table() { + let mut relation = relation(&["ID", "UNUSED"]); + relation.fields[1].unused = true; + enable_skip_unused_columns(&mut relation); + + assert_eq!( + build_snapshot_query("T", None, &BTreeMap::new(), &relation).unwrap(), + "SELECT ID FROM T" + ); + } + + #[test] + fn retains_nonnullable_unused_columns_without_defaults() { + let mut relation = relation(&["ID", "REQUIRED"]); + relation.fields[1].unused = true; + relation.fields[1].columntype.nullable = false; + enable_skip_unused_columns(&mut relation); + + assert_eq!( + build_snapshot_query("T", None, &BTreeMap::new(), &relation).unwrap(), + "SELECT ID, REQUIRED FROM T" + ); + } + + #[test] + fn rejects_empty_table() { + let err = + build_snapshot_query(" ", None, &BTreeMap::new(), &relation(&["ID"])).unwrap_err(); + assert!(err.to_string().contains("table name")); + } + + #[test] + fn maps_snowflake_columns_to_feldera_columns() { + let mapping = BTreeMap::from([ + ("uuid".to_string(), "UUID".to_string()), + ("status".to_string(), "\"applicationStatus\"".to_string()), + ("CamelCase".to_string(), "CAMEL_CASE".to_string()), + ]); + + let query = build_snapshot_query( + "T", + None, + &mapping, + &relation(&[ + "\"uuid\"", + "STATUS", + "CAMELCASE", + "\"CamelCase\"", + "UNCHANGED", + ]), + ) + .unwrap(); + + assert_eq!( + query, + "SELECT UUID AS \"uuid\", \"applicationStatus\" AS STATUS, CAMELCASE, CAMEL_CASE AS \"CamelCase\", UNCHANGED FROM T" + ); + } + + #[test] + fn rejects_unknown_column_mapping_target() { + let mapping = BTreeMap::from([("missing".to_string(), "UUID".to_string())]); + + let error = + build_snapshot_query("T", None, &mapping, &relation(&["\"uuid\""])).unwrap_err(); + + assert!(error + .to_string() + .contains("unknown Feldera column 'missing'")); + } + + #[test] + fn rejects_duplicate_mapping_for_case_insensitive_column() { + let mapping = BTreeMap::from([ + ("status".to_string(), "STATUS_A".to_string()), + ("STATUS".to_string(), "STATUS_B".to_string()), + ]); + + let error = build_snapshot_query("T", None, &mapping, &relation(&["STATUS"])).unwrap_err(); + + assert!(error.to_string().contains("multiple entries")); + } + + #[test] + fn validates_mapped_snowflake_column_identifier() { + let valid_mapping = BTreeMap::from([( + "target".to_string(), + "\"source with \"\"quote\"\"\"".to_string(), + )]); + assert_eq!( + build_snapshot_query("T", None, &valid_mapping, &relation(&["TARGET"]),).unwrap(), + "SELECT \"source with \"\"quote\"\"\" AS TARGET FROM T" + ); + + for source in ["source, other", "source AS other", "unterminated\""] { + let mapping = BTreeMap::from([("target".to_string(), source.to_string())]); + let error = + build_snapshot_query("T", None, &mapping, &relation(&["TARGET"])).unwrap_err(); + assert!(error.to_string().contains("Snowflake column identifier")); + } + } + + #[tokio::test] + async fn streams_arrow_batch_before_gzip_eof() { + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "value", + DataType::Int64, + false, + )])); + let first = + RecordBatch::try_new(schema.clone(), vec![Arc::new(Int64Array::from(vec![1, 2]))]) + .unwrap(); + let second = + RecordBatch::try_new(schema.clone(), vec![Arc::new(Int64Array::from(vec![3, 4]))]) + .unwrap(); + let mut ipc = Vec::new(); + { + let mut writer = StreamWriter::try_new(&mut ipc, &schema).unwrap(); + writer.write(&first).unwrap(); + writer.write(&second).unwrap(); + writer.finish().unwrap(); + } + + let mut encoder = GzEncoder::new(Vec::new(), Compression::default()); + encoder.write_all(&ipc).unwrap(); + let gzip = encoder.finish().unwrap(); + let trailer_start = gzip.len() - 8; + let (mut writer, reader) = duplex(gzip.len()); + let (release_sender, release_receiver) = oneshot::channel(); + tokio::spawn(async move { + writer.write_all(&gzip[..trailer_start]).await.unwrap(); + release_receiver.await.unwrap(); + writer.write_all(&gzip[trailer_start..]).await.unwrap(); + }); + + let reader = GzipDecoder::new(BufReader::new(reader)); + let mut batches = decode_arrow_reader( + reader, + "test stream".to_string(), + SnowflakeNumberMode::Decimal, + ); + let first = timeout(Duration::from_secs(2), batches.next()) + .await + .expect("first batch was not decoded before gzip EOF") + .unwrap() + .unwrap(); + assert_eq!(first.num_rows(), 2); + assert_eq!( + first + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values(), + &[1, 2] + ); + + release_sender.send(()).unwrap(); + let second = batches.next().await.unwrap().unwrap(); + assert_eq!(second.num_rows(), 2); + assert_eq!( + second + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values(), + &[3, 4] + ); + assert!(batches.next().await.is_none()); + } + + #[test] + fn redacts_presigned_chunk_url() { + assert_eq!( + redact_url("https://example.s3.amazonaws.com/chunk?X-Amz-Signature=secret#frag"), + "https://example.s3.amazonaws.com/chunk" + ); + assert_eq!(redact_url("not a url"), ""); + } + + #[test] + fn chunk_retry_uses_bounded_exponential_backoff() { + for _ in 0..100 { + assert!((250..=312).contains(&chunk_retry_backoff(1).as_millis())); + assert!((4_000..=5_000).contains(&chunk_retry_backoff(5).as_millis())); + assert_eq!(chunk_retry_backoff(u32::MAX), Duration::from_secs(8)); + } + } + + #[test] + fn retries_transient_chunk_statuses() { + for status in [429, 500, 502, 503, 504] { + assert!(is_retryable_chunk_status(StatusCode::from_u16(status).unwrap())); + } + for status in [400, 403, 404, 501, 505] { + assert!(!is_retryable_chunk_status( + StatusCode::from_u16(status).unwrap() + )); + } + } +} diff --git a/sql-to-dbsp-compiler/SQL-compiler/src/main/java/org/dbsp/sqlCompiler/compiler/frontend/connectors/ConnectorValidator.java b/sql-to-dbsp-compiler/SQL-compiler/src/main/java/org/dbsp/sqlCompiler/compiler/frontend/connectors/ConnectorValidator.java index fb79cc68e22..f6d82428947 100644 --- a/sql-to-dbsp-compiler/SQL-compiler/src/main/java/org/dbsp/sqlCompiler/compiler/frontend/connectors/ConnectorValidator.java +++ b/sql-to-dbsp-compiler/SQL-compiler/src/main/java/org/dbsp/sqlCompiler/compiler/frontend/connectors/ConnectorValidator.java @@ -244,6 +244,10 @@ public static void validateTransportConfig( validateConfig(transportConfig, outerJson, configPointer, outerStart, S3InputConfig.class, reporter); break; + case "snowflake_input": + validateConfig(transportConfig, outerJson, configPointer, + outerStart, SnowflakeReaderConfig.class, reporter); + break; case "clock": validateConfig(transportConfig, outerJson, configPointer, outerStart, ClockConfig.class, reporter); diff --git a/sql-to-dbsp-compiler/SQL-compiler/src/main/java/org/dbsp/sqlCompiler/compiler/frontend/connectors/config/SnowflakeAuthenticator.java b/sql-to-dbsp-compiler/SQL-compiler/src/main/java/org/dbsp/sqlCompiler/compiler/frontend/connectors/config/SnowflakeAuthenticator.java new file mode 100644 index 00000000000..e603c6fc28d --- /dev/null +++ b/sql-to-dbsp-compiler/SQL-compiler/src/main/java/org/dbsp/sqlCompiler/compiler/frontend/connectors/config/SnowflakeAuthenticator.java @@ -0,0 +1,7 @@ +package org.dbsp.sqlCompiler.compiler.frontend.connectors.config; + +import com.fasterxml.jackson.annotation.JsonProperty; + +public enum SnowflakeAuthenticator { + @JsonProperty("SNOWFLAKE_JWT") SnowflakeJwt, +} diff --git a/sql-to-dbsp-compiler/SQL-compiler/src/main/java/org/dbsp/sqlCompiler/compiler/frontend/connectors/config/SnowflakeIngestMode.java b/sql-to-dbsp-compiler/SQL-compiler/src/main/java/org/dbsp/sqlCompiler/compiler/frontend/connectors/config/SnowflakeIngestMode.java new file mode 100644 index 00000000000..17ea9927425 --- /dev/null +++ b/sql-to-dbsp-compiler/SQL-compiler/src/main/java/org/dbsp/sqlCompiler/compiler/frontend/connectors/config/SnowflakeIngestMode.java @@ -0,0 +1,7 @@ +package org.dbsp.sqlCompiler.compiler.frontend.connectors.config; + +import com.fasterxml.jackson.annotation.JsonProperty; + +public enum SnowflakeIngestMode { + @JsonProperty("snapshot") Snapshot, +} diff --git a/sql-to-dbsp-compiler/SQL-compiler/src/main/java/org/dbsp/sqlCompiler/compiler/frontend/connectors/config/SnowflakeNumberMode.java b/sql-to-dbsp-compiler/SQL-compiler/src/main/java/org/dbsp/sqlCompiler/compiler/frontend/connectors/config/SnowflakeNumberMode.java new file mode 100644 index 00000000000..10cbc497ca6 --- /dev/null +++ b/sql-to-dbsp-compiler/SQL-compiler/src/main/java/org/dbsp/sqlCompiler/compiler/frontend/connectors/config/SnowflakeNumberMode.java @@ -0,0 +1,8 @@ +package org.dbsp.sqlCompiler.compiler.frontend.connectors.config; + +import com.fasterxml.jackson.annotation.JsonProperty; + +public enum SnowflakeNumberMode { + @JsonProperty("decimal") Decimal, + @JsonProperty("double") Double, +} diff --git a/sql-to-dbsp-compiler/SQL-compiler/src/main/java/org/dbsp/sqlCompiler/compiler/frontend/connectors/config/SnowflakeReaderConfig.java b/sql-to-dbsp-compiler/SQL-compiler/src/main/java/org/dbsp/sqlCompiler/compiler/frontend/connectors/config/SnowflakeReaderConfig.java new file mode 100644 index 00000000000..1e0285c2685 --- /dev/null +++ b/sql-to-dbsp-compiler/SQL-compiler/src/main/java/org/dbsp/sqlCompiler/compiler/frontend/connectors/config/SnowflakeReaderConfig.java @@ -0,0 +1,131 @@ +package org.dbsp.sqlCompiler.compiler.frontend.connectors.config; + +import com.fasterxml.jackson.annotation.JsonAlias; +import com.fasterxml.jackson.annotation.JsonProperty; +import org.dbsp.sqlCompiler.compiler.frontend.connectors.ConfigReporter; +import org.dbsp.sqlCompiler.compiler.frontend.connectors.IValidateConfig; + +import javax.annotation.Nullable; +import java.util.HashMap; +import java.util.Map; + +/** Configuration for reading snapshots from Snowflake. */ +@SuppressWarnings("unused") +public class SnowflakeReaderConfig implements IValidateConfig { + @JsonProperty("account") + public String account = ""; + + @JsonProperty("user") + public String user = ""; + + @JsonProperty("authenticator") + public SnowflakeAuthenticator authenticator = SnowflakeAuthenticator.SnowflakeJwt; + + @Nullable + @JsonProperty("role") + public String role = null; + + @Nullable + @JsonProperty("warehouse") + public String warehouse = null; + + @Nullable + @JsonProperty("database") + public String database = null; + + @Nullable + @JsonProperty("schema") + public String schema = null; + + @JsonProperty("private_key_file") + public String privateKeyFile = ""; + + @Nullable + @JsonAlias({"private_key_passphrase", "private_key_file_password"}) + @JsonProperty("private_key_file_pwd") + public String privateKeyFilePwd = null; + + @JsonProperty("table") + public String table = ""; + + @JsonProperty("column_mapping") + public Map columnMapping = new HashMap<>(); + + @JsonProperty("number_mode") + public SnowflakeNumberMode numberMode = SnowflakeNumberMode.Decimal; + + @JsonProperty("mode") + public SnowflakeIngestMode mode = SnowflakeIngestMode.Snapshot; + + @JsonProperty("transaction_mode") + public SnowflakeTransactionMode transactionMode = SnowflakeTransactionMode.None; + + @Nullable + @JsonProperty("snapshot_filter") + public String snapshotFilter = null; + + @JsonProperty("num_parsers") + public int numParsers = 4; + + @Nullable + @JsonProperty("max_concurrent_readers") + public Long maxConcurrentReaders = null; + + @Override + public boolean validate(ConfigReporter reporter) { + boolean ok = true; + ok &= this.checkNonEmpty(reporter, this.account, "account"); + ok &= this.checkNonEmpty(reporter, this.user, "user"); + ok &= this.checkNonEmpty(reporter, this.privateKeyFile, "private_key_file"); + ok &= this.checkNonEmpty(reporter, this.table, "table"); + + if (this.columnMapping == null) { + reporter.warnPath("column_mapping", "Invalid configuration", + "\"column_mapping\" must be an object"); + ok = false; + } else { + for (Map.Entry entry : this.columnMapping.entrySet()) { + if (entry.getKey().isBlank()) { + reporter.warnPath("column_mapping", "Invalid configuration", + "\"column_mapping\" contains an empty Feldera column name"); + ok = false; + } + if (entry.getValue() == null || entry.getValue().isBlank()) { + reporter.warnPath("column_mapping", "Invalid configuration", + "\"column_mapping\" contains an empty Snowflake column name"); + ok = false; + } + } + } + + ok &= checkOptionalNonEmpty(reporter, this.role, "role"); + ok &= checkOptionalNonEmpty(reporter, this.warehouse, "warehouse"); + ok &= checkOptionalNonEmpty(reporter, this.database, "database"); + ok &= checkOptionalNonEmpty(reporter, this.schema, "schema"); + ok &= checkOptionalNonEmpty(reporter, this.privateKeyFilePwd, "private_key_file_pwd"); + ok &= checkOptionalNonEmpty(reporter, this.snapshotFilter, "snapshot_filter"); + + if (this.maxConcurrentReaders != null && this.maxConcurrentReaders <= 0) { + reporter.warnPath("max_concurrent_readers", "Invalid configuration", + "\"max_concurrent_readers\" must be greater than 0"); + ok = false; + } + + if (this.numParsers <= 0) { + reporter.warnPath("num_parsers", "Invalid configuration", + "\"num_parsers\" must be greater than 0"); + ok = false; + } + + return ok; + } + + private boolean checkOptionalNonEmpty(ConfigReporter reporter, @Nullable String value, String name) { + if (value != null && value.isBlank()) { + reporter.warnPath(name, "Invalid configuration", + "field \"" + name + "\" must not be empty when specified"); + return false; + } + return true; + } +} diff --git a/sql-to-dbsp-compiler/SQL-compiler/src/main/java/org/dbsp/sqlCompiler/compiler/frontend/connectors/config/SnowflakeTransactionMode.java b/sql-to-dbsp-compiler/SQL-compiler/src/main/java/org/dbsp/sqlCompiler/compiler/frontend/connectors/config/SnowflakeTransactionMode.java new file mode 100644 index 00000000000..488c42d27eb --- /dev/null +++ b/sql-to-dbsp-compiler/SQL-compiler/src/main/java/org/dbsp/sqlCompiler/compiler/frontend/connectors/config/SnowflakeTransactionMode.java @@ -0,0 +1,8 @@ +package org.dbsp.sqlCompiler.compiler.frontend.connectors.config; + +import com.fasterxml.jackson.annotation.JsonProperty; + +public enum SnowflakeTransactionMode { + @JsonProperty("none") None, + @JsonProperty("snapshot") Snapshot, +} diff --git a/sql-to-dbsp-compiler/SQL-compiler/src/test/java/org/dbsp/sqlCompiler/compiler/sql/ConnectorTests.java b/sql-to-dbsp-compiler/SQL-compiler/src/test/java/org/dbsp/sqlCompiler/compiler/sql/ConnectorTests.java index ff06d06c63d..4cc33322139 100644 --- a/sql-to-dbsp-compiler/SQL-compiler/src/test/java/org/dbsp/sqlCompiler/compiler/sql/ConnectorTests.java +++ b/sql-to-dbsp-compiler/SQL-compiler/src/test/java/org/dbsp/sqlCompiler/compiler/sql/ConnectorTests.java @@ -987,6 +987,82 @@ public void icebergReaderInvalidMode() { "field \"mode\": invalid value \"batch\""); } + // ---- Snowflake transport config ---- + + @Test + public void snowflakeReaderValidSnapshotConfig() { + tableConnectorTest(""" + "transport": { + "name": "snowflake_input", + "config": { + "account": "org-account", + "user": "svc_user", + "authenticator": "SNOWFLAKE_JWT", + "role": "DEVELOPER", + "warehouse": "DEVELOPER", + "private_key_file": "/secrets/key.p8", + "table": "DB.SCHEMA.TABLE", + "column_mapping": { + "x": "UUID" + }, + "number_mode": "double", + "mode": "snapshot", + "transaction_mode": "snapshot", + "num_parsers": 4, + "max_concurrent_readers": 8 + } + }"""); + } + + @Test + public void snowflakeReaderMissingRequiredFields() { + tableConnectorTest(""" + "transport": { + "name": "snowflake_input", + "config": { + "mode": "snapshot" + } + }""", + "required field \"account\" is missing or empty", + "required field \"user\" is missing or empty", + "required field \"private_key_file\" is missing or empty", + "required field \"table\" is missing or empty"); + } + + @Test + public void snowflakeReaderRejectsZeroParsers() { + tableConnectorTest(""" + "transport": { + "name": "snowflake_input", + "config": { + "account": "org-account", + "user": "svc_user", + "private_key_file": "/secrets/key.p8", + "table": "DB.SCHEMA.TABLE", + "num_parsers": 0 + } + }""", + "\"num_parsers\" must be greater than 0"); + } + + @Test + public void snowflakeReaderRejectsEmptyColumnMapping() { + tableConnectorTest(""" + "transport": { + "name": "snowflake_input", + "config": { + "account": "org-account", + "user": "svc_user", + "private_key_file": "/secrets/key.p8", + "table": "DB.SCHEMA.TABLE", + "column_mapping": { + "uuid": "" + } + } + }""", + "\"column_mapping\" contains an empty Snowflake column name"); + } + // ---- File transport config ---- @Test