This is an automated email from the ASF dual-hosted git repository. numinnex pushed a commit to branch iceberg_connector_atleast_once in repository https://gitbox.apache.org/repos/asf/iggy.git
commit c629fb97bd3e22e43d5626b6abefa39db3e0158c Author: Grzegorz Koszyk <[email protected]> AuthorDate: Fri Nov 6 12:47:03 2026 +0100 address review --- .claude/skills/connector-sink/SKILL.md | 2 +- core/connectors/sdk/src/lib.rs | 2 +- core/connectors/sinks/iceberg_sink/README.md | 51 +-- core/connectors/sinks/iceberg_sink/src/lib.rs | 2 +- .../sinks/iceberg_sink/src/router/commit.rs | 420 +++++++++++---------- .../iceberg_sink/src/router/dynamic_router.rs | 326 ++++++++-------- .../sinks/iceberg_sink/src/router/mod.rs | 173 ++++++--- .../sinks/iceberg_sink/src/router/redelivery.rs | 358 ++++++++++++------ .../sinks/iceberg_sink/src/router/static_router.rs | 198 ++++------ core/connectors/sinks/iceberg_sink/src/sink.rs | 36 +- 10 files changed, 884 insertions(+), 684 deletions(-) diff --git a/.claude/skills/connector-sink/SKILL.md b/.claude/skills/connector-sink/SKILL.md index d9df101e4..516d41d0b 100644 --- a/.claude/skills/connector-sink/SKILL.md +++ b/.claude/skills/connector-sink/SKILL.md @@ -96,7 +96,7 @@ Runtime retries on any `Err` - same batch can be redelivered. In-tree guarantees - **MongoDB**: composite `_id` from `stream:topic:partition:message_id`, via `insert_many(docs).ordered(false)`. Per-document duplicate-key errors don't abort the batch - effective idempotency under retry. - **Elasticsearch** (`bulk_index_documents`): bulk action specifies only `_index`, no `_id` -> ES auto-generates ids -> NOT idempotent. Extensions should add `"_id"`. - **HTTP**: endpoint-dependent. Document in the plugin README. -- **Iceberg**: transactional commit per table. `router/commit.rs` tags each snapshot with `iggy.commit-id`, and a commit retry first checks whether the catalog already holds it. At least once, `router/redelivery.rs` records the last message of each source partition in the snapshot summary. A redelivered batch that carries the last message of any append on the current snapshot history appends only the messages after the newest one. +- **Iceberg**: transactional commit per table. `router/commit.rs` tags each snapshot with `iggy.commit-id`, and a commit retry first checks whether the catalog already holds it. At least once, `router/redelivery.rs` records the first offset and the last message of each source partition in the snapshot summary. A redelivered batch that carries the last message of an append on the current snapshot history skips the start of the batch that such appends cover without a gap. When writing a new sink, prefer dedup-on-write using the message `id` field. Mongo's upsert-style is the cleanest reference. diff --git a/core/connectors/sdk/src/lib.rs b/core/connectors/sdk/src/lib.rs index b3decafc5..239e625c9 100644 --- a/core/connectors/sdk/src/lib.rs +++ b/core/connectors/sdk/src/lib.rs @@ -162,7 +162,7 @@ pub trait Sink: Send + Sync { /// /// Returning `Ok` promises that, from now on, a successful [`Sink::consume`] never drops a /// message without failing the call, unless the sink's configuration explicitly tells it to - /// skip such messages and it logs each one with its offset, and that a successful + /// skip messages it cannot write and it logs each one with its offset, and that a successful /// [`Sink::flush`] proves that the destination acknowledged every message of the covered /// partitions handed over so far. /// Return [`Error::Unsupported`] when the sink, or the configuration it was opened with, diff --git a/core/connectors/sinks/iceberg_sink/README.md b/core/connectors/sinks/iceberg_sink/README.md index 1b2bb3f20..70740b7a5 100644 --- a/core/connectors/sinks/iceberg_sink/README.md +++ b/core/connectors/sinks/iceberg_sink/README.md @@ -51,7 +51,7 @@ max_retry_delay = "30s" - **store_region**: The region of the object storage. Required, including for S3-compatible stores that ignore it. - **store_class**: The storage class to use. **Currently, only S3-compatible storage is supported.** - **store_path_style_access**: Use path-style S3 URLs (`http://host/bucket/key`). Defaults to `true`, which MinIO-style endpoints require; set to `false` for stores that only accept virtual-hosted-style URLs. -- **invalid_rows**: What to do with a row that does not match the table schema. `fail`, the default, fails the whole batch. `skip` writes the rows that match and logs each skipped row as a warning with its offset. At least once, it also decides what happens to a message that dynamic routing cannot route. +- **invalid_rows**: What to do with a row that does not match the table schema. `fail`, the default, commits the rows before the first such row and fails the call at it. `skip` writes the rows that match and logs each skipped row as a warning with its offset. At least once, it also decides what happens to a message without a JSON payload and to a message that dynamic routing cannot route. At most once, a message without a JSON payload is always skipped with a warning. - **max_retries**: Total attempts per catalog commit, including the first. Defaults to `5`. See [Commit Retries](#commit-retries). - **retry_delay**: Delay before the first commit retry, doubled for each further retry. Defaults to `1s`. - **max_retry_delay**: Upper bound for one delay between commit attempts. Defaults to `30s`. @@ -62,8 +62,8 @@ Static routing requires a non-empty `tables` list; dynamic routing requires a no ## Static Routing With `dynamic_routing = false`, every batch is copied to every table in `tables`. Startup fails -if none of them loads. A table that does not load at startup is loaded again by each batch. At -most once, a batch skips a table that still does not load. At least once, it fails instead: +if none of them loads. At most once, a table that does not load at startup is skipped until the +connector restarts. At least once, each batch loads it again and fails while it does not load: transiently while the catalog cannot be reached, so that the runtime delivers it again, and for good when the table does not exist. A name without a namespace is skipped at startup, and at least once the sink refuses to start with one. The writer uses the schema and default partition @@ -180,22 +180,28 @@ retries and how a stopped delivery resumes. At least once: `invalid_rows = "skip"` skips such a message instead, with a warning naming its offset. Skipping is an explicit choice that drops the message, and the runtime stores offsets past it. - **Redelivered messages are appended once.** Each append records in its snapshot summary the - source of its rows (`iggy.stream`, `iggy.topic`, `iggy.partition-id`) and its last message - (`iggy.last-offset`, `iggy.last-message-id`, and `iggy.last-timestamp`, the time the server - stored it). When a batch arrives again, and it carries the last message of any append on the - table's current snapshot history, the sink appends only the messages after the newest append. - A crash between the commit and the offset store, a backlog of several unflushed batches, or a - batch that committed to one table and failed on another, therefore does not append rows twice. - Each table keeps its own record, so a batch copied to several tables is appended once to each. -- **Limits.** The record is the newest message a table holds, not a list of its rows. A message at - or below it is skipped even when the table lacks its row, for example a row deleted later, or - one that `invalid_rows = "skip"` left out. To append such messages again, replay them with - `guarantee = "at_most_once"`, which does not skip anything. A batch that carries no appended last - message, for example after `batch_length` changed, cannot be proven and is appended in full. - After a rollback of the table, restart the sink, since it keeps the record of its own commits - in memory. Keep one sink writing each table and partition, and keep the stream and topic - identifiers in its configuration unchanged, since the record names them. At most once, appends - record nothing. + source of its rows (`iggy.stream`, `iggy.topic`, `iggy.partition-id`), the offset of its first + message (`iggy.first-offset`), and its last message (`iggy.last-offset`, `iggy.last-message-id`, + and `iggy.last-timestamp`, the time the server stored it). An append on the table's current + snapshot history counts as proven when a batch that arrives again carries its last message. + The appends after the end of that batch, up to the newest one, count too, so a batch cut at + other offsets than before still finds them. The sink skips the start of the batch that the + counted appends cover without a gap, and appends the rest. A crash between the commit and the + offset store, a backlog of several unflushed batches, or a batch that committed to one table and + failed on another, therefore does not append rows twice. Each table keeps its own record, so a + batch copied to several tables is appended once to each. +- **Limits.** The record is the range of offsets each append covers, not a list of its rows. A + message inside a counted append is skipped even when the table lacks its row, for example a row + deleted later, or one that `invalid_rows = "skip"` left out. To append such messages again, + replay them with `guarantee = "at_most_once"`, which does not skip anything. A message that no + append covers, for example one an at-most-once run lost, is appended again, and so is every + message after it in the batch, even when an append covers some of them. A batch that carries no + appended last message, for example after `batch_length` changed, cannot be proven and is + appended in full. After a rollback of the table, or after it is dropped and created again under + the same name, restart the sink, since it keeps the record of its own commits in memory and, + with static routing, the table it loaded. Keep one sink writing each table and partition, and + keep the stream and topic identifiers in its configuration unchanged, since the record names + them. At most once, appends record nothing. ## Commit Rate and Table Maintenance @@ -218,7 +224,7 @@ At 2,000 messages per second, `batch_length = 100` creates about 20 snapshots pe values above about 2. A poll also returns once its reply reaches 1 MiB, and `max_wait` bounds how long a message waits for its batch. -At least once, each snapshot summary also carries the six `iggy.*` properties described in +At least once, each snapshot summary also carries the seven `iggy.*` properties described in [Delivery Guarantee](#delivery-guarantee), which add a few hundred bytes per snapshot to the table metadata until the snapshot expires. @@ -230,8 +236,9 @@ schedule that keeps the time travel you need. Some managed catalogs run it autom Use `schema = "json"` on the stream. Each JSON object represents a table row; nested objects and arrays are supported when they match the table schema. Unknown fields are ignored. Missing -nullable fields become null. A missing required field or an incompatible value fails the table -write for the whole batch, unless `invalid_rows = "skip"` leaves out only that row. +nullable fields become null. A missing required field or an incompatible value stops the write +at that row: the rows before it are committed and the call fails, unless `invalid_rows = "skip"` +leaves out only that row. Sources that wrap row data in an envelope need a transform when the table schema describes the inner row. Otherwise, the Arrow JSON reader maps envelope keys to table columns, producing nulls diff --git a/core/connectors/sinks/iceberg_sink/src/lib.rs b/core/connectors/sinks/iceberg_sink/src/lib.rs index 64f6a2162..49fbb56bf 100644 --- a/core/connectors/sinks/iceberg_sink/src/lib.rs +++ b/core/connectors/sinks/iceberg_sink/src/lib.rs @@ -40,7 +40,7 @@ const DEFAULT_MAX_RETRY_DELAY: &str = "30s"; #[serde(rename_all = "lowercase")] #[strum(serialize_all = "lowercase")] pub enum InvalidRows { - /// Fail the whole batch. + /// Commit the rows before the first row that does not match, then fail the call at it. #[default] Fail, /// Write the rows that match, and log each skipped row with its offset. diff --git a/core/connectors/sinks/iceberg_sink/src/router/commit.rs b/core/connectors/sinks/iceberg_sink/src/router/commit.rs index 768ecc31e..8f19d651e 100644 --- a/core/connectors/sinks/iceberg_sink/src/router/commit.rs +++ b/core/connectors/sinks/iceberg_sink/src/router/commit.rs @@ -104,7 +104,7 @@ pub(crate) fn is_transient(error: &iceberg::Error) -> bool { } /// The error for a table that does not load from the catalog. An unreachable, failing or -/// throttling catalog is transient. Any other refusal, such as a missing table or denied access, +/// throttling catalog is transient. A missing table, or any other refusal such as denied access, /// fails the same way until the configuration or the catalog changes. pub(crate) fn table_load_error(table: &str, error: &iceberg::Error) -> Error { let reason = format!( @@ -113,8 +113,10 @@ pub(crate) fn table_load_error(table: &str, error: &iceberg::Error) -> Error { ); if is_transient(error) { Error::Connection(reason) - } else { + } else if is_missing(error) { Error::InvalidConfigValue(reason) + } else { + Error::PermanentHttpError(reason) } } @@ -172,8 +174,8 @@ fn holds_commit(table: &Table, commit_id: &str) -> bool { pub(crate) mod tests { use super::*; use crate::InvalidRows; - use crate::router::redelivery::{HeldMessages, LastMessage, snapshot_properties}; - use crate::router::{BatchOrigin, Written, write_data}; + use crate::router::redelivery::{Append, HeldMessages, snapshot_properties}; + use crate::router::{Appender, BatchOrigin, Written}; use async_trait::async_trait; use iceberg::memory::{MEMORY_CATALOG_WAREHOUSE, MemoryCatalogBuilder}; use iceberg::spec::{ @@ -185,9 +187,9 @@ pub(crate) mod tests { CatalogBuilder, MemoryCatalog, Namespace, NamespaceIdent, TableCommit, TableCreation, }; use iceberg_catalog_rest::{REST_CATALOG_PROP_URI, RestCatalog, RestCatalogBuilder}; - use iggy_connector_sdk::Schema; - use iggy_connector_sdk::{ConsumedMessage, Payload}; + use iggy_connector_sdk::{ConsumedMessage, FlushPartition, Payload, Schema, TopicMetadata}; use std::collections::VecDeque; + use std::ops::RangeInclusive; use std::sync::Mutex; use std::sync::atomic::{AtomicU32, Ordering}; use std::time::Duration; @@ -200,6 +202,68 @@ pub(crate) mod tests { max_delay: Duration::from_millis(1), }; + pub(crate) const SINGLE_ATTEMPT: RetryPolicy = RetryPolicy { + max_attempts: 1, + ..FAST_RETRY + }; + + /// A memory catalog holding the table `test.events`, created with `properties`, whose rows + /// hold one required `id` column. + pub(crate) async fn catalog_with_table( + properties: HashMap<String, String>, + ) -> (MemoryCatalog, Table) { + let catalog = MemoryCatalogBuilder::default() + .load( + "memory", + HashMap::from([( + MEMORY_CATALOG_WAREHOUSE.to_owned(), + "memory://warehouse".to_owned(), + )]), + ) + .await + .expect("Failed to create catalog"); + let namespace = NamespaceIdent::new("test".to_owned()); + catalog + .create_namespace(&namespace, HashMap::new()) + .await + .expect("Failed to create namespace"); + let schema = IcebergSchema::builder() + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Long)).into(), + ]) + .build() + .expect("Failed to build schema"); + let table = catalog + .create_table( + &namespace, + TableCreation::builder() + .name("events".to_owned()) + .schema(schema) + .properties(properties) + .build(), + ) + .await + .expect("Failed to create table"); + (catalog, table) + } + + /// A message at `offset`, with an ID and a timestamp that tell it apart. + pub(crate) fn message(offset: u64) -> ConsumedMessage { + ConsumedMessage { + id: u128::from(offset) + 100, + offset, + checksum: 0, + timestamp: offset + 1_000, + origin_timestamp: 0, + headers: None, + payload: Payload::Raw(Vec::new()), + } + } + + pub(crate) fn messages(offsets: RangeInclusive<u64>) -> Vec<ConsumedMessage> { + offsets.map(message).collect() + } + /// What `FlakyCatalog` does with the next `update_table` call. #[derive(Debug, Clone, Copy)] enum CommitFailure { @@ -232,38 +296,7 @@ pub(crate) mod tests { failures: &[CommitFailure], properties: HashMap<String, String>, ) -> (Self, Table) { - let inner = MemoryCatalogBuilder::default() - .load( - "memory", - HashMap::from([( - MEMORY_CATALOG_WAREHOUSE.to_owned(), - "memory://warehouse".to_owned(), - )]), - ) - .await - .expect("Failed to create catalog"); - let namespace = NamespaceIdent::new("test".to_owned()); - inner - .create_namespace(&namespace, HashMap::new()) - .await - .expect("Failed to create namespace"); - let schema = IcebergSchema::builder() - .with_fields(vec![ - NestedField::required(1, "id", Type::Primitive(PrimitiveType::Long)).into(), - ]) - .build() - .expect("Failed to build schema"); - let table = inner - .create_table( - &namespace, - TableCreation::builder() - .name("events".to_owned()) - .schema(schema) - .properties(properties) - .build(), - ) - .await - .expect("Failed to create table"); + let (inner, table) = catalog_with_table(properties).await; let catalog = Self { inner, failures: Mutex::new(failures.iter().copied().collect()), @@ -588,7 +621,6 @@ pub(crate) mod tests { } /// A REST catalog whose server answers every request with `status` and no body. - /// A REST catalog that answers every call with `status`. pub(crate) async fn rest_catalog_answering(status: StatusCode) -> RestCatalog { let listener = TcpListener::bind("127.0.0.1:0") .await @@ -624,23 +656,7 @@ pub(crate) mod tests { partition_id: 2, }; - fn message(offset: u64) -> ConsumedMessage { - ConsumedMessage { - id: u128::from(offset) + 100, - offset, - checksum: 0, - timestamp: offset + 1_000, - origin_timestamp: 0, - headers: None, - payload: Payload::Raw(Vec::new()), - } - } - - fn messages(offsets: std::ops::RangeInclusive<u64>) -> Vec<ConsumedMessage> { - offsets.map(message).collect() - } - - fn json_messages(offsets: std::ops::RangeInclusive<u64>) -> Vec<ConsumedMessage> { + fn json_messages(offsets: RangeInclusive<u64>) -> Vec<ConsumedMessage> { offsets .map(|offset| ConsumedMessage { payload: Payload::Json(simd_json::json!({ "id": offset })), @@ -653,13 +669,17 @@ pub(crate) mod tests { messages.iter().map(|message| message.offset).collect() } - async fn commit_up_to( + fn append(offsets: RangeInclusive<u64>) -> Append { + Append::new(&message(*offsets.start()), &message(*offsets.end())) + } + + async fn commit_append( catalog: &FlakyCatalog, table: &Table, origin: &BatchOrigin<'_>, - offset: u64, + offsets: RangeInclusive<u64>, ) { - let properties = snapshot_properties(origin, LastMessage::from(&message(offset))); + let properties = snapshot_properties(origin, append(offsets)); commit_data_files( table, catalog, @@ -671,50 +691,51 @@ pub(crate) mod tests { .expect("The append should commit"); } - async fn current_table(catalog: &FlakyCatalog, table: &Table) -> Table { + async fn current_table(catalog: &dyn Catalog, table: &Table) -> Table { catalog - .inner .load_table(table.identifier()) .await .expect("Failed to load table") } - #[allow(clippy::too_many_arguments)] + /// An appender of one attempt per commit to the table of a `FlakyCatalog` that fails its + /// commits as `failures` script. + async fn appender( + failures: &[CommitFailure], + invalid_rows: InvalidRows, + at_least_once: bool, + ) -> (Appender, Table) { + let (catalog, table) = FlakyCatalog::with_table(failures).await; + let mut appender = Appender::new(Box::new(catalog), invalid_rows, SINGLE_ATTEMPT); + if at_least_once { + appender.enable_at_least_once(); + } + (appender, table) + } + async fn write( - catalog: &FlakyCatalog, + appender: &Appender, table: &Table, batch: &[ConsumedMessage], - invalid_rows: InvalidRows, - retry: RetryPolicy, - held: Option<&HeldMessages>, ) -> Result<Written, Error> { - write_data( - &ORIGIN, - batch, - table, - None, - catalog, - Schema::Json, - invalid_rows, - retry, - held, - ) - .await + appender + .append(&ORIGIN, table, false, Schema::Json, batch) + .await } #[tokio::test] async fn given_a_table_holding_part_of_a_batch_should_append_only_the_messages_after_it() { let (catalog, table) = FlakyCatalog::with_table(&[]).await; - commit_up_to(&catalog, &table, &ORIGIN, 4).await; + commit_append(&catalog, &table, &ORIGIN, 0..=4).await; let other_partition = BatchOrigin { partition_id: 3, ..ORIGIN }; - commit_up_to(&catalog, &table, &other_partition, 9).await; + commit_append(&catalog, &table, &other_partition, 0..=9).await; let batch = messages(3..=5); let unheld = HeldMessages::default() - .unheld(&ORIGIN, &table, None, &catalog, &batch) + .unheld(&ORIGIN, &table, false, &catalog, &batch) .await .expect("The table should load"); @@ -724,19 +745,19 @@ pub(crate) mod tests { #[tokio::test] async fn given_a_backlog_batch_carrying_an_older_append_should_skip_up_to_the_newest() { let (catalog, table) = FlakyCatalog::with_table(&[]).await; - for offset in [1, 3, 5] { - commit_up_to(&catalog, &table, &ORIGIN, offset).await; + for offsets in [0..=1, 2..=3, 4..=5] { + commit_append(&catalog, &table, &ORIGIN, offsets).await; } let first_batch = messages(0..=1); let longer_batch = messages(0..=6); let held = HeldMessages::default(); let first = held - .unheld(&ORIGIN, &table, None, &catalog, &first_batch) + .unheld(&ORIGIN, &table, false, &catalog, &first_batch) .await .expect("The table should load"); let longer = held - .unheld(&ORIGIN, &table, None, &catalog, &longer_batch) + .unheld(&ORIGIN, &table, false, &catalog, &longer_batch) .await .expect("The table should load"); @@ -747,13 +768,13 @@ pub(crate) mod tests { #[tokio::test] async fn given_a_rolled_back_table_should_not_count_the_appends_left_behind() { let (catalog, table) = FlakyCatalog::with_table(&[]).await; - commit_up_to(&catalog, &table, &ORIGIN, 4).await; + commit_append(&catalog, &table, &ORIGIN, 0..=4).await; let kept = current_table(&catalog, &table) .await .metadata() .current_snapshot_id() .expect("The append should be the current snapshot"); - commit_up_to(&catalog, &table, &ORIGIN, 9).await; + commit_append(&catalog, &table, &ORIGIN, 5..=9).await; let current = current_table(&catalog, &table).await; let rolled_back = TableMetadataBuilder::new_from_metadata(current.metadata().clone(), None) .set_ref( @@ -774,7 +795,7 @@ pub(crate) mod tests { let batch = messages(3..=10); let unheld = HeldMessages::default() - .unheld(&ORIGIN, &rolled_back, Some(&rolled_back), &catalog, &batch) + .unheld(&ORIGIN, &rolled_back, true, &catalog, &batch) .await .expect("The current table needs no catalog"); @@ -786,11 +807,11 @@ pub(crate) mod tests { let (_, table) = FlakyCatalog::with_table(&[]).await; let unreachable = rest_catalog_answering(StatusCode::SERVICE_UNAVAILABLE).await; let held = HeldMessages::default(); - held.committed(&ORIGIN, &table, LastMessage::from(&message(4))); + held.committed(&ORIGIN, &table, append(0..=4)); let batch = messages(3..=5); let unheld = held - .unheld(&ORIGIN, &table, None, &unreachable, &batch) + .unheld(&ORIGIN, &table, false, &unreachable, &batch) .await .expect("The known commit should need no catalog"); @@ -802,17 +823,56 @@ pub(crate) mod tests { let (_, table) = FlakyCatalog::with_table(&[]).await; let unreachable = rest_catalog_answering(StatusCode::SERVICE_UNAVAILABLE).await; let held = HeldMessages::default(); - held.committed(&ORIGIN, &table, LastMessage::from(&message(4))); + held.committed(&ORIGIN, &table, append(0..=4)); held.forget(&ORIGIN, &table); let error = held - .unheld(&ORIGIN, &table, None, &unreachable, &messages(3..=5)) + .unheld(&ORIGIN, &table, false, &unreachable, &messages(3..=5)) .await .expect_err("A forgotten commit should be read from the catalog"); assert!(!error.is_permanent(), "{error}"); } + #[tokio::test] + async fn given_a_flushed_offset_should_forget_the_appends_up_to_it_in_that_partition() { + let (_, table) = FlakyCatalog::with_table(&[]).await; + let unreachable = rest_catalog_answering(StatusCode::SERVICE_UNAVAILABLE).await; + let other_partition = BatchOrigin { + partition_id: 3, + ..ORIGIN + }; + let held = HeldMessages::default(); + for origin in [&ORIGIN, &other_partition] { + held.committed(origin, &table, append(0..=4)); + held.committed(origin, &table, append(5..=9)); + } + let topic_metadata = TopicMetadata { + stream: ORIGIN.stream.to_owned(), + topic: ORIGIN.topic.to_owned(), + }; + + held.flushed( + &topic_metadata, + &[FlushPartition { + partition_id: ORIGIN.partition_id, + offset: 4, + }], + ); + + let batch = messages(3..=9); + let flushed = held + .unheld(&ORIGIN, &table, false, &unreachable, &batch) + .await + .expect("The partition should stay known"); + let other = held + .unheld(&other_partition, &table, false, &unreachable, &batch) + .await + .expect("The partition should stay known"); + assert_eq!(offsets(flushed), offsets(&batch)); + assert!(other.is_empty(), "{:?}", offsets(other)); + } + #[tokio::test] async fn given_a_catalog_that_cannot_load_the_table_should_fail_by_the_cause() { let (_, table) = FlakyCatalog::with_table(&[]).await; @@ -825,7 +885,7 @@ pub(crate) mod tests { let catalog = rest_catalog_answering(status).await; let error = HeldMessages::default() - .unheld(&ORIGIN, &table, None, &catalog, &messages(3..=5)) + .unheld(&ORIGIN, &table, false, &catalog, &messages(3..=5)) .await .expect_err("The table should not load"); @@ -833,71 +893,34 @@ pub(crate) mod tests { } } - #[tokio::test] - async fn given_a_batch_below_the_newest_append_when_recording_should_keep_the_newest() { - let (_, table) = FlakyCatalog::with_table(&[]).await; - let held = HeldMessages::default(); - held.committed(&ORIGIN, &table, LastMessage::from(&message(9))); - - let lower = held.last_to_record(&ORIGIN, &table, LastMessage::from(&message(3))); - let higher = held.last_to_record(&ORIGIN, &table, LastMessage::from(&message(12))); - - assert_eq!(lower, LastMessage::from(&message(9))); - assert_eq!(higher, LastMessage::from(&message(12))); - } - #[tokio::test] async fn given_a_commit_that_landed_but_failed_when_delivered_again_should_not_append_it_twice() { - let (catalog, table) = FlakyCatalog::with_table(&[CommitFailure::ResponseLost]).await; - let single_attempt = RetryPolicy { - max_attempts: 1, - ..FAST_RETRY - }; - let held = HeldMessages::default(); + let (appender, table) = + appender(&[CommitFailure::ResponseLost], InvalidRows::Fail, true).await; let batch = json_messages(0..=2); - write( - &catalog, - &table, - &batch, - InvalidRows::Fail, - single_attempt, - Some(&held), - ) - .await - .expect_err("The lost response should fail the commit"); - let again = write( - &catalog, - &table, - &batch, - InvalidRows::Fail, - single_attempt, - Some(&held), - ) - .await - .expect("The batch delivered again should succeed"); + write(&appender, &table, &batch) + .await + .expect_err("The lost response should fail the commit"); + let again = write(&appender, &table, &batch) + .await + .expect("The batch delivered again should succeed"); assert_eq!(again.committed_rows, 0, "the table already holds the batch"); - assert_eq!(catalog.snapshot_count(&table).await, 1); + let current = current_table(appender.catalog.as_ref(), &table).await; + assert_eq!(current.metadata().snapshots().count(), 1); } #[tokio::test] async fn given_a_row_that_does_not_match_when_failing_rows_should_commit_the_rows_before_it() { - let (catalog, table) = FlakyCatalog::with_table(&[]).await; + let (appender, table) = appender(&[], InvalidRows::Fail, true).await; let mut batch = json_messages(0..=3); batch[2].payload = Payload::Json(simd_json::json!({ "id": "two" })); - let written = write( - &catalog, - &table, - &batch, - InvalidRows::Fail, - FAST_RETRY, - Some(&HeldMessages::default()), - ) - .await - .expect("The rows before the rejected one should commit"); + let written = write(&appender, &table, &batch) + .await + .expect("The rows before the rejected one should commit"); assert_eq!(written.committed_rows, 2); let rejection = written.rejection.expect("The row should be rejected"); @@ -908,58 +931,52 @@ pub(crate) mod tests { "{}", rejection.error ); - assert_eq!(catalog.snapshot_count(&table).await, 1); + let current = current_table(appender.catalog.as_ref(), &table).await; + assert_eq!(current.metadata().snapshots().count(), 1); } #[tokio::test] - async fn given_a_payload_without_json_should_reject_it_or_skip_it_as_configured() { - let (catalog, table) = FlakyCatalog::with_table(&[]).await; + async fn given_a_payload_without_json_should_reject_it_only_when_failing_rows_at_least_once() { let mut batch = json_messages(0..=3); batch[1].payload = Payload::Proto("id: 1".to_owned()); - let failed = write( - &catalog, - &table, - &batch, - InvalidRows::Fail, - FAST_RETRY, - None, - ) - .await - .expect("The rows before the rejected one should commit"); - let skipped = write( - &catalog, - &table, - &batch, - InvalidRows::Skip, - FAST_RETRY, - None, - ) - .await - .expect("Skipping should write the other rows"); + let mut written = Vec::new(); + for (invalid_rows, at_least_once) in [ + (InvalidRows::Fail, true), + (InvalidRows::Fail, false), + (InvalidRows::Skip, true), + ] { + let (appender, table) = appender(&[], invalid_rows, at_least_once).await; + written.push( + write(&appender, &table, &batch) + .await + .expect("A payload without JSON is not a failed write"), + ); + } - assert_eq!(failed.committed_rows, 1); - assert_eq!(failed.rejection.map(|rejection| rejection.offset), Some(1)); - assert_eq!(skipped.committed_rows, 3); - assert!(skipped.rejection.is_none()); + assert_eq!(written[0].committed_rows, 1); + assert_eq!( + written[0] + .rejection + .as_ref() + .map(|rejection| rejection.offset), + Some(1) + ); + for skipped in &written[1..] { + assert_eq!(skipped.committed_rows, 3); + assert!(skipped.rejection.is_none()); + } } #[tokio::test] async fn given_at_most_once_should_record_no_delivered_messages_in_the_snapshot() { - let (catalog, table) = FlakyCatalog::with_table(&[]).await; + let (appender, table) = appender(&[], InvalidRows::Fail, false).await; - write( - &catalog, - &table, - &json_messages(0..=2), - InvalidRows::Fail, - FAST_RETRY, - None, - ) - .await - .expect("The batch should commit"); + write(&appender, &table, &json_messages(0..=2)) + .await + .expect("The batch should commit"); - let current = current_table(&catalog, &table).await; + let current = current_table(appender.catalog.as_ref(), &table).await; let snapshot = current .metadata() .current_snapshot() @@ -967,7 +984,9 @@ pub(crate) mod tests { let properties = &snapshot.summary().additional_properties; assert!(properties.contains_key(COMMIT_ID_PROPERTY)); assert!( - !properties.keys().any(|key| key.starts_with("iggy.last-")), + !properties + .keys() + .any(|key| key.starts_with("iggy.") && key != COMMIT_ID_PROPERTY), "{properties:?}" ); } @@ -976,20 +995,12 @@ pub(crate) mod tests { async fn given_an_empty_call_should_succeed_without_the_catalog() { let (_, table) = FlakyCatalog::with_table(&[]).await; let unreachable = rest_catalog_answering(StatusCode::SERVICE_UNAVAILABLE).await; + let mut appender = Appender::new(Box::new(unreachable), InvalidRows::Fail, FAST_RETRY); + appender.enable_at_least_once(); - let written = write_data( - &ORIGIN, - &[], - &table, - None, - &unreachable, - Schema::Json, - InvalidRows::Fail, - FAST_RETRY, - Some(&HeldMessages::default()), - ) - .await - .expect("An empty call has nothing to write"); + let written = write(&appender, &table, &[]) + .await + .expect("An empty call has nothing to write"); assert_eq!(written.committed_rows, 0); } @@ -1006,8 +1017,17 @@ pub(crate) mod tests { assert!(is_missing(&missing)); assert!(!is_missing(&forbidden)); - assert!(table_load_error("test.events", &missing).is_permanent()); - assert!(!table_load_error("test.events", &unreachable).is_permanent()); - assert!(table_load_error("test.events", &forbidden).is_permanent()); + assert!(matches!( + table_load_error("test.events", &missing), + Error::InvalidConfigValue(_) + )); + assert!(matches!( + table_load_error("test.events", &unreachable), + Error::Connection(_) + )); + assert!(matches!( + table_load_error("test.events", &forbidden), + Error::PermanentHttpError(_) + )); } } diff --git a/core/connectors/sinks/iceberg_sink/src/router/dynamic_router.rs b/core/connectors/sinks/iceberg_sink/src/router/dynamic_router.rs index 57f035cb5..c28e72d19 100644 --- a/core/connectors/sinks/iceberg_sink/src/router/dynamic_router.rs +++ b/core/connectors/sinks/iceberg_sink/src/router/dynamic_router.rs @@ -17,16 +17,13 @@ use crate::InvalidRows; use crate::router::commit::{is_missing, table_load_error}; -use crate::router::redelivery::HeldMessages; use crate::router::{ - BatchOrigin, Rejection, Router, before, first_rejection, format_error_chain, - is_valid_namespaced_table, load_table, write_data, + Appender, BatchOrigin, Rejection, Router, before, first_rejection, format_error_chain, + is_valid_namespaced_table, load_table, }; use async_trait::async_trait; -use iceberg::Catalog; use iceberg::table::Table; -use iggy_connector_sdk::retry::RetryPolicy; -use iggy_connector_sdk::{ConsumedMessage, Error, Payload, Schema}; +use iggy_connector_sdk::{ConsumedMessage, Error, FlushPartition, Payload, Schema, TopicMetadata}; use simd_json::base::ValueAsObject; use std::collections::HashMap; use std::collections::hash_map::Entry; @@ -34,26 +31,23 @@ use tracing::{info, warn}; #[derive(Debug)] pub struct DynamicRouter { - catalog: Box<dyn Catalog>, route_field: String, - invalid_rows: InvalidRows, - commit_retry: RetryPolicy, - held: Option<HeldMessages>, + appender: Appender, +} + +/// Where the messages carrying one route value go. +enum Destination { + /// The table at this index of the batch's tables. + Table(usize), + /// No table, for this reason. + Unroutable(String), } impl DynamicRouter { - pub fn new( - catalog: Box<dyn Catalog>, - route_field: String, - invalid_rows: InvalidRows, - commit_retry: RetryPolicy, - ) -> Self { + pub(crate) fn new(appender: Appender, route_field: String) -> Self { Self { - catalog, route_field, - invalid_rows, - commit_retry, - held: None, + appender, } } @@ -61,6 +55,53 @@ impl DynamicRouter { route_value(&message.payload, &self.route_field) } + /// Where the messages carrying `route_field_val` go. Route values that load the same table + /// share its entry in `tables`, so that each table gets its messages in one ascending run, + /// which is what tells the messages it holds from the ones it does not. + async fn destination( + &self, + route_field_val: &str, + tables: &mut Vec<(Table, Vec<ConsumedMessage>)>, + ) -> Result<Destination, Error> { + if !is_valid_namespaced_table(route_field_val) { + return Ok(Destination::Unroutable(format!( + "its route field value: {route_field_val} names no table. It needs at least one namespace before the table, separated by '.'" + ))); + } + let table = match load_table(route_field_val, self.appender.catalog.as_ref()).await { + Ok(table) => table, + Err(error) if is_missing(&error) => { + return Ok(Destination::Unroutable(format!( + "table: {route_field_val} does not exist in the catalog: {}", + format_error_chain(&error) + ))); + } + // An unreachable catalog or a refused load says nothing about the message, so at + // least once the batch fails without skipping it. + Err(error) if self.appender.held.is_some() => { + return Err(table_load_error(route_field_val, &error)); + } + Err(error) => { + return Ok(Destination::Unroutable(format!( + "table: {route_field_val} does not load from the catalog: {}", + format_error_chain(&error) + ))); + } + }; + let uuid = table.metadata().uuid(); + let index = match tables + .iter() + .position(|(loaded, _)| loaded.metadata().uuid() == uuid) + { + Some(index) => index, + None => { + tables.push((table, Vec::new())); + tables.len() - 1 + } + }; + Ok(Destination::Table(index)) + } + /// Drops a message that names no table it can be written to. At least once, it rejects the /// message instead, unless `invalid_rows` skips rows, so that the call stops at it. fn unroutable( @@ -69,17 +110,17 @@ impl DynamicRouter { message: &ConsumedMessage, reason: &str, ) -> Option<Rejection> { - if self.held.is_some() && self.invalid_rows == InvalidRows::Fail { + if self.appender.held.is_some() && self.appender.invalid_rows == InvalidRows::Fail { return Some(Rejection { offset: message.offset, error: Error::InvalidRecordValue(format!( - "The message at offset {} ({origin}) cannot be routed: {reason}", + "The message at offset: {} ({origin}) cannot be routed: {reason}", message.offset )), }); } warn!( - "Dropping the message at offset {} ({origin}): {reason}", + "Dropping the message at offset: {} ({origin}): {reason}", message.offset ); None @@ -106,92 +147,50 @@ impl Router for DynamicRouter { messages_schema: Schema, messages: Vec<ConsumedMessage>, ) -> Result<(), crate::Error> { - let mut tables: HashMap<String, (Table, Vec<ConsumedMessage>)> = HashMap::new(); - // Route values that name no table, with the reason, so that a batch loads each once. - let mut unroutable: HashMap<String, String> = HashMap::new(); + let mut tables: Vec<(Table, Vec<ConsumedMessage>)> = Vec::new(); + // Where each route value leads, so that a batch loads each table once. + let mut destinations: HashMap<String, Destination> = HashMap::new(); + let without_route_field = Destination::Unroutable(format!( + "it is not a JSON object with the route field: {}", + self.route_field + )); let mut rejection = None; for mut message in messages { // Normalised once here so the routing read and the later write both // work off the document, rather than parsing the same proto text twice. message.payload = message.payload.into_json_document(); - let Some(route_field_val) = self.extract_route_field(&message) else { - let reason = format!( - "it is not a JSON object with the route field: {}", - self.route_field - ); - rejection = self.unroutable(origin, &message, &reason); - if rejection.is_some() { - break; - } - continue; - }; - if let Some(reason) = unroutable.get(&route_field_val) { - rejection = self.unroutable(origin, &message, reason); - if rejection.is_some() { - break; - } - continue; - } - - let entry = match tables.entry(route_field_val) { - Entry::Occupied(mut entry) => { - entry.get_mut().1.push(message); - continue; - } - Entry::Vacant(entry) => entry, - }; - let route_field_val = entry.key(); - - let reason = if is_valid_namespaced_table(route_field_val) { - match load_table(route_field_val, self.catalog.as_ref()).await { - Ok(table) => { - entry.insert((table, vec![message])); - continue; + let destination: &Destination = match self.extract_route_field(&message) { + Some(route_field_val) => match destinations.entry(route_field_val) { + Entry::Occupied(entry) => entry.into_mut(), + Entry::Vacant(entry) => { + let destination = self.destination(entry.key(), &mut tables).await?; + entry.insert(destination) } - Err(error) if is_missing(&error) => format!( - "table: {route_field_val} does not exist in the catalog: {}", - format_error_chain(&error) - ), - // An unreachable catalog or a refused load says nothing about the message, - // so at least once the batch fails without skipping it. - Err(error) if self.held.is_some() => { - return Err(table_load_error(route_field_val, &error)); + }, + None => &without_route_field, + }; + match destination { + Destination::Table(index) => tables[*index].1.push(message), + Destination::Unroutable(reason) => { + rejection = self.unroutable(origin, &message, reason); + if rejection.is_some() { + break; } - Err(error) => format!( - "table: {route_field_val} does not load from the catalog: {}", - format_error_chain(&error) - ), } - } else { - format!( - "its route field value: {route_field_val} names no table. It needs at least one namespace before the table, separated by '.'" - ) - }; - rejection = self.unroutable(origin, &message, &reason); - if rejection.is_some() { - break; } - unroutable.insert(entry.into_key(), reason); } - for (table_name, (table, batch_messages)) in tables { + for (table, table_messages) in &tables { // Once a message is rejected, the tables still need every message before it. - let messages = before(&batch_messages, rejection.as_ref()); - let written = write_data( - origin, - messages, - &table, - Some(&table), - self.catalog.as_ref(), - messages_schema, - self.invalid_rows, - self.commit_retry, - self.held.as_ref(), - ) - .await?; + let messages = before(table_messages, rejection.as_ref()); + let written = self + .appender + .append(origin, table, true, messages_schema, messages) + .await?; info!( - "Dynamically routed {} input rows to iceberg table {table_name}, committed {}", + "Dynamically routed {} input rows to iceberg table {}, committed {}", messages.len(), + table.identifier(), written.committed_rows ); rejection = first_rejection(rejection, written.rejection); @@ -204,18 +203,22 @@ impl Router for DynamicRouter { } fn enable_at_least_once(&mut self) -> Result<(), Error> { - self.held = Some(HeldMessages::default()); + self.appender.enable_at_least_once(); Ok(()) } + + fn flushed(&self, topic_metadata: &TopicMetadata, partitions: &[FlushPartition]) { + self.appender.flushed(topic_metadata, partitions); + } } #[cfg(test)] mod tests { use super::*; - use iceberg::memory::{MEMORY_CATALOG_WAREHOUSE, MemoryCatalogBuilder}; - use iceberg::spec::{NestedField, PrimitiveType, Schema as IcebergSchema, Type}; - use iceberg::{CatalogBuilder, NamespaceIdent, TableCreation, TableIdent}; - use std::time::Duration; + use crate::router::commit::tests::{ + SINGLE_ATTEMPT, catalog_with_table, rest_catalog_answering, + }; + use iceberg::TableIdent; const ORIGIN: BatchOrigin<'static> = BatchOrigin { connector_id: 1, @@ -225,44 +228,11 @@ mod tests { }; async fn router(invalid_rows: InvalidRows, at_least_once: bool) -> DynamicRouter { - let catalog = MemoryCatalogBuilder::default() - .load( - "memory", - HashMap::from([( - MEMORY_CATALOG_WAREHOUSE.to_owned(), - "memory://warehouse".to_owned(), - )]), - ) - .await - .expect("Failed to create catalog"); - let namespace = NamespaceIdent::new("test".to_owned()); - catalog - .create_namespace(&namespace, HashMap::new()) - .await - .expect("Failed to create namespace"); - let schema = IcebergSchema::builder() - .with_fields(vec![ - NestedField::required(1, "id", Type::Primitive(PrimitiveType::Long)).into(), - ]) - .build() - .expect("Failed to build schema"); - catalog - .create_table( - &namespace, - TableCreation::builder() - .name("events".to_owned()) - .schema(schema) - .build(), - ) - .await - .expect("Failed to create table"); - let retry = RetryPolicy { - max_attempts: 1, - base_delay: Duration::from_millis(1), - max_delay: Duration::from_millis(1), - }; - let mut router = - DynamicRouter::new(Box::new(catalog), "table".to_owned(), invalid_rows, retry); + let (catalog, _) = catalog_with_table(HashMap::new()).await; + let mut router = DynamicRouter::new( + Appender::new(Box::new(catalog), invalid_rows, SINGLE_ATTEMPT), + "table".to_owned(), + ); if at_least_once { router .enable_at_least_once() @@ -303,7 +273,7 @@ mod tests { .expect_err("An unroutable message should fail the batch"); assert!(error.is_permanent(), "{error}"); - assert!(error.to_string().contains("offset 7"), "{error}"); + assert!(error.to_string().contains("offset: 7"), "{error}"); } } @@ -350,6 +320,15 @@ mod tests { assert_eq!(route_value(&payload, "region"), None); } + async fn events(router: &DynamicRouter) -> Table { + router + .appender + .catalog + .load_table(&TableIdent::from_strs(["test", "events"]).expect("a valid table name")) + .await + .expect("Failed to load table") + } + #[tokio::test] async fn given_at_least_once_when_a_message_cannot_be_routed_should_commit_the_messages_before_it() { @@ -365,12 +344,8 @@ mod tests { .await .expect_err("The unroutable message should stop the batch"); - assert!(error.to_string().contains("offset 7"), "{error}"); - let table = router - .catalog - .load_table(&TableIdent::from_strs(["test", "events"]).expect("a valid table name")) - .await - .expect("Failed to load table"); + assert!(error.to_string().contains("offset: 7"), "{error}"); + let table = events(&router).await; let snapshot = table .metadata() .current_snapshot() @@ -384,20 +359,61 @@ mod tests { ); } + #[tokio::test] + async fn given_route_values_that_load_one_table_should_append_their_messages_together() { + let router = router(InvalidRows::Fail, true).await; + let table = events(&router).await; + let alias = TableIdent::from_strs(["test", "alias"]).expect("a valid table name"); + router + .appender + .catalog + .register_table( + &alias, + table + .metadata_location() + .expect("A created table has a metadata location") + .to_owned(), + ) + .await + .expect("Failed to register the alias"); + let batch = (0..4) + .map(|offset| ConsumedMessage { + id: u128::from(offset), + offset, + ..message(simd_json::json!({ + "table": if offset % 2 == 0 { "test.events" } else { "test.alias" }, + "id": offset, + })) + }) + .collect(); + + router + .route_data(&ORIGIN, Schema::Json, batch) + .await + .expect("Both route values should load the table"); + + let table = events(&router).await; + let snapshot = table + .metadata() + .current_snapshot() + .expect("The batch should be appended"); + assert_eq!( + snapshot + .summary() + .additional_properties + .get("added-records"), + Some(&"4".to_owned()), + "one append should hold the messages of both route values" + ); + } + #[tokio::test] async fn given_a_catalog_that_refuses_the_load_when_skipping_rows_should_fail_without_skipping() { - let catalog = - crate::router::commit::tests::rest_catalog_answering(http::StatusCode::FORBIDDEN).await; + let catalog = rest_catalog_answering(http::StatusCode::FORBIDDEN).await; let mut router = DynamicRouter::new( - Box::new(catalog), + Appender::new(Box::new(catalog), InvalidRows::Skip, SINGLE_ATTEMPT), "table".to_owned(), - InvalidRows::Skip, - RetryPolicy { - max_attempts: 1, - base_delay: Duration::from_millis(1), - max_delay: Duration::from_millis(1), - }, ); router .enable_at_least_once() @@ -410,6 +426,6 @@ mod tests { .expect_err("A refused load says nothing about the message, so it is not skipped"); assert!(error.is_permanent(), "{error}"); - assert!(matches!(error, Error::InvalidConfigValue(_)), "{error}"); + assert!(matches!(error, Error::PermanentHttpError(_)), "{error}"); } } diff --git a/core/connectors/sinks/iceberg_sink/src/router/mod.rs b/core/connectors/sinks/iceberg_sink/src/router/mod.rs index de3865f08..2c729e2af 100644 --- a/core/connectors/sinks/iceberg_sink/src/router/mod.rs +++ b/core/connectors/sinks/iceberg_sink/src/router/mod.rs @@ -18,7 +18,7 @@ use crate::InvalidRows; use crate::router::arrow_streamer::JsonArrowReader; use crate::router::commit::commit_data_files; -use crate::router::redelivery::{HeldMessages, LastMessage, snapshot_properties}; +use crate::router::redelivery::{Append, HeldMessages, snapshot_properties}; use crate::slice_user_table; use arrow_array::RecordBatch; use arrow_json::ReaderBuilder; @@ -38,7 +38,7 @@ use iceberg::{ writer::file_writer::location_generator::{DefaultFileNameGenerator, DefaultLocationGenerator}, }; use iggy_connector_sdk::retry::RetryPolicy; -use iggy_connector_sdk::{ConsumedMessage, Error, Schema}; +use iggy_connector_sdk::{ConsumedMessage, Error, FlushPartition, Schema, TopicMetadata}; use parquet::file::properties::WriterProperties; use std::borrow::Cow; use std::collections::HashMap; @@ -139,73 +139,131 @@ fn before<'m>( } } -/// Writes and commits `messages` to `table`. With `held`, the sink delivers at least once: the -/// messages that `table` already holds are not appended again, and the commit records the last -/// message it holds. `current` is a copy of `table` just loaded from the catalog. -#[allow(clippy::too_many_arguments)] -async fn write_data( - origin: &BatchOrigin<'_>, - messages: &[ConsumedMessage], - table: &Table, - current: Option<&Table>, - catalog: &dyn Catalog, - messages_schema: Schema, +/// What both routers need to append messages to a table. +#[derive(Debug)] +pub(crate) struct Appender { + catalog: Box<dyn Catalog>, invalid_rows: InvalidRows, commit_retry: RetryPolicy, - held: Option<&HeldMessages>, -) -> Result<Written, Error> { - if messages.is_empty() { - return Ok(Written::default()); - } - let messages = match held { - Some(held) => { - held.unheld(origin, table, current, catalog, messages) - .await? + /// Set when the sink delivers at least once. + held: Option<HeldMessages>, +} + +impl Appender { + pub(crate) fn new( + catalog: Box<dyn Catalog>, + invalid_rows: InvalidRows, + commit_retry: RetryPolicy, + ) -> Self { + Self { + catalog, + invalid_rows, + commit_retry, + held: None, } - None => messages, - }; - if messages.is_empty() { - return Ok(Written::default()); - } - let (data_files, rejection) = - write_data_files(origin, messages, table, messages_schema, invalid_rows).await?; - let handled = before(messages, rejection.as_ref()); - let (Some(last), false) = (handled.last(), data_files.is_empty()) else { - // Every row was skipped or rejected, and an empty append would still add a snapshot. - return Ok(Written { - committed_rows: 0, - rejection, - }); - }; - let committed_rows = data_files.iter().map(DataFile::record_count).sum(); - let recorded = held.map(|held| held.last_to_record(origin, table, LastMessage::from(last))); - let properties = recorded - .map(|last| snapshot_properties(origin, last)) - .unwrap_or_default(); - let committed = commit_data_files(table, catalog, data_files, properties, commit_retry).await; - if let (Some(held), Some(recorded)) = (held, recorded) { - match &committed { - Ok(()) => held.committed(origin, table, recorded), - Err(_) => held.forget(origin, table), + } + + fn enable_at_least_once(&mut self) { + self.held = Some(HeldMessages::default()); + } + + fn flushed(&self, topic_metadata: &TopicMetadata, partitions: &[FlushPartition]) { + if let Some(held) = &self.held { + held.flushed(topic_metadata, partitions); + } + } + + /// Writes and commits `messages` to `table`. `fresh` says that `table` was just loaded from + /// the catalog. At least once, the messages that `table` already holds are not appended + /// again, and the commit records the messages it appends. + async fn append( + &self, + origin: &BatchOrigin<'_>, + table: &Table, + fresh: bool, + messages_schema: Schema, + messages: &[ConsumedMessage], + ) -> Result<Written, Error> { + if messages.is_empty() { + return Ok(Written::default()); + } + let messages = match &self.held { + Some(held) => { + held.unheld(origin, table, fresh, self.catalog.as_ref(), messages) + .await? + } + None => messages, + }; + if messages.is_empty() { + return Ok(Written::default()); } + // At most once, a failed call loses the rest of the run, so a payload without a JSON + // document is skipped whatever `invalid_rows` says. + let without_json = if self.held.is_some() { + self.invalid_rows + } else { + InvalidRows::Skip + }; + let (data_files, rejection) = write_data_files( + origin, + messages, + table, + messages_schema, + self.invalid_rows, + without_json, + ) + .await?; + let handled = before(messages, rejection.as_ref()); + let (Some(first), Some(last), false) = + (handled.first(), handled.last(), data_files.is_empty()) + else { + // Every row was skipped or rejected, and an empty append would still add a snapshot. + return Ok(Written { + committed_rows: 0, + rejection, + }); + }; + let committed_rows = data_files.iter().map(DataFile::record_count).sum(); + let append = Append::new(first, last); + let properties = if self.held.is_some() { + snapshot_properties(origin, append) + } else { + HashMap::new() + }; + let committed = commit_data_files( + table, + self.catalog.as_ref(), + data_files, + properties, + self.commit_retry, + ) + .await; + if let Some(held) = &self.held { + match &committed { + Ok(()) => held.committed(origin, table, append), + Err(_) => held.forget(origin, table), + } + } + committed?; + Ok(Written { + committed_rows, + rejection, + }) } - committed?; - Ok(Written { - committed_rows, - rejection, - }) } /// Writes the JSON payloads as Parquet data files under the table location, /// laid out by the table's default partition spec. Nothing is committed here. With /// `invalid_rows = "fail"`, the first message that the table cannot take ends the write, and the -/// data files hold the messages before it. +/// data files hold the messages before it. `without_json` decides the same for a payload that +/// holds no JSON document. async fn write_data_files( origin: &BatchOrigin<'_>, messages: &[ConsumedMessage], table: &Table, messages_schema: Schema, invalid_rows: InvalidRows, + without_json: InvalidRows, ) -> Result<(Vec<DataFile>, Option<Rejection>), Error> { let location = DefaultLocationGenerator::new(table.metadata()).map_err(|err| { error!( @@ -252,7 +310,7 @@ async fn write_data_files( message.offset, table.identifier() ); - if invalid_rows == InvalidRows::Skip { + if without_json == InvalidRows::Skip { warn!("Skipping a row. {reason}"); continue; } @@ -543,6 +601,10 @@ pub trait Router: std::fmt::Debug + Sync + Send { /// does not load, and appends only the messages a table does not hold yet when a batch is /// delivered again. Refuses when the configuration cannot keep that promise. fn enable_at_least_once(&mut self) -> Result<(), Error>; + + /// A flush proved `partitions` delivered up to their offsets, so the runtime never hands those + /// messages over again. + fn flushed(&self, topic_metadata: &TopicMetadata, partitions: &[FlushPartition]); } #[cfg(test)] @@ -642,6 +704,7 @@ mod tests { table, Schema::Json, invalid_rows, + invalid_rows, )) } @@ -700,6 +763,7 @@ mod tests { &table, Schema::Proto, InvalidRows::Fail, + InvalidRows::Fail, )) .expect("a rejected message is not a failed write"); @@ -789,6 +853,7 @@ mod tests { &table, Schema::Text, InvalidRows::Skip, + InvalidRows::Skip, )) .expect("skipping every row is not an error"); diff --git a/core/connectors/sinks/iceberg_sink/src/router/redelivery.rs b/core/connectors/sinks/iceberg_sink/src/router/redelivery.rs index f90618ea6..4a2246cd7 100644 --- a/core/connectors/sinks/iceberg_sink/src/router/redelivery.rs +++ b/core/connectors/sinks/iceberg_sink/src/router/redelivery.rs @@ -18,8 +18,8 @@ //! Which messages of a partition a table already holds. At least once, the runtime hands a //! message over again when it cannot prove it delivered, for example after a crash between the //! commit and the offset store, or after a batch committed to one table and failed on another. -//! Every append records the last message it holds in its snapshot summary, and a delivery of -//! the same messages again appends only the ones after the newest append. +//! Every append records in its snapshot summary the first offset and the last message it holds, +//! and a delivery of the same messages again skips the ones that proven appends cover. use crate::router::BatchOrigin; use crate::router::commit::table_load_error; @@ -28,7 +28,7 @@ use iceberg::spec::Snapshot; use iceberg::table::Table; use iceberg::util::snapshot::ancestors_of; use iceberg::{Catalog, TableIdent}; -use iggy_connector_sdk::{ConsumedMessage, Error}; +use iggy_connector_sdk::{ConsumedMessage, Error, FlushPartition, TopicMetadata}; use std::collections::{BTreeMap, HashMap}; use tracing::{info, warn}; use uuid::Uuid; @@ -36,15 +36,45 @@ use uuid::Uuid; const STREAM_PROPERTY: &str = "iggy.stream"; const TOPIC_PROPERTY: &str = "iggy.topic"; const PARTITION_ID_PROPERTY: &str = "iggy.partition-id"; +const FIRST_OFFSET_PROPERTY: &str = "iggy.first-offset"; const LAST_OFFSET_PROPERTY: &str = "iggy.last-offset"; const LAST_MESSAGE_ID_PROPERTY: &str = "iggy.last-message-id"; const LAST_TIMESTAMP_PROPERTY: &str = "iggy.last-timestamp"; +/// The messages of a partition that one append holds: every message from offset `first` to +/// `last`, except the ones the sink skipped. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct Append { + first: u64, + last: LastMessage, +} + +impl Append { + /// The append of the messages from `first` to `last`. + pub(crate) fn new(first: &ConsumedMessage, last: &ConsumedMessage) -> Self { + Self { + first: first.offset, + last: LastMessage::from(last), + } + } + + fn from_properties(properties: &HashMap<String, String>) -> Option<Self> { + Some(Self { + first: properties.get(FIRST_OFFSET_PROPERTY)?.parse().ok()?, + last: LastMessage::from_properties(properties)?, + }) + } + + fn offsets(&self) -> (u64, u64) { + (self.first, self.last.offset) + } +} + /// The last message of a partition that an append holds. The message ID and the time the /// server stored the message tell it apart from another message at the same offset in a later /// history of the partition. #[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(crate) struct LastMessage { +struct LastMessage { offset: u64, id: u128, timestamp: u64, @@ -70,11 +100,10 @@ impl LastMessage { } } -/// The snapshot summary properties that record `last` as the last message of `origin` the -/// append holds. +/// The snapshot summary properties that record `append` as the messages of `origin` it holds. pub(crate) fn snapshot_properties( origin: &BatchOrigin<'_>, - last: LastMessage, + append: Append, ) -> HashMap<String, String> { HashMap::from([ (STREAM_PROPERTY.to_owned(), origin.stream.to_owned()), @@ -83,17 +112,25 @@ pub(crate) fn snapshot_properties( PARTITION_ID_PROPERTY.to_owned(), origin.partition_id.to_string(), ), - (LAST_OFFSET_PROPERTY.to_owned(), last.offset.to_string()), - (LAST_MESSAGE_ID_PROPERTY.to_owned(), last.id.to_string()), + (FIRST_OFFSET_PROPERTY.to_owned(), append.first.to_string()), + ( + LAST_OFFSET_PROPERTY.to_owned(), + append.last.offset.to_string(), + ), + ( + LAST_MESSAGE_ID_PROPERTY.to_owned(), + append.last.id.to_string(), + ), ( LAST_TIMESTAMP_PROPERTY.to_owned(), - last.timestamp.to_string(), + append.last.timestamp.to_string(), ), ]) } /// A partition delivered to a table. The table UUID tells a table recreated under the same -/// name apart. +/// name apart, once a router loads the table again. The static router keeps the table it loaded +/// first, so a recreated table needs a restart of the sink. #[derive(Debug, Clone, PartialEq, Eq, Hash)] struct TablePartition { table: Uuid, @@ -116,16 +153,17 @@ impl TablePartition { /// The appends of one partition that a table holds. #[derive(Debug, Default)] struct Appends { - /// The last message of each append, by offset. Where two histories of the partition - /// appended at the same offset, the newer one is kept. - last_messages: BTreeMap<u64, LastMessage>, - newest: Option<LastMessage>, + /// Each append by the offset of its last message. Where two appends end at the same offset, + /// as two histories of the partition can, the newer one is kept. + by_last: BTreeMap<u64, Append>, + /// The offset of the last message of the append committed last. + newest: Option<u64>, } impl Appends { - fn record(&mut self, last: LastMessage) { - self.last_messages.insert(last.offset, last); - self.newest = Some(last); + fn record(&mut self, append: Append) { + self.by_last.insert(append.last.offset, append); + self.newest = Some(append.last.offset); } } @@ -138,73 +176,71 @@ pub(crate) struct HeldMessages { } impl HeldMessages { - /// The tail of `messages` that `table` does not hold yet. `current` is a copy of `table` - /// just loaded from the catalog, which saves loading it again. + /// The tail of `messages` that `table` does not hold yet. `fresh` says that `table` was just + /// loaded from the catalog, which saves loading it again. pub(crate) async fn unheld<'m>( &self, origin: &BatchOrigin<'_>, table: &Table, - current: Option<&Table>, + fresh: bool, catalog: &dyn Catalog, messages: &'m [ConsumedMessage], ) -> Result<&'m [ConsumedMessage], Error> { let key = TablePartition::new(origin, table); - if !self.known.contains_key(&key) { - let appends = match current { - Some(current) => appends_in(current, origin), - None => { - let identifier = table.identifier(); - let current = catalog - .load_table(identifier) - .await - .map_err(|error| table_load_error(&identifier.to_string(), &error))?; - appends_in(¤t, origin) - } - }; - self.known.entry(key.clone()).or_insert(appends); - } - let held = match self.known.get_mut(&key) { - Some(mut appends) => { - match held_prefix(origin, table.identifier(), messages, &appends) { - Prefix::Held(held) => held, - Prefix::NewHistory => { - *appends = Appends::default(); - 0 - } - Prefix::Unproven => 0, - } + let loaded = if self.known.contains_key(&key) { + None + } else if fresh { + Some(appends_in(table, origin)) + } else { + let identifier = table.identifier(); + let current = catalog + .load_table(identifier) + .await + .map_err(|error| table_load_error(&identifier.to_string(), &error))?; + Some(appends_in(¤t, origin)) + }; + let mut appends = self + .known + .entry(key) + .or_insert_with(|| loaded.unwrap_or_default()); + let held = match held_prefix(origin, table.identifier(), messages, &appends) { + Prefix::Held(held) => held, + Prefix::NewHistory => { + *appends = Appends::default(); + 0 } - None => 0, + Prefix::Unproven => 0, }; Ok(&messages[held..]) } - /// The last message of a commit to record in its snapshot. A batch that the table could not - /// prove it held, but that ends below the newest append, records that append again, so that - /// the record never moves back within one history of the partition. - pub(crate) fn last_to_record( - &self, - origin: &BatchOrigin<'_>, - table: &Table, - last: LastMessage, - ) -> LastMessage { - self.known - .get(&TablePartition::new(origin, table)) - .and_then(|appends| appends.newest) - .filter(|newest| newest.offset > last.offset) - .unwrap_or(last) - } - - pub(crate) fn committed(&self, origin: &BatchOrigin<'_>, table: &Table, last: LastMessage) { + pub(crate) fn committed(&self, origin: &BatchOrigin<'_>, table: &Table, append: Append) { self.known .entry(TablePartition::new(origin, table)) .or_default() - .record(last); + .record(append); } pub(crate) fn forget(&self, origin: &BatchOrigin<'_>, table: &Table) { self.known.remove(&TablePartition::new(origin, table)); } + + /// Forgets the appends that end at or before the offset each of `partitions` flushed. The + /// runtime never hands those messages over again, and a restart reads the table anew. + pub(crate) fn flushed(&self, topic_metadata: &TopicMetadata, partitions: &[FlushPartition]) { + for mut known in self.known.iter_mut() { + let (key, appends) = known.pair_mut(); + if key.stream != topic_metadata.stream || key.topic != topic_metadata.topic { + continue; + } + if let Some(flushed) = partitions + .iter() + .find(|partition| partition.partition_id == key.partition_id) + { + appends.by_last.retain(|&last, _| last > flushed.offset); + } + } + } } /// The appends of `origin` on the current snapshot's ancestry. Snapshots that a rollback or a @@ -220,17 +256,17 @@ fn appends_in(table: &Table, origin: &BatchOrigin<'_>) -> Appends { continue; } let properties = &snapshot.summary().additional_properties; - let Some(last) = LastMessage::from_properties(properties) else { + let Some(append) = Append::from_properties(properties) else { warn!( - "Iceberg table {} has snapshot {} of {origin} without a readable last message. Skipping it", + "Iceberg table {} has snapshot {} of {origin} without readable offsets. Skipping it", table.identifier(), snapshot.snapshot_id() ); continue; }; // The walk runs from the newest snapshot back, so the first one seen at an offset wins. - appends.last_messages.entry(last.offset).or_insert(last); - appends.newest.get_or_insert(last); + appends.by_last.entry(append.last.offset).or_insert(append); + appends.newest.get_or_insert(append.last.offset); } appends } @@ -249,59 +285,120 @@ enum Prefix { Held(usize), /// None: the table holds other messages at the batch's offsets. NewHistory, - /// None that can be proven, although the batch reaches offsets the table holds. + /// None that can be proven, although appends cover some of the batch's offsets. Unproven, } -/// What the table holds of the start of `messages`. Messages count as held only when the batch -/// carries the last message of one of the table's appends, which proves that it is the history -/// the table holds. Then every message up to the newest append is held. +/// What the table holds of the start of `messages`. An append counts only when the batch carries +/// its last message, which proves that the batch is the history the append holds. Once one does, +/// the appends after the end of the batch count too, up to the newest one, since a batch cut at +/// other offsets than before carries none of their last messages. The messages held are the +/// start of the batch that the counted appends cover without a gap, so a message that no append +/// holds, for example one an at-most-once run lost, is appended again. fn held_prefix( origin: &BatchOrigin<'_>, table: &TableIdent, messages: &[ConsumedMessage], appends: &Appends, ) -> Prefix { - let Some(newest) = appends.newest else { + let (Some(first), Some(last)) = (messages.first(), messages.last()) else { return Prefix::Held(0); }; - // A run of one partition carries ascending offsets. - let held = messages.partition_point(|message| message.offset <= newest.offset); - if held == 0 { - return Prefix::Held(0); - } + let mut proven = Vec::new(); let mut other_history = false; - for message in &messages[..held] { - match appends.last_messages.get(&message.offset) { - Some(last) if *last == LastMessage::from(message) => { - info!( - "Skipping {held} messages from offset: {} to offset: {} ({origin}) that Iceberg table {table} already holds", - messages[0].offset, - messages[held - 1].offset, - ); - return Prefix::Held(held); + for message in messages { + match appends.by_last.get(&message.offset) { + Some(append) if append.last == LastMessage::from(message) => { + proven.push(append.offsets()); } Some(_) => other_history = true, None => {} } } + if proven.is_empty() { + return prefix_without_proof( + origin, + table, + first.offset, + last.offset, + appends, + other_history, + ); + } + if let Some(newest) = appends.newest + && newest > last.offset + { + proven.extend( + appends + .by_last + .range(last.offset + 1..=newest) + .map(|(_, append)| append.offsets()), + ); + } + proven.sort_unstable(); + let held = covered_prefix(messages, &proven); + if held > 0 { + info!( + "Skipping {held} messages from offset: {} to offset: {} ({origin}) that Iceberg table {table} already holds", + first.offset, + messages[held - 1].offset, + ); + } + Prefix::Held(held) +} + +/// What the table holds of a batch from offset `first` to `last` that carries the last message +/// of no append. +fn prefix_without_proof( + origin: &BatchOrigin<'_>, + table: &TableIdent, + first: u64, + last: u64, + appends: &Appends, + other_history: bool, +) -> Prefix { if other_history { warn!( "Iceberg table {table} holds other messages at the offsets of the batch ({origin}), so the partition has a new history. Appending every message of the batch" ); return Prefix::NewHistory; } + let covered = appends + .by_last + .range(first..) + .any(|(_, append)| append.first <= last); + if !covered { + return Prefix::Held(0); + } warn!( - "Iceberg table {table} holds messages up to offset: {} ({origin}), but the batch carries no message that proves which ones. Appending every message of the batch, which can append some twice", - newest.offset + "Iceberg table {table} holds messages between offset: {first} and offset: {last} ({origin}), but the batch carries no message that proves which ones. Appending every message of the batch, which can append some twice" ); Prefix::Unproven } +/// How many messages from the start of `messages` the offset ranges cover without a gap. The +/// ranges are sorted by their first offset. +fn covered_prefix(messages: &[ConsumedMessage], ranges: &[(u64, u64)]) -> usize { + let mut next = 0; + let mut covered_to = None; + messages + .iter() + .take_while(|message| { + while let Some(&(first, last)) = ranges.get(next) + && first <= message.offset + { + covered_to = covered_to.max(Some(last)); + next += 1; + } + covered_to.is_some_and(|to| message.offset <= to) + }) + .count() +} + #[cfg(test)] mod tests { use super::*; - use iggy_connector_sdk::Payload; + use crate::router::commit::tests::{message, messages}; const ORIGIN: BatchOrigin<'static> = BatchOrigin { connector_id: 1, @@ -314,26 +411,11 @@ mod tests { TableIdent::from_strs(["ns", "table"]).expect("a valid table name") } - fn message(offset: u64) -> ConsumedMessage { - ConsumedMessage { - id: u128::from(offset) + 100, - offset, - checksum: 0, - timestamp: offset + 1_000, - origin_timestamp: 0, - headers: None, - payload: Payload::Raw(Vec::new()), - } - } - - fn messages(offsets: std::ops::RangeInclusive<u64>) -> Vec<ConsumedMessage> { - offsets.map(message).collect() - } - - fn appends(last_offsets: &[u64]) -> Appends { + /// Appends of the messages from each first to each last offset, committed in this order. + fn appends(offsets: &[(u64, u64)]) -> Appends { let mut appends = Appends::default(); - for &offset in last_offsets { - appends.record(LastMessage::from(&message(offset))); + for &(first, last) in offsets { + appends.record(Append::new(&message(first), &message(last))); } appends } @@ -347,14 +429,55 @@ mod tests { #[test] fn given_a_batch_carrying_the_newest_append_should_hold_up_to_it() { - let prefix = held_prefix(&ORIGIN, &table(), &messages(3..=5), &appends(&[4])); + let prefix = held_prefix(&ORIGIN, &table(), &messages(3..=5), &appends(&[(0, 4)])); assert_eq!(prefix, Prefix::Held(2)); } #[test] - fn given_a_backlog_batch_carrying_an_older_append_should_hold_up_to_the_newest() { - let prefix = held_prefix(&ORIGIN, &table(), &messages(0..=1), &appends(&[1, 3, 5])); + fn given_a_backlog_batch_carrying_an_older_append_should_hold_the_appends_after_it() { + let prefix = held_prefix( + &ORIGIN, + &table(), + &messages(0..=2), + &appends(&[(0, 1), (2, 3), (4, 5)]), + ); + + assert_eq!(prefix, Prefix::Held(3)); + } + + #[test] + fn given_a_gap_between_appends_should_hold_only_the_messages_before_it() { + let prefix = held_prefix( + &ORIGIN, + &table(), + &messages(0..=7), + &appends(&[(0, 1), (6, 7)]), + ); + + assert_eq!(prefix, Prefix::Held(2)); + } + + #[test] + fn given_a_batch_ending_after_the_newest_append_should_not_count_older_appends_after_it() { + let prefix = held_prefix( + &ORIGIN, + &table(), + &messages(5..=12), + &appends(&[(10, 19), (0, 9)]), + ); + + assert_eq!(prefix, Prefix::Held(5)); + } + + #[test] + fn given_a_batch_ending_before_the_newest_append_should_not_count_older_appends_after_it() { + let prefix = held_prefix( + &ORIGIN, + &table(), + &messages(3..=7), + &appends(&[(5, 19), (0, 4), (9, 10)]), + ); assert_eq!(prefix, Prefix::Held(2)); } @@ -364,21 +487,26 @@ mod tests { let mut batch = messages(3..=4); batch[1].timestamp += 1; - let prefix = held_prefix(&ORIGIN, &table(), &batch, &appends(&[4])); + let prefix = held_prefix(&ORIGIN, &table(), &batch, &appends(&[(0, 4)])); assert_eq!(prefix, Prefix::NewHistory); } #[test] fn given_a_batch_after_the_newest_append_should_append_every_message() { - let prefix = held_prefix(&ORIGIN, &table(), &messages(5..=6), &appends(&[4])); + let prefix = held_prefix(&ORIGIN, &table(), &messages(5..=6), &appends(&[(0, 4)])); assert_eq!(prefix, Prefix::Held(0)); } #[test] fn given_a_batch_without_an_appended_message_should_append_every_message() { - let prefix = held_prefix(&ORIGIN, &table(), &messages(2..=3), &appends(&[1, 4])); + let prefix = held_prefix( + &ORIGIN, + &table(), + &messages(2..=3), + &appends(&[(0, 1), (2, 4)]), + ); assert_eq!(prefix, Prefix::Unproven); } diff --git a/core/connectors/sinks/iceberg_sink/src/router/static_router.rs b/core/connectors/sinks/iceberg_sink/src/router/static_router.rs index 7d3d8cc58..73f8a39ff 100644 --- a/core/connectors/sinks/iceberg_sink/src/router/static_router.rs +++ b/core/connectors/sinks/iceberg_sink/src/router/static_router.rs @@ -15,39 +15,22 @@ // specific language governing permissions and limitations // under the License. -use crate::InvalidRows; use crate::router::commit::table_load_error; -use crate::router::redelivery::HeldMessages; use crate::router::{ - BatchOrigin, Router, before, first_rejection, format_error_chain, is_valid_namespaced_table, - load_table, write_data, + Appender, BatchOrigin, Router, before, first_rejection, format_error_chain, + is_valid_namespaced_table, load_table, }; use async_trait::async_trait; -use iceberg::Catalog; use iceberg::table::Table; -use iggy_connector_sdk::retry::RetryPolicy; -use iggy_connector_sdk::{ConsumedMessage, Error, Schema}; -use std::sync::{Arc, Mutex, MutexGuard, PoisonError}; +use iggy_connector_sdk::{ConsumedMessage, Error, FlushPartition, Schema, TopicMetadata}; +use std::sync::OnceLock; use tracing::{error, info, warn}; -/// A table that the configuration declares. A table that does not load is loaded again by the -/// next batch, so that a catalog that was unreachable at startup does not drop it for good. +/// A table that the configuration declares, set once it loads. #[derive(Debug)] struct DeclaredTable { name: String, - table: Mutex<Option<Arc<Table>>>, -} - -impl DeclaredTable { - fn loaded(&self) -> Option<Arc<Table>> { - self.lock().clone() - } - - fn lock(&self) -> MutexGuard<'_, Option<Arc<Table>>> { - // The guarded value is replaced whole, so a panic while the lock was held cannot leave it - // half written. - self.table.lock().unwrap_or_else(PoisonError::into_inner) - } + table: OnceLock<Table>, } #[derive(Debug)] @@ -55,19 +38,11 @@ pub(crate) struct StaticRouter { tables: Vec<DeclaredTable>, /// Declared names that name no table, which no load can fix. invalid_names: Vec<String>, - catalog: Box<dyn Catalog>, - invalid_rows: InvalidRows, - commit_retry: RetryPolicy, - held: Option<HeldMessages>, + appender: Appender, } impl StaticRouter { - pub async fn new( - catalog: Box<dyn Catalog>, - declared_tables: &Vec<String>, - invalid_rows: InvalidRows, - commit_retry: RetryPolicy, - ) -> Result<Self, Error> { + pub async fn new(appender: Appender, declared_tables: &Vec<String>) -> Result<Self, Error> { let mut tables = Vec::with_capacity(declared_tables.len()); let mut invalid_names = Vec::new(); for declared_table in declared_tables { @@ -80,24 +55,24 @@ impl StaticRouter { continue; } - let table = match load_table(declared_table, catalog.as_ref()).await { - Ok(table) => Some(Arc::new(table)), + let table = match load_table(declared_table, appender.catalog.as_ref()).await { + Ok(table) => OnceLock::from(table), Err(err) => { warn!( - "Declared table {declared_table} does not load from the configured catalog, so each batch tries again. {}", + "Declared table {declared_table} does not load from the configured catalog. At most once it is skipped until the connector restarts, at least once each batch loads it again. {}", format_error_chain(&err) ); - None + OnceLock::new() } }; tables.push(DeclaredTable { name: declared_table.clone(), - table: Mutex::new(table), + table, }); } let loaded = tables .iter() - .filter(|declared| declared.loaded().is_some()) + .filter(|declared| declared.table.get().is_some()) .count(); info!( "Static router found {loaded} tables on iceberg catalog from {} tables declared", @@ -110,37 +85,25 @@ impl StaticRouter { Ok(StaticRouter { tables, invalid_names, - catalog, - invalid_rows, - commit_retry, - held: None, + appender, }) } - /// The table of `declared`, loaded from the catalog when no earlier batch could. At most - /// once, a table that does not load is skipped for this batch. At least once, it fails the - /// batch, so that its messages are delivered again instead of lost. - async fn table(&self, declared: &DeclaredTable) -> Result<Option<Arc<Table>>, Error> { - if let Some(table) = declared.loaded() { + /// The table of `declared`. At most once, a table that did not load at startup is skipped, + /// so that a batch does not pay a catalog call for it. At least once, it is loaded again, and + /// a failed load fails the batch, so that its messages are delivered again instead of lost. + async fn table<'d>(&self, declared: &'d DeclaredTable) -> Result<Option<&'d Table>, Error> { + if let Some(table) = declared.table.get() { return Ok(Some(table)); } - match load_table(&declared.name, self.catalog.as_ref()).await { - Ok(table) => { - let table = Arc::new(table); - info!("Declared table {} loaded from the catalog", declared.name); - *declared.lock() = Some(Arc::clone(&table)); - Ok(Some(table)) - } - Err(err) if self.held.is_some() => Err(table_load_error(&declared.name, &err)), - Err(err) => { - warn!( - "Skipping declared table {} for this batch, it does not load from the catalog. {}", - declared.name, - format_error_chain(&err) - ); - Ok(None) - } + if self.appender.held.is_none() { + return Ok(None); } + let table = load_table(&declared.name, self.appender.catalog.as_ref()) + .await + .map_err(|err| table_load_error(&declared.name, &err))?; + info!("Declared table {} loaded from the catalog", declared.name); + Ok(Some(declared.table.get_or_init(|| table))) } } @@ -159,18 +122,10 @@ impl Router for StaticRouter { }; // Once a table rejects a message, the others still need every message before it. let messages = before(&messages, rejection.as_ref()); - let written = write_data( - origin, - messages, - &table, - None, - self.catalog.as_ref(), - messages_schema, - self.invalid_rows, - self.commit_retry, - self.held.as_ref(), - ) - .await?; + let written = self + .appender + .append(origin, table, false, messages_schema, messages) + .await?; info!( "Routed {} input rows to iceberg table {}, committed {}", messages.len(), @@ -193,20 +148,23 @@ impl Router for StaticRouter { self.invalid_names ))); } - self.held = Some(HeldMessages::default()); + self.appender.enable_at_least_once(); Ok(()) } + + fn flushed(&self, topic_metadata: &TopicMetadata, partitions: &[FlushPartition]) { + self.appender.flushed(topic_metadata, partitions); + } } #[cfg(test)] mod tests { use super::*; - use iceberg::memory::{MEMORY_CATALOG_WAREHOUSE, MemoryCatalogBuilder}; - use iceberg::spec::{NestedField, PrimitiveType, Schema as IcebergSchema, Type}; - use iceberg::{CatalogBuilder, NamespaceIdent, TableCreation, TableIdent}; + use crate::InvalidRows; + use crate::router::commit::tests::{SINGLE_ATTEMPT, catalog_with_table}; + use iceberg::{TableCreation, TableIdent}; use iggy_connector_sdk::Payload; use std::collections::HashMap; - use std::time::Duration; const ORIGIN: BatchOrigin<'static> = BatchOrigin { connector_id: 1, @@ -217,44 +175,10 @@ mod tests { /// A router over `declared` tables, of which the catalog holds only `test.events`. async fn router(declared: &[&str]) -> StaticRouter { - let catalog = MemoryCatalogBuilder::default() - .load( - "memory", - HashMap::from([( - MEMORY_CATALOG_WAREHOUSE.to_owned(), - "memory://warehouse".to_owned(), - )]), - ) - .await - .expect("Failed to create catalog"); - let namespace = NamespaceIdent::new("test".to_owned()); - catalog - .create_namespace(&namespace, HashMap::new()) - .await - .expect("Failed to create namespace"); - let schema = IcebergSchema::builder() - .with_fields(vec![ - NestedField::required(1, "id", Type::Primitive(PrimitiveType::Long)).into(), - ]) - .build() - .expect("Failed to build schema"); - catalog - .create_table( - &namespace, - TableCreation::builder() - .name("events".to_owned()) - .schema(schema) - .build(), - ) - .await - .expect("Failed to create table"); + let (catalog, _) = catalog_with_table(HashMap::new()).await; let declared: Vec<String> = declared.iter().map(|name| (*name).to_owned()).collect(); - let retry = RetryPolicy { - max_attempts: 1, - base_delay: Duration::from_millis(1), - max_delay: Duration::from_millis(1), - }; - StaticRouter::new(Box::new(catalog), &declared, InvalidRows::Fail, retry) + let appender = Appender::new(Box::new(catalog), InvalidRows::Fail, SINGLE_ATTEMPT); + StaticRouter::new(appender, &declared) .await .expect("One declared table should load") } @@ -271,10 +195,11 @@ mod tests { }] } - async fn snapshots(router: &StaticRouter) -> usize { + async fn snapshots(router: &StaticRouter, table: &str) -> usize { router + .appender .catalog - .load_table(&TableIdent::from_strs(["test", "events"]).expect("a valid table name")) + .load_table(&TableIdent::from_strs(["test", table]).expect("a valid table name")) .await .expect("Failed to load table") .metadata() @@ -291,7 +216,38 @@ mod tests { .await .expect("At most once, a table that does not load is skipped"); - assert_eq!(snapshots(&router).await, 1); + assert_eq!(snapshots(&router, "events").await, 1); + } + + #[tokio::test] + async fn given_a_declared_table_created_after_startup_when_at_most_once_should_keep_skipping_it() + { + let router = router(&["test.events", "test.later"]).await; + let events = router + .appender + .catalog + .load_table(&TableIdent::from_strs(["test", "events"]).expect("a valid table name")) + .await + .expect("Failed to load table"); + router + .appender + .catalog + .create_table( + events.identifier().namespace(), + TableCreation::builder() + .name("later".to_owned()) + .schema(events.metadata().current_schema().as_ref().clone()) + .build(), + ) + .await + .expect("Failed to create table"); + + router + .route_data(&ORIGIN, Schema::Json, messages()) + .await + .expect("At most once, a table that did not load at startup is skipped"); + + assert_eq!(snapshots(&router, "later").await, 0); } #[tokio::test] diff --git a/core/connectors/sinks/iceberg_sink/src/sink.rs b/core/connectors/sinks/iceberg_sink/src/sink.rs index 8530a8096..f76155136 100644 --- a/core/connectors/sinks/iceberg_sink/src/sink.rs +++ b/core/connectors/sinks/iceberg_sink/src/sink.rs @@ -18,11 +18,13 @@ use crate::{ IcebergSink, catalog::init_catalog, - router::{BatchOrigin, dynamic_router::DynamicRouter, static_router::StaticRouter}, + router::{Appender, BatchOrigin, dynamic_router::DynamicRouter, static_router::StaticRouter}, }; use async_trait::async_trait; use iceberg::Catalog; -use iggy_connector_sdk::{ConsumedMessage, Error, MessagesMetadata, Sink, TopicMetadata}; +use iggy_connector_sdk::{ + ConsumedMessage, Error, FlushPartition, MessagesMetadata, Sink, TopicMetadata, +}; use tracing::{debug, error, info}; #[async_trait] @@ -71,23 +73,16 @@ impl Sink for IcebergSink { let catalog: Box<dyn Catalog> = init_catalog(&self.config).await?; let invalid_rows = self.config.invalid_rows.unwrap_or_default(); + let appender = Appender::new(catalog, invalid_rows, self.commit_retry); if self.config.dynamic_routing { self.router = Some(Box::new(DynamicRouter::new( - catalog, + appender, self.config.dynamic_route_field.clone(), - invalid_rows, - self.commit_retry, ))) } else { self.router = Some(Box::new( - StaticRouter::new( - catalog, - &self.config.tables, - invalid_rows, - self.commit_retry, - ) - .await?, + StaticRouter::new(appender, &self.config.tables).await?, )); } @@ -130,8 +125,7 @@ impl Sink for IcebergSink { Ok(()) } - /// The sink commits every batch inside `consume`, so the default flush has nothing left to - /// prove. + /// The sink commits every batch inside `consume`, so a flush has nothing left to prove. fn enable_at_least_once(&mut self) -> Result<(), Error> { let Some(router) = self.router.as_mut() else { return Err(Error::InvalidState); @@ -139,6 +133,20 @@ impl Sink for IcebergSink { router.enable_at_least_once() } + /// Writes nothing, since `consume` committed every message already. It lets the router forget + /// what the tables hold of the flushed messages, which the runtime never hands over again. + async fn flush( + &self, + topic_metadata: &TopicMetadata, + partitions: &[FlushPartition], + ) -> Result<(), Error> { + let Some(router) = &self.router else { + return Err(Error::InvalidState); + }; + router.flushed(topic_metadata, partitions); + Ok(()) + } + async fn close(&mut self) -> Result<(), Error> { info!("Iceberg sink connector with ID: {} is closed.", self.id); Ok(())
