From b1da9734e0fb85c85ec243c9b9496be4fdda4ba0 Mon Sep 17 00:00:00 2001 From: Song Huang Date: Wed, 8 Jul 2026 00:11:48 -0400 Subject: [PATCH 1/7] Add Snowflake snapshot input connector Add an integrated Snowflake input connector backed by Arrow REST results. Support JWT key-pair authentication, projection pushdown, connector-managed snapshot transactions, concurrent chunk reads, streaming decode, and parallel parsing. --- Cargo.lock | 28 + Cargo.toml | 3 + crates/adapters/Cargo.toml | 3 + crates/adapters/src/controller.rs | 2 +- crates/adapters/src/integrated.rs | 4 + crates/adapters/src/transport.rs | 1 + crates/feldera-types/src/config.rs | 3 + crates/feldera-types/src/transport.rs | 1 + .../feldera-types/src/transport/snowflake.rs | 291 +++++++ .../pipeline-manager/src/db/types/program.rs | 1 + crates/rest-api/build.rs | 16 + crates/snowflake/Cargo.toml | 36 + crates/snowflake/src/arrow.rs | 435 ++++++++++ crates/snowflake/src/client.rs | 806 ++++++++++++++++++ crates/snowflake/src/input.rs | 386 +++++++++ crates/snowflake/src/lib.rs | 23 + crates/snowflake/src/query.rs | 154 ++++ .../connectors/ConnectorValidator.java | 4 + .../config/SnowflakeAuthenticator.java | 7 + .../config/SnowflakeIngestMode.java | 7 + .../config/SnowflakeReaderConfig.java | 107 +++ .../config/SnowflakeTransactionMode.java | 8 + .../compiler/sql/ConnectorTests.java | 54 ++ 23 files changed, 2379 insertions(+), 1 deletion(-) create mode 100644 crates/feldera-types/src/transport/snowflake.rs create mode 100644 crates/snowflake/Cargo.toml create mode 100644 crates/snowflake/src/arrow.rs create mode 100644 crates/snowflake/src/client.rs create mode 100644 crates/snowflake/src/input.rs create mode 100644 crates/snowflake/src/lib.rs create mode 100644 crates/snowflake/src/query.rs create mode 100644 sql-to-dbsp-compiler/SQL-compiler/src/main/java/org/dbsp/sqlCompiler/compiler/frontend/connectors/config/SnowflakeAuthenticator.java create mode 100644 sql-to-dbsp-compiler/SQL-compiler/src/main/java/org/dbsp/sqlCompiler/compiler/frontend/connectors/config/SnowflakeIngestMode.java create mode 100644 sql-to-dbsp-compiler/SQL-compiler/src/main/java/org/dbsp/sqlCompiler/compiler/frontend/connectors/config/SnowflakeReaderConfig.java create mode 100644 sql-to-dbsp-compiler/SQL-compiler/src/main/java/org/dbsp/sqlCompiler/compiler/frontend/connectors/config/SnowflakeTransactionMode.java diff --git a/Cargo.lock b/Cargo.lock index f1e318a8ef7..272f7064735 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", @@ -5480,6 +5481,33 @@ 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", + "openssl", + "reqwest 0.12.24", + "serde", + "serde_json", + "tokio", + "tokio-util", + "uuid", +] + [[package]] name = "feldera-sqllib" version = "0.324.0" diff --git a/Cargo.toml b/Cargo.toml index 22f32397157..6739d796c3a 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" } 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..50bc3d6545b --- /dev/null +++ b/crates/feldera-types/src/transport/snowflake.rs @@ -0,0 +1,291 @@ +use serde::{Deserialize, Serialize}; +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. + #[serde(rename = "snapshot")] + Snapshot, +} + +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. + pub private_key_file: String, + + /// RSA private key passphrase for `SNOWFLAKE_JWT` authentication. + /// + /// Required when `private_key_file` contains an encrypted 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, + + /// 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, + + /// Don't read unused columns from the Snowflake table. + /// + /// The connector always reads every column declared in the Feldera SQL + /// table by default. Set this option to avoid reading declared columns that + /// are not used by the pipeline and can be replaced with `NULL` or a + /// declared default value. + #[serde(default)] + pub skip_unused_columns: bool, + + /// 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 (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(), + mode: SnowflakeIngestMode::Snapshot, + transaction_mode: SnowflakeTransactionMode::None, + snapshot_filter: None, + skip_unused_columns: false, + 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); + 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(); + } +} 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..8b357d22ad8 --- /dev/null +++ b/crates/snowflake/Cargo.toml @@ -0,0 +1,36 @@ +[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 } +openssl = { 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..a861903840d --- /dev/null +++ b/crates/snowflake/src/arrow.rs @@ -0,0 +1,435 @@ +use std::io::Cursor; + +use anyhow::{anyhow, bail, Context, Result as AnyResult}; +use arrow::{ + array::{ + Array, Decimal128Array, 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 _}; + +/// Decode an inline Snowflake `rowsetBase64` payload into Arrow record batches. +pub(crate) fn decode_base64_ipc_stream(rowset_base64: &str) -> AnyResult> { + let bytes = BASE64_STANDARD + .decode(rowset_base64) + .context("invalid base64 in Snowflake Arrow rowset")?; + decode_ipc_stream(&bytes) +} + +/// Decode a raw Snowflake Arrow chunk into Arrow record batches. +pub(crate) fn decode_ipc_stream(bytes: &[u8]) -> 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(normalize_snowflake_batch)) + .collect() +} + +/// Convert Snowflake-specific Arrow encodings into standard Arrow types understood by Feldera. +pub(crate) fn normalize_snowflake_batch(batch: RecordBatch) -> 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 decimals = normalize_fixed(field.as_ref(), column.as_ref())?; + fields.push( + Field::new( + field.name(), + decimals.data_type().clone(), + field.is_nullable(), + ) + .with_metadata(field.metadata().clone()), + ); + columns.push(std::sync::Arc::new(decimals) as _); + } 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| signed_integer_at(column, row, field.name()).map(|value| value.map(i128::from))) + .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_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::{ArrayRef, 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).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).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).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).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)); + } +} diff --git a/crates/snowflake/src/client.rs b/crates/snowflake/src/client.rs new file mode 100644 index 00000000000..2d2624f35ec --- /dev/null +++ b/crates/snowflake/src/client.rs @@ -0,0 +1,806 @@ +use crate::arrow::{decode_base64_ipc_stream, normalize_snowflake_batch}; +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 base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine as _}; +use bytes::BytesMut; +use feldera_types::transport::snowflake::{SnowflakeAuthenticator, SnowflakeReaderConfig}; +use futures_util::{stream, stream::BoxStream, StreamExt, TryStreamExt}; +use jsonwebtoken::{encode, Algorithm, EncodingKey, Header}; +use openssl::{ + pkey::{PKey, Private}, + sha::sha256, +}; +use reqwest::{ + header::{ + HeaderMap, HeaderName, HeaderValue, ACCEPT, AUTHORIZATION, CONTENT_ENCODING, USER_AGENT, + }, + Client, StatusCode, Url, +}; +use serde::{Deserialize, Serialize}; +use serde_json::{json, Value as JsonValue}; +use std::{ + collections::HashMap, + env, fs, io, + pin::Pin, + sync::atomic::{AtomicU64, Ordering}, + time::{Duration, SystemTime, UNIX_EPOCH}, +}; +use tokio::io::{AsyncBufReadExt, AsyncRead, AsyncReadExt, BufReader}; +use tokio_util::io::StreamReader as AsyncStreamReader; +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 JWT_LIFETIME_SECONDS: u64 = 60; +const QUERY_RESULT_FORMAT: &str = "ARROW_FORCE"; +const QUERY_IN_PROGRESS_CODE: &str = "333333"; +const QUERY_IN_PROGRESS_ASYNC_CODE: &str = "333334"; +const CHUNK_DOWNLOAD_MAX_ATTEMPTS: u32 = 6; +const ARROW_DECODE_BUFFER_SIZE: usize = 256 * 1024; + +pub(crate) struct SnowflakeClient { + http: Client, + base_url: Url, + session_token: String, + sequence_id: AtomicU64, +} + +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>, +} + +#[derive(Deserialize)] +struct SnowflakeApiResponse { + #[serde(default)] + success: Option, + #[serde(default)] + code: Option, + #[serde(default)] + message: Option, + data: Option, +} + +#[derive(Deserialize)] +struct LoginData { + #[serde(default)] + token: Option, + #[serde(default, rename = "sessionToken")] + session_token: Option, +} + +#[derive(Deserialize)] +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, +} + +#[derive(Debug, Serialize)] +struct JwtClaims { + iss: String, + sub: String, + iat: u64, + exp: u64, +} + +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(crate) async fn query_arrow_batch_stream( + &self, + sql: &str, + max_concurrent_readers: usize, + ) -> 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(decode_base64_ipc_stream) + .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())), + ) + .flatten_unordered(max_concurrent_readers.max(1)); + + Ok(SnowflakeArrowBatchStream { + metadata, + batches: inline_stream.chain(downloaded_batches).boxed(), + }) + } + + fn stream_chunk<'a>( + &'a self, + chunk: SnowflakeChunk, + headers: HeaderMap, + ) -> 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); + while let Some(batch) = batches.next().await { + yield batch?; + } + } + .boxed() + } + + 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") + } + + 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"); + } + + fn endpoint(&self, path: &str) -> AnyResult { + self.base_url + .join(path.trim_start_matches('/')) + .with_context(|| format!("invalid Snowflake endpoint path: {path}")) + } + + 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) + } +} + +/// Incrementally decode Arrow IPC batches from an asynchronous byte source. +fn decode_arrow_reader( + mut reader: R, + source: String, +) -> 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).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 { + Duration::from_millis(250 * 2_u64.pow(attempt - 1)) +} + +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_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 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, + } + }) +} + +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 + .private_key_to_pem_pkcs8() + .context("error serializing Snowflake RSA private key")?; + let signing_key = EncodingKey::from_rsa_pem(&signing_key) + .context("error parsing Snowflake RSA private key")?; + encode(&header, &claims, &signing_key).context("error signing Snowflake JWT") +} + +fn parse_private_key(pem: &[u8], passphrase: Option<&str>) -> AnyResult> { + PKey::private_key_from_pem_passphrase(pem, passphrase.unwrap_or_default().as_bytes()).context( + "error parsing Snowflake RSA private key; provide 'private_key_file_pwd' if the key is encrypted", + ) +} + +fn public_key_fingerprint(key: &PKey) -> AnyResult { + let public_key_der = key + .public_key_to_der() + .context("error extracting Snowflake public key DER")?; + Ok(BASE64_STANDARD.encode(sha256(&public_key_der))) +} + +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 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::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_seconds() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs() +} + +fn epoch_millis() -> u128 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_millis() +} + +fn user_agent() -> String { + format!("{CLIENT_APP_ID}/{CLIENT_APP_VERSION}") +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::{ + array::Int64Array, + datatypes::{DataType, Field, Schema}, + ipc::writer::StreamWriter, + }; + use flate2::{write::GzEncoder, Compression}; + use openssl::{rsa::Rsa, symm::Cipher}; + use std::{io::Write, sync::Arc}; + use tokio::{ + io::{duplex, AsyncWriteExt}, + sync::oneshot, + time::timeout, + }; + + #[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()); + } + + #[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 = PKey::from_rsa(Rsa::generate(2048).unwrap()).unwrap(); + let encrypted = key + .private_key_to_pem_pkcs8_passphrase(Cipher::aes_256_cbc(), b"secret") + .unwrap(); + + assert!(parse_private_key(&encrypted, None).is_err()); + let decrypted = parse_private_key(&encrypted, Some("secret")).unwrap(); + assert_eq!( + public_key_fingerprint(&decrypted).unwrap(), + public_key_fingerprint(&key).unwrap() + ); + } + + #[tokio::test] + async fn streams_arrow_batch_before_gzip_eof() { + let schema = Arc::new(Schema::new(vec![Field::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()); + 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"), ""); + } + + #[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(), + mode: Default::default(), + transaction_mode: Default::default(), + snapshot_filter: None, + skip_unused_columns: false, + num_parsers: 4, + max_concurrent_readers: None, + }; + let client = SnowflakeClient::from_reader_config(&config).await?; + let query = + env::var("SNOWFLAKE_TEST_QUERY").unwrap_or_else(|_| "select 1::int as ID".to_string()); + let result = client.query_arrow_batch_stream(&query, 1).await?; + let batches = result.batches.collect::>().await; + let batches = batches.into_iter().collect::>>()?; + let rows = batches + .into_iter() + .map(|batch| batch.num_rows()) + .sum::(); + assert_eq!(rows, 1); + Ok(()) + } +} diff --git a/crates/snowflake/src/input.rs b/crates/snowflake/src/input.rs new file mode 100644 index 00000000000..a8e6c4dbef1 --- /dev/null +++ b/crates/snowflake/src/input.rs @@ -0,0 +1,386 @@ +use crate::{client::SnowflakeClient, query::build_snapshot_query, snowflake_input_serde_config}; +use anyhow::{anyhow, Context, Result as AnyResult}; +use arrow::record_batch::RecordBatch; +use chrono::Utc; +use dbsp::circuit::tokio::TOKIO; +use feldera_adapterlib::{ + catalog::{ArrowStream, InputCollectionHandle}, + errors::journal::ControllerError, + format::{InputBuffer, ParseError}, + transport::{ + InputConsumer, InputEndpoint, InputQueue, InputQueueEntry, InputReader, InputReaderCommand, + IntegratedInputEndpoint, Watermark, + }, + PipelineState, +}; +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 std::{sync::Arc, thread}; +use tokio::{ + select, + sync::{ + mpsc, + watch::{channel, Receiver, Sender}, + }, +}; + +/// 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 { + None + } +} + +impl IntegratedInputEndpoint for SnowflakeInputEndpoint { + fn open( + self: Box, + input_handle: &InputCollectionHandle, + _seek: Option, + ) -> AnyResult> { + Ok(Box::new(SnowflakeInputReader::new( + &self.inner, + input_handle, + )?)) + } +} + +struct SnowflakeInputReader { + sender: Sender, + inner: Arc, +} + +impl SnowflakeInputReader { + fn new( + endpoint: &Arc, + input_handle: &InputCollectionHandle, + ) -> AnyResult { + endpoint.config.validate().map_err(|e| { + ControllerError::invalid_transport_configuration(&endpoint.endpoint_name, &e) + })?; + + let (sender, receiver) = channel(PipelineState::Paused); + let endpoint_clone = endpoint.clone(); + let receiver_clone = receiver.clone(); + + let (init_status_sender, mut init_status_receiver) = + mpsc::channel::>(1); + + 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(); + + 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 { .. } => { + let (total, _, timestamps) = self.inner.queue.flush_with_aux(); + self.inner.consumer.extended( + total, + None, + timestamps + .into_iter() + .map(|(timestamp, ())| Watermark::new(timestamp, None)) + .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, +} + +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, + } + } + + 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.skip_unused_columns, + &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()) + .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") + } + }) + .buffered(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, ()) + .with_buffer(buffer) + .with_start_transaction(transaction_label.clone()); + + self.queue.push_entry(entry, errors); + timestamp = Utc::now(); + } + + if transaction_label.is_some() { + self.queue.push_entry( + InputQueueEntry::new_with_aux(timestamp, ()) + .with_start_transaction(transaction_label) + .with_commit_transaction(true), + 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(), ()).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..c435a4ffd5a --- /dev/null +++ b/crates/snowflake/src/lib.rs @@ -0,0 +1,23 @@ +mod arrow; +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..9274cdf0643 --- /dev/null +++ b/crates/snowflake/src/query.rs @@ -0,0 +1,154 @@ +use anyhow::{bail, Result as AnyResult}; +use feldera_types::program_schema::Relation; + +pub(crate) fn build_snapshot_query( + table: &str, + snapshot_filter: Option<&str>, + skip_unused_columns: bool, + relation: &Relation, +) -> AnyResult { + let table = table.trim(); + if table.is_empty() { + bail!("Snowflake table name must not be empty"); + } + + let skip_unused_columns = + skip_unused_columns || relation.get_property("skip_unused_columns") == Some("true"); + let columns = relation + .fields + .iter() + .filter(|field| { + !skip_unused_columns + || !field.unused + || (!field.columntype.nullable && field.default.is_none()) + }) + .map(|field| 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) +} + +#[cfg(test)] +mod tests { + use super::*; + use feldera_types::program_schema::{ + ColumnType, Field, PropertyValue, Relation, SourcePosition, SqlIdentifier, + }; + + 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, + } + } + + #[test] + fn builds_snapshot_query() { + let query = build_snapshot_query( + "DB.SCHEMA.T", + Some("ID > 10"), + false, + &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, false, &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; + 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, + }, + ); + + assert_eq!( + build_snapshot_query("T", None, false, &relation).unwrap(), + "SELECT ID FROM T" + ); + } + + #[test] + fn skips_unused_columns_when_requested_by_connector() { + let mut relation = relation(&["ID", "UNUSED"]); + relation.fields[1].unused = true; + + assert_eq!( + build_snapshot_query("T", None, true, &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; + + assert_eq!( + build_snapshot_query("T", None, true, &relation).unwrap(), + "SELECT ID, REQUIRED FROM T" + ); + } + + #[test] + fn rejects_empty_table() { + let err = build_snapshot_query(" ", None, false, &relation(&["ID"])).unwrap_err(); + assert!(err.to_string().contains("table name")); + } +} 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/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..cfffe49b9a0 --- /dev/null +++ b/sql-to-dbsp-compiler/SQL-compiler/src/main/java/org/dbsp/sqlCompiler/compiler/frontend/connectors/config/SnowflakeReaderConfig.java @@ -0,0 +1,107 @@ +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; + +/** 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("mode") + public SnowflakeIngestMode mode = SnowflakeIngestMode.Snapshot; + + @JsonProperty("transaction_mode") + public SnowflakeTransactionMode transactionMode = SnowflakeTransactionMode.None; + + @Nullable + @JsonProperty("snapshot_filter") + public String snapshotFilter = null; + + @JsonProperty("skip_unused_columns") + public boolean skipUnusedColumns = false; + + @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"); + + 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..ca7774d21f2 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,60 @@ 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", + "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"); + } + // ---- File transport config ---- @Test From 9e2340558b156d2d1bb5a5eb195cf641e09eef8f Mon Sep 17 00:00:00 2001 From: Song Huang Date: Wed, 8 Jul 2026 15:21:50 -0400 Subject: [PATCH 2/7] Organize and tune Snowflake connector --- Cargo.lock | 15 +- Cargo.toml | 2 + .../feldera-types/src/transport/snowflake.rs | 9 +- crates/snowflake/Cargo.toml | 5 +- crates/snowflake/src/auth.rs | 139 +++++ crates/snowflake/src/client.rs | 559 +----------------- crates/snowflake/src/input.rs | 2 +- crates/snowflake/src/lib.rs | 1 + crates/snowflake/src/query.rs | 439 +++++++++++++- 9 files changed, 624 insertions(+), 547 deletions(-) create mode 100644 crates/snowflake/src/auth.rs diff --git a/Cargo.lock b/Cargo.lock index 272f7064735..df2768e4996 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4618,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" @@ -5499,10 +5508,13 @@ dependencies = [ "futures-util", "jsonwebtoken", "log", - "openssl", + "pkcs8 0.10.2", + "rand 0.8.6", "reqwest 0.12.24", + "rsa", "serde", "serde_json", + "sha2 0.10.9", "tokio", "tokio-util", "uuid", @@ -8978,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 6739d796c3a..3433d9be7cc 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -216,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"] } @@ -247,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/feldera-types/src/transport/snowflake.rs b/crates/feldera-types/src/transport/snowflake.rs index 50bc3d6545b..c38cccee911 100644 --- a/crates/feldera-types/src/transport/snowflake.rs +++ b/crates/feldera-types/src/transport/snowflake.rs @@ -35,6 +35,10 @@ pub enum SnowflakeTransactionMode { 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, } @@ -80,11 +84,14 @@ pub struct SnowflakeReaderConfig { 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 private key. + /// Required when `private_key_file` contains an encrypted PKCS#8 private + /// key. #[serde( default, alias = "private_key_passphrase", diff --git a/crates/snowflake/Cargo.toml b/crates/snowflake/Cargo.toml index 8b357d22ad8..ae153a8b26a 100644 --- a/crates/snowflake/Cargo.toml +++ b/crates/snowflake/Cargo.toml @@ -29,8 +29,11 @@ arrow = { workspace = true, features = ["ipc"] } futures-util = { workspace = true } base64 = { workspace = true } jsonwebtoken = { workspace = true } -openssl = { workspace = true } +pkcs8 = { workspace = true } +rsa = { workspace = true } +sha2 = { workspace = true } uuid = { workspace = true, features = ["v4"] } [dev-dependencies] flate2 = { workspace = true } +rand = { workspace = true } 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 index 2d2624f35ec..e5c6ee0815e 100644 --- a/crates/snowflake/src/client.rs +++ b/crates/snowflake/src/client.rs @@ -1,73 +1,37 @@ -use crate::arrow::{decode_base64_ipc_stream, normalize_snowflake_batch}; +use crate::auth::{jwt_token, login_account_name}; 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 base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine as _}; -use bytes::BytesMut; -use feldera_types::transport::snowflake::{SnowflakeAuthenticator, SnowflakeReaderConfig}; -use futures_util::{stream, stream::BoxStream, StreamExt, TryStreamExt}; -use jsonwebtoken::{encode, Algorithm, EncodingKey, Header}; -use openssl::{ - pkey::{PKey, Private}, - sha::sha256, -}; +use feldera_types::transport::snowflake::SnowflakeReaderConfig; use reqwest::{ - header::{ - HeaderMap, HeaderName, HeaderValue, ACCEPT, AUTHORIZATION, CONTENT_ENCODING, USER_AGENT, - }, - Client, StatusCode, Url, + header::{HeaderMap, HeaderValue, ACCEPT, AUTHORIZATION, USER_AGENT}, + Client, Url, }; -use serde::{Deserialize, Serialize}; +use serde::Deserialize; use serde_json::{json, Value as JsonValue}; -use std::{ - collections::HashMap, - env, fs, io, - pin::Pin, - sync::atomic::{AtomicU64, Ordering}, - time::{Duration, SystemTime, UNIX_EPOCH}, -}; -use tokio::io::{AsyncBufReadExt, AsyncRead, AsyncReadExt, BufReader}; -use tokio_util::io::StreamReader as AsyncStreamReader; +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 JWT_LIFETIME_SECONDS: u64 = 60; const QUERY_RESULT_FORMAT: &str = "ARROW_FORCE"; -const QUERY_IN_PROGRESS_CODE: &str = "333333"; -const QUERY_IN_PROGRESS_ASYNC_CODE: &str = "333334"; -const CHUNK_DOWNLOAD_MAX_ATTEMPTS: u32 = 6; -const ARROW_DECODE_BUFFER_SIZE: usize = 256 * 1024; pub(crate) struct SnowflakeClient { - http: Client, - base_url: Url, + pub(super) http: Client, + pub(super) base_url: Url, session_token: String, - sequence_id: AtomicU64, -} - -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>, + pub(super) sequence_id: AtomicU64, } #[derive(Deserialize)] -struct SnowflakeApiResponse { +pub(super) struct SnowflakeApiResponse { #[serde(default)] success: Option, #[serde(default)] - code: Option, + pub(super) code: Option, #[serde(default)] message: Option, - data: Option, + pub(super) data: Option, } #[derive(Deserialize)] @@ -78,39 +42,6 @@ struct LoginData { session_token: Option, } -#[derive(Deserialize)] -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, -} - -#[derive(Debug, Serialize)] -struct JwtClaims { - iss: String, - sub: String, - iat: u64, - exp: u64, -} - impl SnowflakeClient { pub(crate) async fn from_reader_config(config: &SnowflakeReaderConfig) -> AnyResult { config.validate().map_err(|error| anyhow!(error))?; @@ -161,191 +92,13 @@ impl SnowflakeClient { Ok(client) } - pub(crate) async fn query_arrow_batch_stream( - &self, - sql: &str, - max_concurrent_readers: usize, - ) -> 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(decode_base64_ipc_stream) - .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())), - ) - .flatten_unordered(max_concurrent_readers.max(1)); - - Ok(SnowflakeArrowBatchStream { - metadata, - batches: inline_stream.chain(downloaded_batches).boxed(), - }) - } - - fn stream_chunk<'a>( - &'a self, - chunk: SnowflakeChunk, - headers: HeaderMap, - ) -> 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); - while let Some(batch) = batches.next().await { - yield batch?; - } - } - .boxed() - } - - 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") - } - - 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"); - } - - fn endpoint(&self, path: &str) -> AnyResult { + pub(super) fn endpoint(&self, path: &str) -> AnyResult { self.base_url .join(path.trim_start_matches('/')) .with_context(|| format!("invalid Snowflake endpoint path: {path}")) } - fn session_headers(&self) -> AnyResult { + 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())?); @@ -358,51 +111,6 @@ impl SnowflakeClient { } } -/// Incrementally decode Arrow IPC batches from an asynchronous byte source. -fn decode_arrow_reader( - mut reader: R, - source: String, -) -> 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).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 { - Duration::from_millis(250 * 2_u64.pow(attempt - 1)) -} - fn account_base_url(account: &str) -> AnyResult { let account = account.trim(); validate_account_identifier(account)?; @@ -439,17 +147,6 @@ fn validate_account_identifier(account: &str) -> AnyResult<()> { Ok(()) } -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 login_url(base_url: &Url, config: &SnowflakeReaderConfig) -> AnyResult { let mut url = base_url.join("session/v1/login-request")?; { @@ -504,54 +201,7 @@ fn login_body(config: &SnowflakeReaderConfig, jwt: &str) -> JsonValue { }) } -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 - .private_key_to_pem_pkcs8() - .context("error serializing Snowflake RSA private key")?; - let signing_key = EncodingKey::from_rsa_pem(&signing_key) - .context("error parsing Snowflake RSA private key")?; - encode(&header, &claims, &signing_key).context("error signing Snowflake JWT") -} - -fn parse_private_key(pem: &[u8], passphrase: Option<&str>) -> AnyResult> { - PKey::private_key_from_pem_passphrase(pem, passphrase.unwrap_or_default().as_bytes()).context( - "error parsing Snowflake RSA private key; provide 'private_key_file_pwd' if the key is encrypted", - ) -} - -fn public_key_fingerprint(key: &PKey) -> AnyResult { - let public_key_der = key - .public_key_to_der() - .context("error extracting Snowflake public key DER")?; - Ok(BASE64_STANDARD.encode(sha256(&public_key_der))) -} - -fn response_data(response: SnowflakeApiResponse, operation: &str) -> AnyResult { +pub(super) fn response_data(response: SnowflakeApiResponse, operation: &str) -> AnyResult { if response.success == Some(false) { bail!( "{operation} failed: code={} message={}", @@ -569,73 +219,6 @@ fn response_data(response: SnowflakeApiResponse, operation: &str) -> AnyRe }) } -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::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_seconds() -> u64 { - SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap_or_default() - .as_secs() -} - -fn epoch_millis() -> u128 { - SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap_or_default() - .as_millis() -} - fn user_agent() -> String { format!("{CLIENT_APP_ID}/{CLIENT_APP_VERSION}") } @@ -643,19 +226,8 @@ fn user_agent() -> String { #[cfg(test)] mod tests { use super::*; - use arrow::{ - array::Int64Array, - datatypes::{DataType, Field, Schema}, - ipc::writer::StreamWriter, - }; - use flate2::{write::GzEncoder, Compression}; - use openssl::{rsa::Rsa, symm::Cipher}; - use std::{io::Write, sync::Arc}; - use tokio::{ - io::{duplex, AsyncWriteExt}, - sync::oneshot, - time::timeout, - }; + use feldera_types::transport::snowflake::SnowflakeAuthenticator; + use futures_util::StreamExt; #[test] fn derives_account_urls() { @@ -671,103 +243,6 @@ mod tests { assert!(account_base_url("xy12345.snowflakecomputing.com").is_err()); } - #[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 = PKey::from_rsa(Rsa::generate(2048).unwrap()).unwrap(); - let encrypted = key - .private_key_to_pem_pkcs8_passphrase(Cipher::aes_256_cbc(), b"secret") - .unwrap(); - - assert!(parse_private_key(&encrypted, None).is_err()); - let decrypted = parse_private_key(&encrypted, Some("secret")).unwrap(); - assert_eq!( - public_key_fingerprint(&decrypted).unwrap(), - public_key_fingerprint(&key).unwrap() - ); - } - - #[tokio::test] - async fn streams_arrow_batch_before_gzip_eof() { - let schema = Arc::new(Schema::new(vec![Field::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()); - 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"), ""); - } - #[ignore] #[tokio::test] async fn snowflake_arrow_smoke_test() -> AnyResult<()> { diff --git a/crates/snowflake/src/input.rs b/crates/snowflake/src/input.rs index a8e6c4dbef1..1a1e937a6c7 100644 --- a/crates/snowflake/src/input.rs +++ b/crates/snowflake/src/input.rs @@ -306,7 +306,7 @@ impl SnowflakeInputEndpointInner { .context("Snowflake Arrow batch parser task failed") } }) - .buffered(self.config.num_parsers as usize); + .buffer_unordered(self.config.num_parsers as usize); while let Some(parsed_batch) = parsed_batches.next().await { wait_running(&mut receiver).await; diff --git a/crates/snowflake/src/lib.rs b/crates/snowflake/src/lib.rs index c435a4ffd5a..c61a79002a3 100644 --- a/crates/snowflake/src/lib.rs +++ b/crates/snowflake/src/lib.rs @@ -1,4 +1,5 @@ mod arrow; +mod auth; mod client; mod input; mod query; diff --git a/crates/snowflake/src/query.rs b/crates/snowflake/src/query.rs index 9274cdf0643..33f13ed4efc 100644 --- a/crates/snowflake/src/query.rs +++ b/crates/snowflake/src/query.rs @@ -1,5 +1,45 @@ -use anyhow::{bail, Result as AnyResult}; +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; +use futures_util::{stream, stream::BoxStream, StreamExt, TryStreamExt}; +use reqwest::{ + header::{HeaderMap, HeaderName, HeaderValue, CONTENT_ENCODING}, + StatusCode, Url, +}; +use serde::Deserialize; +use serde_json::json; +use std::{ + collections::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 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>, +} pub(crate) fn build_snapshot_query( table: &str, @@ -46,12 +86,334 @@ pub(crate) fn build_snapshot_query( 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, + ) -> 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(decode_base64_ipc_stream) + .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())), + ) + .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, + ) -> 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); + 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, +) -> 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).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 { + Duration::from_millis(250 * 2_u64.pow(attempt - 1)) +} + +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::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 { @@ -151,4 +513,79 @@ mod tests { let err = build_snapshot_query(" ", None, false, &relation(&["ID"])).unwrap_err(); assert!(err.to_string().contains("table name")); } + + #[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()); + 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"), ""); + } } From ee86feca469fe57d439d74afa0b2499af1604d26 Mon Sep 17 00:00:00 2001 From: Song Huang Date: Mon, 13 Jul 2026 14:21:43 -0400 Subject: [PATCH 3/7] Support Snowflake column mapping --- .../feldera-types/src/transport/snowflake.rs | 71 ++++++++ crates/snowflake/src/client.rs | 1 + crates/snowflake/src/input.rs | 1 + crates/snowflake/src/query.rs | 153 +++++++++++++++++- .../config/SnowflakeReaderConfig.java | 24 +++ .../compiler/sql/ConnectorTests.java | 21 +++ 6 files changed, 264 insertions(+), 7 deletions(-) diff --git a/crates/feldera-types/src/transport/snowflake.rs b/crates/feldera-types/src/transport/snowflake.rs index c38cccee911..b21caafc55d 100644 --- a/crates/feldera-types/src/transport/snowflake.rs +++ b/crates/feldera-types/src/transport/snowflake.rs @@ -1,4 +1,5 @@ use serde::{Deserialize, Serialize}; +use std::collections::BTreeMap; use utoipa::ToSchema; /// Snowflake table read mode. @@ -107,6 +108,20 @@ pub struct SnowflakeReaderConfig { /// `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, + /// Table read mode. /// /// Only `"snapshot"` is supported. @@ -169,6 +184,17 @@ impl SnowflakeReaderConfig { } } + 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), @@ -225,6 +251,7 @@ mod tests { private_key_file: "/secrets/key.p8".to_string(), private_key_file_pwd: None, table: "DB.SCHEMA.TABLE".to_string(), + column_mapping: BTreeMap::new(), mode: SnowflakeIngestMode::Snapshot, transaction_mode: SnowflakeTransactionMode::None, snapshot_filter: None, @@ -295,4 +322,48 @@ mod tests { 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 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/snowflake/src/client.rs b/crates/snowflake/src/client.rs index e5c6ee0815e..2cf8de786e4 100644 --- a/crates/snowflake/src/client.rs +++ b/crates/snowflake/src/client.rs @@ -258,6 +258,7 @@ mod tests { .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(), mode: Default::default(), transaction_mode: Default::default(), snapshot_filter: None, diff --git a/crates/snowflake/src/input.rs b/crates/snowflake/src/input.rs index 1a1e937a6c7..4d139aa5479 100644 --- a/crates/snowflake/src/input.rs +++ b/crates/snowflake/src/input.rs @@ -215,6 +215,7 @@ impl SnowflakeInputEndpointInner { &self.config.table, self.config.snapshot_filter.as_deref(), self.config.skip_unused_columns, + &self.config.column_mapping, &schema, ), }; diff --git a/crates/snowflake/src/query.rs b/crates/snowflake/src/query.rs index 33f13ed4efc..2bd7e80f462 100644 --- a/crates/snowflake/src/query.rs +++ b/crates/snowflake/src/query.rs @@ -16,7 +16,7 @@ use reqwest::{ use serde::Deserialize; use serde_json::json; use std::{ - collections::HashMap, + collections::{BTreeMap, HashMap}, io, pin::Pin, sync::atomic::Ordering, @@ -41,10 +41,40 @@ pub(crate) struct SnowflakeArrowBatchStream<'a> { 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>, skip_unused_columns: bool, + column_mapping: &BTreeMap, relation: &Relation, ) -> AnyResult { let table = table.trim(); @@ -54,6 +84,34 @@ pub(crate) fn build_snapshot_query( let skip_unused_columns = 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() @@ -62,7 +120,12 @@ pub(crate) fn build_snapshot_query( || !field.unused || (!field.columntype.nullable && field.default.is_none()) }) - .map(|field| field.name.sql_name()) + .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() { @@ -439,6 +502,7 @@ mod tests { "DB.SCHEMA.T", Some("ID > 10"), false, + &BTreeMap::new(), &relation(&["ID", "\"customer name\""]), ) .unwrap(); @@ -455,7 +519,7 @@ mod tests { relation.fields[1].unused = true; assert_eq!( - build_snapshot_query("T", None, false, &relation).unwrap(), + build_snapshot_query("T", None, false, &BTreeMap::new(), &relation).unwrap(), "SELECT ID, UNUSED FROM T" ); } @@ -480,7 +544,7 @@ mod tests { ); assert_eq!( - build_snapshot_query("T", None, false, &relation).unwrap(), + build_snapshot_query("T", None, false, &BTreeMap::new(), &relation).unwrap(), "SELECT ID FROM T" ); } @@ -491,7 +555,7 @@ mod tests { relation.fields[1].unused = true; assert_eq!( - build_snapshot_query("T", None, true, &relation).unwrap(), + build_snapshot_query("T", None, true, &BTreeMap::new(), &relation).unwrap(), "SELECT ID FROM T" ); } @@ -503,17 +567,92 @@ mod tests { relation.fields[1].columntype.nullable = false; assert_eq!( - build_snapshot_query("T", None, true, &relation).unwrap(), + build_snapshot_query("T", None, true, &BTreeMap::new(), &relation).unwrap(), "SELECT ID, REQUIRED FROM T" ); } #[test] fn rejects_empty_table() { - let err = build_snapshot_query(" ", None, false, &relation(&["ID"])).unwrap_err(); + let err = build_snapshot_query(" ", None, false, &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, + false, + &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, false, &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, false, &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, false, &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, false, &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( 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 index cfffe49b9a0..8b5726f6ffc 100644 --- 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 @@ -6,6 +6,8 @@ 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") @@ -46,6 +48,9 @@ public class SnowflakeReaderConfig implements IValidateConfig { @JsonProperty("table") public String table = ""; + @JsonProperty("column_mapping") + public Map columnMapping = new HashMap<>(); + @JsonProperty("mode") public SnowflakeIngestMode mode = SnowflakeIngestMode.Snapshot; @@ -74,6 +79,25 @@ public boolean validate(ConfigReporter reporter) { 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"); 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 ca7774d21f2..b1b1a3c83b3 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 @@ -1002,6 +1002,9 @@ public void snowflakeReaderValidSnapshotConfig() { "warehouse": "DEVELOPER", "private_key_file": "/secrets/key.p8", "table": "DB.SCHEMA.TABLE", + "column_mapping": { + "x": "UUID" + }, "mode": "snapshot", "transaction_mode": "snapshot", "num_parsers": 4, @@ -1041,6 +1044,24 @@ public void snowflakeReaderRejectsZeroParsers() { "\"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 From e6a5d252187192339c42607017c131c1773e77b1 Mon Sep 17 00:00:00 2001 From: Song Huang Date: Tue, 14 Jul 2026 12:06:16 -0400 Subject: [PATCH 4/7] Support configurable Snowflake number conversion --- .../feldera-types/src/transport/snowflake.rs | 40 ++++++ crates/snowflake/src/arrow.rs | 135 +++++++++++++++--- crates/snowflake/src/client.rs | 68 +++++++-- crates/snowflake/src/input.rs | 6 +- crates/snowflake/src/query.rs | 19 ++- .../config/SnowflakeNumberMode.java | 8 ++ .../config/SnowflakeReaderConfig.java | 3 + .../compiler/sql/ConnectorTests.java | 1 + 8 files changed, 246 insertions(+), 34 deletions(-) create mode 100644 sql-to-dbsp-compiler/SQL-compiler/src/main/java/org/dbsp/sqlCompiler/compiler/frontend/connectors/config/SnowflakeNumberMode.java diff --git a/crates/feldera-types/src/transport/snowflake.rs b/crates/feldera-types/src/transport/snowflake.rs index b21caafc55d..5908c9d512b 100644 --- a/crates/feldera-types/src/transport/snowflake.rs +++ b/crates/feldera-types/src/transport/snowflake.rs @@ -44,6 +44,21 @@ pub enum SnowflakeTransactionMode { 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 { @@ -122,6 +137,13 @@ pub struct SnowflakeReaderConfig { #[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. @@ -252,6 +274,7 @@ mod tests { 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, @@ -303,6 +326,7 @@ mod tests { .unwrap(); assert_eq!(config.transaction_mode, SnowflakeTransactionMode::Snapshot); + assert_eq!(config.number_mode, SnowflakeNumberMode::Decimal); config.validate().unwrap(); } @@ -342,6 +366,22 @@ mod tests { 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(); diff --git a/crates/snowflake/src/arrow.rs b/crates/snowflake/src/arrow.rs index a861903840d..66838b23159 100644 --- a/crates/snowflake/src/arrow.rs +++ b/crates/snowflake/src/arrow.rs @@ -3,25 +3,32 @@ use std::io::Cursor; use anyhow::{anyhow, bail, Context, Result as AnyResult}; use arrow::{ array::{ - Array, Decimal128Array, Int16Array, Int32Array, Int64Array, Int8Array, StructArray, - Time64NanosecondArray, TimestampMicrosecondArray, + 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) -> AnyResult> { +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) + decode_ipc_stream(&bytes, number_mode) } /// Decode a raw Snowflake Arrow chunk into Arrow record batches. -pub(crate) fn decode_ipc_stream(bytes: &[u8]) -> AnyResult> { +pub(crate) fn decode_ipc_stream( + bytes: &[u8], + number_mode: SnowflakeNumberMode, +) -> AnyResult> { if bytes.is_empty() { return Ok(Vec::new()); } @@ -31,12 +38,15 @@ pub(crate) fn decode_ipc_stream(bytes: &[u8]) -> AnyResult> { reader .map(|batch| batch.context("error reading Snowflake Arrow IPC record batch")) - .map(|batch| batch.and_then(normalize_snowflake_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) -> AnyResult { +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()); @@ -48,16 +58,23 @@ pub(crate) fn normalize_snowflake_batch(batch: RecordBatch) -> AnyResult { + 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(), - decimals.data_type().clone(), + normalized.data_type().clone(), field.is_nullable(), ) .with_metadata(field.metadata().clone()), ); - columns.push(std::sync::Arc::new(decimals) as _); + columns.push(normalized); } else if matches!( logical_type, Some("TIMESTAMP_NTZ" | "TIMESTAMP_LTZ" | "TIMESTAMP_TZ") @@ -115,7 +132,7 @@ fn normalize_fixed(field: &Field, column: &dyn Array) -> AnyResult>>()?; Decimal128Array::from(values) @@ -128,6 +145,26 @@ fn normalize_fixed(field: &Field, column: &dyn Array) -> AnyResult 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, @@ -288,7 +325,7 @@ fn struct_child<'a, T: Array + 'static>( mod tests { use super::*; use arrow::{ - array::{ArrayRef, Int32Array, Int64Array, StringArray, StructArray}, + array::{Int32Array, Int64Array, StringArray, StructArray}, datatypes::{DataType, Field, Schema}, ipc::writer::StreamWriter, }; @@ -316,7 +353,7 @@ mod tests { writer.finish().unwrap(); } - let batches = decode_ipc_stream(&bytes).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"); @@ -358,7 +395,7 @@ mod tests { ) .unwrap(); - let normalized = normalize_snowflake_batch(batch).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())) @@ -391,7 +428,7 @@ mod tests { ) .unwrap(); - let normalized = normalize_snowflake_batch(batch).unwrap(); + let normalized = normalize_snowflake_batch(batch, SnowflakeNumberMode::Decimal).unwrap(); let timestamp = normalized .column(0) .as_any() @@ -419,7 +456,7 @@ mod tests { ) .unwrap(); - let normalized = normalize_snowflake_batch(batch).unwrap(); + let normalized = normalize_snowflake_batch(batch, SnowflakeNumberMode::Decimal).unwrap(); assert_eq!( normalized.schema().field(0).data_type(), &DataType::Decimal128(10, 2) @@ -432,4 +469,70 @@ mod tests { 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/client.rs b/crates/snowflake/src/client.rs index 2cf8de786e4..da79bc5c179 100644 --- a/crates/snowflake/src/client.rs +++ b/crates/snowflake/src/client.rs @@ -226,7 +226,11 @@ fn user_agent() -> String { #[cfg(test)] mod tests { use super::*; - use feldera_types::transport::snowflake::SnowflakeAuthenticator; + use arrow::{ + array::{Array, Decimal128Array, Float64Array}, + datatypes::DataType, + }; + use feldera_types::transport::snowflake::{SnowflakeAuthenticator, SnowflakeNumberMode}; use futures_util::StreamExt; #[test] @@ -259,6 +263,7 @@ mod tests { 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, @@ -267,16 +272,57 @@ mod tests { max_concurrent_readers: None, }; let client = SnowflakeClient::from_reader_config(&config).await?; - let query = - env::var("SNOWFLAKE_TEST_QUERY").unwrap_or_else(|_| "select 1::int as ID".to_string()); - let result = client.query_arrow_batch_stream(&query, 1).await?; - let batches = result.batches.collect::>().await; - let batches = batches.into_iter().collect::>>()?; - let rows = batches - .into_iter() - .map(|batch| batch.num_rows()) - .sum::(); - assert_eq!(rows, 1); + 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 index 4d139aa5479..eb4f25e15c1 100644 --- a/crates/snowflake/src/input.rs +++ b/crates/snowflake/src/input.rs @@ -252,7 +252,11 @@ impl SnowflakeInputEndpointInner { }; let result_stream = match client - .query_arrow_batch_stream(&query, self.config.max_concurrent_readers()) + .query_arrow_batch_stream( + &query, + self.config.max_concurrent_readers(), + self.config.number_mode, + ) .await { Ok(result) => result, diff --git a/crates/snowflake/src/query.rs b/crates/snowflake/src/query.rs index 2bd7e80f462..baaca2a3574 100644 --- a/crates/snowflake/src/query.rs +++ b/crates/snowflake/src/query.rs @@ -7,7 +7,7 @@ use arrow::{buffer::Buffer, ipc::reader::StreamDecoder, record_batch::RecordBatc use async_compression::tokio::bufread::GzipDecoder; use async_stream::try_stream; use bytes::BytesMut; -use feldera_types::program_schema::Relation; +use feldera_types::{program_schema::Relation, transport::snowflake::SnowflakeNumberMode}; use futures_util::{stream, stream::BoxStream, StreamExt, TryStreamExt}; use reqwest::{ header::{HeaderMap, HeaderName, HeaderValue, CONTENT_ENCODING}, @@ -179,6 +179,7 @@ impl SnowflakeClient { &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(); @@ -189,7 +190,7 @@ impl SnowflakeClient { let inline_batches = data .rowset_base64 .as_deref() - .map(decode_base64_ipc_stream) + .map(|rowset| decode_base64_ipc_stream(rowset, number_mode)) .transpose()?; let chunk_headers = chunk_headers(&data)?; let metadata = SnowflakeArrowQueryMetadata { @@ -206,7 +207,7 @@ impl SnowflakeClient { let downloaded_batches = stream::iter( data.chunks .into_iter() - .map(move |chunk| self.stream_chunk(chunk, chunk_headers.clone())), + .map(move |chunk| self.stream_chunk(chunk, chunk_headers.clone(), number_mode)), ) .flatten_unordered(max_concurrent_readers.max(1)); @@ -282,6 +283,7 @@ impl SnowflakeClient { &'a self, chunk: SnowflakeChunk, headers: HeaderMap, + number_mode: SnowflakeNumberMode, ) -> BoxStream<'a, AnyResult> { try_stream! { let chunk_url = redact_url(&chunk.url); @@ -308,7 +310,7 @@ impl SnowflakeClient { } else { Box::pin(body) }; - let mut batches = decode_arrow_reader(reader, chunk_url); + let mut batches = decode_arrow_reader(reader, chunk_url, number_mode); while let Some(batch) = batches.next().await { yield batch?; } @@ -358,6 +360,7 @@ impl SnowflakeClient { fn decode_arrow_reader( mut reader: R, source: String, + number_mode: SnowflakeNumberMode, ) -> BoxStream<'static, AnyResult> where R: AsyncRead + Send + Unpin + 'static, @@ -381,7 +384,7 @@ where .decode(&mut buffer) .with_context(|| format!("error decoding Arrow IPC data from {source}"))? { - yield normalize_snowflake_batch(batch).with_context(|| { + yield normalize_snowflake_batch(batch, number_mode).with_context(|| { format!("error normalizing Arrow data from {source}") })?; } @@ -687,7 +690,11 @@ mod tests { }); let reader = GzipDecoder::new(BufReader::new(reader)); - let mut batches = decode_arrow_reader(reader, "test stream".to_string()); + 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") 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 index 8b5726f6ffc..0ecd871e407 100644 --- 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 @@ -51,6 +51,9 @@ public class SnowflakeReaderConfig implements IValidateConfig { @JsonProperty("column_mapping") public Map columnMapping = new HashMap<>(); + @JsonProperty("number_mode") + public SnowflakeNumberMode numberMode = SnowflakeNumberMode.Decimal; + @JsonProperty("mode") public SnowflakeIngestMode mode = SnowflakeIngestMode.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 b1b1a3c83b3..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 @@ -1005,6 +1005,7 @@ public void snowflakeReaderValidSnapshotConfig() { "column_mapping": { "x": "UUID" }, + "number_mode": "double", "mode": "snapshot", "transaction_mode": "snapshot", "num_parsers": 4, From b2a1f58598d791062788af71dcf1aa72f9b978dc Mon Sep 17 00:00:00 2001 From: Song Huang Date: Tue, 14 Jul 2026 12:41:04 -0400 Subject: [PATCH 5/7] Use table property for skipping unused Snowflake columns --- .../feldera-types/src/transport/snowflake.rs | 10 --- crates/snowflake/src/client.rs | 1 - crates/snowflake/src/input.rs | 1 - crates/snowflake/src/query.rs | 72 ++++++++----------- .../config/SnowflakeReaderConfig.java | 3 - 5 files changed, 30 insertions(+), 57 deletions(-) diff --git a/crates/feldera-types/src/transport/snowflake.rs b/crates/feldera-types/src/transport/snowflake.rs index 5908c9d512b..ee51e26ec8f 100644 --- a/crates/feldera-types/src/transport/snowflake.rs +++ b/crates/feldera-types/src/transport/snowflake.rs @@ -165,15 +165,6 @@ pub struct SnowflakeReaderConfig { #[serde(default, skip_serializing_if = "Option::is_none")] pub snapshot_filter: Option, - /// Don't read unused columns from the Snowflake table. - /// - /// The connector always reads every column declared in the Feldera SQL - /// table by default. Set this option to avoid reading declared columns that - /// are not used by the pipeline and can be replaced with `NULL` or a - /// declared default value. - #[serde(default)] - pub skip_unused_columns: bool, - /// Number of parallel tasks used to deserialize Arrow record batches. /// /// Recommended range: 1–10. Default: 4. @@ -278,7 +269,6 @@ mod tests { mode: SnowflakeIngestMode::Snapshot, transaction_mode: SnowflakeTransactionMode::None, snapshot_filter: None, - skip_unused_columns: false, num_parsers: default_num_parsers(), max_concurrent_readers: None, } diff --git a/crates/snowflake/src/client.rs b/crates/snowflake/src/client.rs index da79bc5c179..851c0786118 100644 --- a/crates/snowflake/src/client.rs +++ b/crates/snowflake/src/client.rs @@ -267,7 +267,6 @@ mod tests { mode: Default::default(), transaction_mode: Default::default(), snapshot_filter: None, - skip_unused_columns: false, num_parsers: 4, max_concurrent_readers: None, }; diff --git a/crates/snowflake/src/input.rs b/crates/snowflake/src/input.rs index eb4f25e15c1..01760562be2 100644 --- a/crates/snowflake/src/input.rs +++ b/crates/snowflake/src/input.rs @@ -214,7 +214,6 @@ impl SnowflakeInputEndpointInner { SnowflakeIngestMode::Snapshot => build_snapshot_query( &self.config.table, self.config.snapshot_filter.as_deref(), - self.config.skip_unused_columns, &self.config.column_mapping, &schema, ), diff --git a/crates/snowflake/src/query.rs b/crates/snowflake/src/query.rs index baaca2a3574..8cbeac5d4cc 100644 --- a/crates/snowflake/src/query.rs +++ b/crates/snowflake/src/query.rs @@ -73,7 +73,6 @@ fn snowflake_column_identifier(identifier: &str) -> AnyResult<&str> { pub(crate) fn build_snapshot_query( table: &str, snapshot_filter: Option<&str>, - skip_unused_columns: bool, column_mapping: &BTreeMap, relation: &Relation, ) -> AnyResult { @@ -82,8 +81,7 @@ pub(crate) fn build_snapshot_query( bail!("Snowflake table name must not be empty"); } - let skip_unused_columns = - skip_unused_columns || relation.get_property("skip_unused_columns") == Some("true"); + 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(); @@ -499,12 +497,28 @@ mod tests { } } + 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"), - false, &BTreeMap::new(), &relation(&["ID", "\"customer name\""]), ) @@ -522,7 +536,7 @@ mod tests { relation.fields[1].unused = true; assert_eq!( - build_snapshot_query("T", None, false, &BTreeMap::new(), &relation).unwrap(), + build_snapshot_query("T", None, &BTreeMap::new(), &relation).unwrap(), "SELECT ID, UNUSED FROM T" ); } @@ -531,34 +545,10 @@ mod tests { fn skips_unused_columns_when_requested_by_table() { let mut relation = relation(&["ID", "UNUSED"]); relation.fields[1].unused = true; - 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, - }, - ); - - assert_eq!( - build_snapshot_query("T", None, false, &BTreeMap::new(), &relation).unwrap(), - "SELECT ID FROM T" - ); - } - - #[test] - fn skips_unused_columns_when_requested_by_connector() { - let mut relation = relation(&["ID", "UNUSED"]); - relation.fields[1].unused = true; + enable_skip_unused_columns(&mut relation); assert_eq!( - build_snapshot_query("T", None, true, &BTreeMap::new(), &relation).unwrap(), + build_snapshot_query("T", None, &BTreeMap::new(), &relation).unwrap(), "SELECT ID FROM T" ); } @@ -568,17 +558,18 @@ mod tests { 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, true, &BTreeMap::new(), &relation).unwrap(), + 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, false, &BTreeMap::new(), &relation(&["ID"])) - .unwrap_err(); + let err = + build_snapshot_query(" ", None, &BTreeMap::new(), &relation(&["ID"])).unwrap_err(); assert!(err.to_string().contains("table name")); } @@ -593,7 +584,6 @@ mod tests { let query = build_snapshot_query( "T", None, - false, &mapping, &relation(&[ "\"uuid\"", @@ -616,7 +606,7 @@ mod tests { let mapping = BTreeMap::from([("missing".to_string(), "UUID".to_string())]); let error = - build_snapshot_query("T", None, false, &mapping, &relation(&["\"uuid\""])).unwrap_err(); + build_snapshot_query("T", None, &mapping, &relation(&["\"uuid\""])).unwrap_err(); assert!(error .to_string() @@ -630,8 +620,7 @@ mod tests { ("STATUS".to_string(), "STATUS_B".to_string()), ]); - let error = - build_snapshot_query("T", None, false, &mapping, &relation(&["STATUS"])).unwrap_err(); + let error = build_snapshot_query("T", None, &mapping, &relation(&["STATUS"])).unwrap_err(); assert!(error.to_string().contains("multiple entries")); } @@ -643,15 +632,14 @@ mod tests { "\"source with \"\"quote\"\"\"".to_string(), )]); assert_eq!( - build_snapshot_query("T", None, false, &valid_mapping, &relation(&["TARGET"]),) - .unwrap(), + 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, false, &mapping, &relation(&["TARGET"])) - .unwrap_err(); + let error = + build_snapshot_query("T", None, &mapping, &relation(&["TARGET"])).unwrap_err(); assert!(error.to_string().contains("Snowflake column identifier")); } } 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 index 0ecd871e407..1e0285c2685 100644 --- 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 @@ -64,9 +64,6 @@ public class SnowflakeReaderConfig implements IValidateConfig { @JsonProperty("snapshot_filter") public String snapshotFilter = null; - @JsonProperty("skip_unused_columns") - public boolean skipUnusedColumns = false; - @JsonProperty("num_parsers") public int numParsers = 4; From 07a094926cf86faa216d503a761d78684faffe4a Mon Sep 17 00:00:00 2001 From: Song Huang Date: Tue, 21 Jul 2026 18:19:33 -0400 Subject: [PATCH 6/7] Add full snapshot Snowflake checkpoint support --- crates/snowflake/src/input.rs | 154 +++++++++++++++++++++++++--------- 1 file changed, 116 insertions(+), 38 deletions(-) diff --git a/crates/snowflake/src/input.rs b/crates/snowflake/src/input.rs index 01760562be2..ee615178a1f 100644 --- a/crates/snowflake/src/input.rs +++ b/crates/snowflake/src/input.rs @@ -1,17 +1,17 @@ use crate::{client::SnowflakeClient, query::build_snapshot_query, snowflake_input_serde_config}; -use anyhow::{anyhow, Context, Result as AnyResult}; +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, Watermark, + IntegratedInputEndpoint, Resume, Watermark, parse_resume_info, }, - PipelineState, }; use feldera_types::{ adapter_stats::ConnectorHealth, @@ -21,12 +21,16 @@ use feldera_types::{ }; use futures_util::StreamExt; use log::{debug, info}; -use std::{sync::Arc, thread}; +use serde::{Deserialize, Serialize}; +use std::{ + sync::{Arc, Mutex}, + thread, +}; use tokio::{ select, sync::{ mpsc, - watch::{channel, Receiver, Sender}, + watch::{Receiver, Sender, channel}, }, }; @@ -53,7 +57,9 @@ impl SnowflakeInputEndpoint { impl InputEndpoint for SnowflakeInputEndpoint { fn fault_tolerance(&self) -> Option { - None + // This test-only fault tolerance supports full-refresh checkpoints, not recovery + // within a snapshot: before EOI, recovery reruns the entire snapshot. + Some(FtModel::AtLeastOnce) } } @@ -61,15 +67,37 @@ impl IntegratedInputEndpoint for SnowflakeInputEndpoint { fn open( self: Box, input_handle: &InputCollectionHandle, - _seek: Option, + 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, @@ -79,17 +107,28 @@ 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 (sender, receiver) = channel(PipelineState::Paused); - let endpoint_clone = endpoint.clone(); - let receiver_clone = receiver.clone(); + 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 (init_status_sender, mut init_status_receiver) = - mpsc::channel::>(1); + let (sender, receiver) = channel(PipelineState::Paused); let input_stream = input_handle .handle @@ -97,20 +136,32 @@ impl SnowflakeInputReader { let schema = input_handle.schema.clone(); let endpoint_name = endpoint.endpoint_name.clone(); - 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; + 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"); + .expect("failed to spawn snowflake-input tokio wrapper thread"); - init_status_receiver.blocking_recv().ok_or_else(|| { - anyhow!("worker thread terminated unexpectedly during initialization") - })??; + init_status_receiver.blocking_recv().ok_or_else(|| { + anyhow!("worker thread terminated unexpectedly during initialization") + })??; + } Ok(Self { sender, @@ -135,14 +186,39 @@ impl InputReader for SnowflakeInputReader { InputReaderCommand::Pause => { let _ = self.sender.send_replace(PipelineState::Paused); } - InputReaderCommand::Queue { .. } => { - let (total, _, timestamps) = self.inner.queue.flush_with_aux(); + 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, - None, - timestamps + Some(resume), + resume_info .into_iter() - .map(|(timestamp, ())| Watermark::new(timestamp, None)) + .map(|(timestamp, resume_info)| { + Watermark::new( + timestamp, + resume_info + .map(|resume_info| serde_json::to_value(resume_info).unwrap()), + ) + }) .collect(), ); } @@ -167,7 +243,8 @@ struct SnowflakeInputEndpointInner { endpoint_name: String, config: SnowflakeReaderConfig, consumer: Box, - queue: InputQueue, + queue: InputQueue>, + last_resume_status: Mutex>, } impl SnowflakeInputEndpointInner { @@ -182,6 +259,7 @@ impl SnowflakeInputEndpointInner { config, consumer, queue, + last_resume_status: Mutex::new(Some(SnowflakeResumeInfo::initial())), } } @@ -326,7 +404,7 @@ impl SnowflakeInputEndpointInner { } }; - let entry = InputQueueEntry::new_with_aux(timestamp, ()) + let entry = InputQueueEntry::new_with_aux(timestamp, None) .with_buffer(buffer) .with_start_transaction(transaction_label.clone()); @@ -334,14 +412,14 @@ impl SnowflakeInputEndpointInner { timestamp = Utc::now(); } + let mut completion_entry = + InputQueueEntry::new_with_aux(timestamp, Some(SnowflakeResumeInfo::eoi())); if transaction_label.is_some() { - self.queue.push_entry( - InputQueueEntry::new_with_aux(timestamp, ()) - .with_start_transaction(transaction_label) - .with_commit_transaction(true), - Vec::new(), - ); + 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: '{}')", @@ -356,7 +434,7 @@ impl SnowflakeInputEndpointInner { // 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(), ()).with_commit_transaction(true), + InputQueueEntry::new_with_aux(Utc::now(), None).with_commit_transaction(true), Vec::new(), ); } From 337271d2d951446a750a7585b24559117b7303ef Mon Sep 17 00:00:00 2001 From: Song Huang Date: Tue, 21 Jul 2026 18:52:46 -0400 Subject: [PATCH 7/7] Add jittered retries for Snowflake chunk downloads --- crates/snowflake/Cargo.toml | 2 +- crates/snowflake/src/query.rs | 33 ++++++++++++++++++++++++++++++++- 2 files changed, 33 insertions(+), 2 deletions(-) diff --git a/crates/snowflake/Cargo.toml b/crates/snowflake/Cargo.toml index ae153a8b26a..25849937ed8 100644 --- a/crates/snowflake/Cargo.toml +++ b/crates/snowflake/Cargo.toml @@ -30,10 +30,10 @@ 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 } -rand = { workspace = true } diff --git a/crates/snowflake/src/query.rs b/crates/snowflake/src/query.rs index 8cbeac5d4cc..aa9a148861c 100644 --- a/crates/snowflake/src/query.rs +++ b/crates/snowflake/src/query.rs @@ -9,6 +9,7 @@ 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, @@ -29,6 +30,8 @@ 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 { @@ -397,7 +400,13 @@ where } fn chunk_retry_backoff(attempt: u32) -> Duration { - Duration::from_millis(250 * 2_u64.pow(attempt - 1)) + 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 { @@ -437,6 +446,7 @@ fn response_is_gzip_encoded(response: &reqwest::Response, url: &str) -> bool { 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 @@ -722,4 +732,25 @@ mod tests { ); 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() + )); + } + } }