Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
155 changes: 80 additions & 75 deletions crates/adapters/src/test/iceberg.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ use feldera_types::{
};
use serde_json::json;

use std::{collections::HashMap, time::Instant};
use std::time::Instant;
use tempfile::NamedTempFile;
use tracing::info;
use tracing_subscriber::{EnvFilter, layer::SubscriberExt, util::SubscriberInitExt};
Expand Down Expand Up @@ -64,7 +64,10 @@ fn data_to_ndjson(data: Vec<IcebergTestStruct>) -> NamedTempFile {
}

/// Read a snapshot of an Iceberg table with records of type `T` to a temporary JSON file.
fn iceberg_snapshot_to_json<T>(schema: &[Field], config: &HashMap<String, String>) -> NamedTempFile
///
/// `config` is the connector's transport config as a JSON object. This function
/// forces `mode = snapshot`.
fn iceberg_snapshot_to_json<T>(schema: &[Field], config: serde_json::Value) -> NamedTempFile
where
T: DBData
+ SerializeWithContext<SqlSerdeConfig>
Expand All @@ -78,11 +81,14 @@ where
json_file.path().display()
);

let mut config = config.clone();
config.insert("mode".to_string(), "snapshot".to_string());
let mut config = config;
config
.as_object_mut()
.expect("iceberg connector config must be a JSON object")
.insert("mode".to_string(), json!("snapshot"));

let (input_pipeline, err_receiver) =
iceberg_input_pipeline::<T>(schema, &config, &json_file.path().display().to_string());
iceberg_input_pipeline::<T>(schema, config, &json_file.path().display().to_string());
input_pipeline.start();
wait(
|| input_pipeline.status().pipeline_complete() || err_receiver.len() > 0,
Expand All @@ -102,7 +108,7 @@ where
/// Build a pipeline that reads from an Iceberg table and writes to a JSON file.
fn iceberg_input_pipeline<T>(
schema: &[Field],
config: &HashMap<String, String>,
config: serde_json::Value,
output_file_path: &str,
) -> (Controller, Receiver<String>)
where
Expand Down Expand Up @@ -198,36 +204,63 @@ fn data(n_records: usize) -> Vec<IcebergTestStruct> {
#[test]
#[cfg(feature = "iceberg-tests-fs")]
fn iceberg_localfs_input_test_unordered() {
iceberg_localfs_input_test(&[], &|_| true);
iceberg_localfs_input_test(1_000_000, json!({}), &|_| true);
}

#[test]
#[cfg(feature = "iceberg-tests-fs")]
fn iceberg_localfs_input_test_ordered() {
iceberg_localfs_input_test(
&[("timestamp_column".to_string(), "ts".to_string())],
&|_| true,
);
iceberg_localfs_input_test(1_000_000, json!({ "timestamp_column": "ts" }), &|_| true);
}

#[test]
#[cfg(feature = "iceberg-tests-fs")]
fn iceberg_localfs_input_test_ordered_with_filter() {
iceberg_localfs_input_test(
&[
("timestamp_column".to_string(), "ts".to_string()),
("snapshot_filter".to_string(), "i >= 10000".to_string()),
],
1_000_000,
json!({ "timestamp_column": "ts", "snapshot_filter": "i >= 10000" }),
&|x| x.i >= 10000,
);
}

/// A single parser task must ingest the whole snapshot correctly (the parallel
/// path defaults to 4 parsers and is covered by the tests above).
#[test]
#[cfg(feature = "iceberg-tests-fs")]
fn iceberg_localfs_input_test_single_parser() {
iceberg_localfs_input_test(100_000, json!({ "num_parsers": 1 }), &|_| true);
}

/// `transaction_mode = snapshot` on an unordered read ingests the whole snapshot
/// in one Feldera transaction; the ingested data must be identical to a
/// non-transactional read.
#[test]
#[cfg(feature = "iceberg-tests-fs")]
fn iceberg_localfs_input_test_transactional() {
iceberg_localfs_input_test(100_000, json!({ "transaction_mode": "snapshot" }), &|_| {
Comment thread
swanandx marked this conversation as resolved.
true
});
}

/// `transaction_mode = snapshot` on an ordered read ingests one Feldera
/// transaction per lateness range; the ingested data must still be complete.
#[test]
#[cfg(feature = "iceberg-tests-fs")]
fn iceberg_localfs_input_test_ordered_transactional() {
iceberg_localfs_input_test(
100_000,
json!({ "timestamp_column": "ts", "transaction_mode": "snapshot" }),
&|_| true,
);
}

#[cfg(feature = "iceberg-tests-fs")]
fn iceberg_localfs_input_test(
extra_config: &[(String, String)],
n_records: usize,
extra_config: serde_json::Value,
filter: &dyn Fn(&IcebergTestStruct) -> bool,
) {
let data = data(1_000_000);
let data = data(n_records);

let table_dir = tempfile::TempDir::new().unwrap();
let table_path = table_dir.path().display().to_string();
Expand Down Expand Up @@ -269,12 +302,18 @@ fn iceberg_localfs_input_test(
.unwrap()
.to_string();

let mut config = json!({ "metadata_location": metadata_path });
let config_obj = config.as_object_mut().unwrap();
for (key, value) in extra_config
.as_object()
.expect("extra_config must be a JSON object")
{
config_obj.insert(key.clone(), value.clone());
}

let mut json_file = iceberg_snapshot_to_json::<IcebergTestStruct>(
&IcebergTestStruct::schema_with_lateness(),
&[("metadata_location".to_string(), metadata_path.to_string())]
.into_iter()
.chain(extra_config.into_iter().cloned())
.collect::<HashMap<_, _>>(),
config,
);

let expected_zset = dbsp::OrdZSet::from_tuples(
Expand All @@ -297,37 +336,17 @@ fn iceberg_glue_s3_input_test() {
// Read delta table unordered.
let mut json_file = iceberg_snapshot_to_json::<IcebergTestStruct>(
&IcebergTestStruct::schema_with_lateness(),
&[
("catalog_type".to_string(), "glue".to_string()),
(
"glue.warehouse".to_string(),
"s3://feldera-iceberg-test/".to_string(),
),
(
"table_name".to_string(),
"iceberg_test.test_table".to_string(),
),
(
"glue.access-key-id".to_string(),
std::env::var("ICEBERG_TEST_AWS_ACCESS_KEY_ID").unwrap(),
),
(
"glue.secret-access-key".to_string(),
std::env::var("ICEBERG_TEST_AWS_SECRET_ACCESS_KEY").unwrap(),
),
("glue.region".to_string(), "us-east-1".to_string()),
(
"s3.access-key-id".to_string(),
std::env::var("ICEBERG_TEST_AWS_ACCESS_KEY_ID").unwrap(),
),
(
"s3.secret-access-key".to_string(),
std::env::var("ICEBERG_TEST_AWS_SECRET_ACCESS_KEY").unwrap(),
),
("s3.region".to_string(), "us-east-1".to_string()),
]
.into_iter()
.collect::<HashMap<_, _>>(),
json!({
"catalog_type": "glue",
"glue.warehouse": "s3://feldera-iceberg-test/",
"table_name": "iceberg_test.test_table",
"glue.access-key-id": std::env::var("ICEBERG_TEST_AWS_ACCESS_KEY_ID").unwrap(),
"glue.secret-access-key": std::env::var("ICEBERG_TEST_AWS_SECRET_ACCESS_KEY").unwrap(),
"glue.region": "us-east-1",
"s3.access-key-id": std::env::var("ICEBERG_TEST_AWS_ACCESS_KEY_ID").unwrap(),
"s3.secret-access-key": std::env::var("ICEBERG_TEST_AWS_SECRET_ACCESS_KEY").unwrap(),
"s3.region": "us-east-1",
}),
);

let zset = file_to_zset::<IcebergTestStruct>(json_file.as_file_mut());
Expand Down Expand Up @@ -381,29 +400,15 @@ fn iceberg_rest_s3_input_test() {
// Read delta table unordered.
let mut json_file = iceberg_snapshot_to_json::<IcebergTestStruct>(
&IcebergTestStruct::schema_with_lateness(),
&[
("catalog_type".to_string(), "rest".to_string()),
("rest.uri".to_string(), "http://localhost:8181".to_string()),
(
"rest.warehouse".to_string(),
"s3://feldera-iceberg-test/".to_string(),
),
(
"table_name".to_string(),
"iceberg_test.test_table".to_string(),
),
(
"s3.access-key-id".to_string(),
std::env::var("ICEBERG_TEST_AWS_ACCESS_KEY_ID").unwrap(),
),
(
"s3.secret-access-key".to_string(),
std::env::var("ICEBERG_TEST_AWS_SECRET_ACCESS_KEY").unwrap(),
),
("s3.region".to_string(), "us-east-1".to_string()),
]
.into_iter()
.collect::<HashMap<_, _>>(),
json!({
"catalog_type": "rest",
"rest.uri": "http://localhost:8181",
"rest.warehouse": "s3://feldera-iceberg-test/",
"table_name": "iceberg_test.test_table",
"s3.access-key-id": std::env::var("ICEBERG_TEST_AWS_ACCESS_KEY_ID").unwrap(),
"s3.secret-access-key": std::env::var("ICEBERG_TEST_AWS_SECRET_ACCESS_KEY").unwrap(),
"s3.region": "us-east-1",
}),
);

let zset = file_to_zset::<IcebergTestStruct>(json_file.as_file_mut());
Expand Down
48 changes: 48 additions & 0 deletions crates/feldera-types/src/transport/iceberg.rs
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,29 @@ pub enum IcebergCatalogType {
S3Tables,
}

/// Iceberg table transaction mode.
///
/// Determines how the connector breaks up its input into Feldera transactions.
///
/// * `none` - the connector does not break up its input into transactions.
/// * `snapshot` - ingest the initial snapshot of the table in one or several transactions.
///
/// # How the table snapshot is ingested using transactions
///
/// When `transaction_mode` is set to `snapshot`, the connector ingests the snapshot in one
/// or several transactions, depending on `timestamp_column`. If `timestamp_column` is not set,
/// the whole snapshot is ingested in a single Feldera transaction. If `timestamp_column` is set,
/// the connector ingests the snapshot in a series of timestamp ranges of width equal to the
/// `LATENESS` attribute of the column, each range in a separate transaction.
#[derive(Debug, Clone, Eq, PartialEq, Deserialize, Serialize, ToSchema, Default)]
pub enum IcebergTransactionMode {
#[default]
#[serde(rename = "none")]
None,
#[serde(rename = "snapshot")]
Snapshot,
}

/// AWS Glue catalog config.
#[derive(Debug, Clone, Eq, PartialEq, Deserialize, Serialize, ToSchema)]
pub struct GlueCatalogConfig {
Expand Down Expand Up @@ -186,6 +209,13 @@ pub struct IcebergReaderConfig {
/// Table read mode.
pub mode: IcebergIngestMode,

/// Transaction mode.
///
/// Determines how the connector breaks up its input into Feldera transactions.
/// See [`IcebergTransactionMode`]. Defaults to [`IcebergTransactionMode::None`].
#[serde(default)]
pub transaction_mode: IcebergTransactionMode,

/// Table column that serves as an event timestamp.
///
/// When this option is specified, and `mode` is one of `snapshot` or `snapshot_and_follow`,
Expand Down Expand Up @@ -575,6 +605,24 @@ mod test {
assert_eq!(config.max_retries(), 0);
}

#[test]
fn transaction_mode_defaults_to_none() {
let config: IcebergReaderConfig = serde_json::from_str(
r#"{"mode":"snapshot","metadata_location":"file:///tmp/t/metadata.json"}"#,
)
.unwrap();
assert_eq!(config.transaction_mode, IcebergTransactionMode::None);
}

#[test]
fn transaction_mode_snapshot_parses() {
let config: IcebergReaderConfig = serde_json::from_str(
r#"{"mode":"snapshot","metadata_location":"file:///tmp/t/metadata.json","transaction_mode":"snapshot"}"#,
)
.unwrap();
assert_eq!(config.transaction_mode, IcebergTransactionMode::Snapshot);
}

#[test]
fn reader_config_roundtrips() {
let config: IcebergReaderConfig = serde_json::from_str(
Expand Down
Loading
Loading