diff --git a/crates/iceberg/public-api.txt b/crates/iceberg/public-api.txt index b6461d77f5..28bc70c3a0 100644 --- a/crates/iceberg/public-api.txt +++ b/crates/iceberg/public-api.txt @@ -3693,7 +3693,9 @@ pub struct iceberg::transaction::ActionCommit impl iceberg::transaction::ActionCommit pub fn iceberg::transaction::ActionCommit::new(updates: alloc::vec::Vec, requirements: alloc::vec::Vec) -> Self pub fn iceberg::transaction::ActionCommit::take_requirements(&mut self) -> alloc::vec::Vec +pub fn iceberg::transaction::ActionCommit::take_unchecked_requirements(&mut self) -> alloc::vec::Vec pub fn iceberg::transaction::ActionCommit::take_updates(&mut self) -> alloc::vec::Vec +pub fn iceberg::transaction::ActionCommit::with_unchecked_requirements(self, requirements: alloc::vec::Vec) -> Self pub struct iceberg::transaction::AddColumn impl iceberg::transaction::AddColumn pub fn iceberg::transaction::AddColumn::optional(name: impl alloc::string::ToString, field_type: iceberg::spec::Type) -> Self @@ -3714,6 +3716,7 @@ pub fn iceberg::transaction::FastAppendAction::with_check_duplicate(self, v: boo pub struct iceberg::transaction::OverwriteAction impl iceberg::transaction::OverwriteAction pub fn iceberg::transaction::OverwriteAction::add_data_files(self, data_files: impl core::iter::traits::collect::IntoIterator) -> Self +pub fn iceberg::transaction::OverwriteAction::assert_requirements(self, requirements: impl core::iter::traits::collect::IntoIterator) -> Self pub fn iceberg::transaction::OverwriteAction::delete_data_files(self, data_files: impl core::iter::traits::collect::IntoIterator) -> Self pub fn iceberg::transaction::OverwriteAction::set_commit_uuid(self, commit_uuid: uuid::Uuid) -> Self pub fn iceberg::transaction::OverwriteAction::set_key_metadata(self, key_metadata: alloc::vec::Vec) -> Self diff --git a/crates/iceberg/src/transaction/action.rs b/crates/iceberg/src/transaction/action.rs index aa0a05d0d9..fc10d83201 100644 --- a/crates/iceberg/src/transaction/action.rs +++ b/crates/iceberg/src/transaction/action.rs @@ -81,6 +81,7 @@ impl ApplyTransactionAction for T { pub struct ActionCommit { updates: Vec, requirements: Vec, + unchecked_requirements: Vec, } impl ActionCommit { @@ -89,18 +90,41 @@ impl ActionCommit { Self { updates, requirements, + unchecked_requirements: vec![], } } + /// Attaches requirements that should be included in the final commit but must *not* be + /// locally validated against the table this action is being applied to. + /// + /// This is for caller-supplied requirements (e.g. via `OverwriteAction::assert_requirements`) + /// that may deliberately assert something other than what the in-memory `Table` currently + /// reports — the whole point of such an assertion is to let the catalog (or whatever + /// external component eventually submits the commit) be the arbiter, not this crate's own + /// local, `table.metadata()`-derived sanity check. Requirements passed to [`Self::new`] + /// remain locally checked, since those are always derived from — and so trivially consistent + /// with — the very table metadata they'd be checked against. + pub fn with_unchecked_requirements(mut self, requirements: Vec) -> Self { + self.unchecked_requirements = requirements; + self + } + /// Consumes and returns the list of table updates. pub fn take_updates(&mut self) -> Vec { take(&mut self.updates) } - /// Consumes and returns the list of table requirements. + /// Consumes and returns the list of table requirements that should be locally validated + /// against the table this action was applied to. pub fn take_requirements(&mut self) -> Vec { take(&mut self.requirements) } + + /// Consumes and returns the list of table requirements that should be included in the final + /// commit but skipped during local validation. See [`Self::with_unchecked_requirements`]. + pub fn take_unchecked_requirements(&mut self) -> Vec { + take(&mut self.unchecked_requirements) + } } #[cfg(test)] diff --git a/crates/iceberg/src/transaction/append.rs b/crates/iceberg/src/transaction/append.rs index 881bd5d752..2269ba0921 100644 --- a/crates/iceberg/src/transaction/append.rs +++ b/crates/iceberg/src/transaction/append.rs @@ -93,7 +93,8 @@ impl TransactionAction for FastAppendAction { self.commit_uuid.unwrap_or_else(Uuid::now_v7), self.snapshot_properties.clone(), self.dedupe_added_files(), - vec![], + vec![], // deleted_data_files: fast-append never deletes + vec![], // extra_requirements: not exposed on FastAppendAction ); // validate added files diff --git a/crates/iceberg/src/transaction/mod.rs b/crates/iceberg/src/transaction/mod.rs index b41f21f38e..0094305676 100644 --- a/crates/iceberg/src/transaction/mod.rs +++ b/crates/iceberg/src/transaction/mod.rs @@ -120,6 +120,7 @@ impl Transaction { ) -> Result { let updates = action_commit.take_updates(); let requirements = action_commit.take_requirements(); + let unchecked_requirements = action_commit.take_unchecked_requirements(); for requirement in &requirements { requirement.check(Some(table.metadata()))?; @@ -129,6 +130,7 @@ impl Transaction { existing_updates.extend(updates); existing_requirements.extend(requirements); + existing_requirements.extend(unchecked_requirements); Ok(updated_table) } diff --git a/crates/iceberg/src/transaction/overwrite.rs b/crates/iceberg/src/transaction/overwrite.rs index 051c753753..c62c49b1bb 100644 --- a/crates/iceberg/src/transaction/overwrite.rs +++ b/crates/iceberg/src/transaction/overwrite.rs @@ -31,7 +31,7 @@ use crate::transaction::snapshot::{ DefaultManifestProcess, SnapshotProduceOperation, SnapshotProducer, }; use crate::transaction::{ActionCommit, TransactionAction}; -use crate::{Error, ErrorKind}; +use crate::{Error, ErrorKind, TableRequirement}; /// OverwriteAction is a transaction action for overwriting data files in the table. /// @@ -43,6 +43,7 @@ pub struct OverwriteAction { commit_uuid: Option, key_metadata: Option>, snapshot_properties: HashMap, + extra_requirements: Vec, added_data_files: Vec, deleted_data_files: Vec, } @@ -54,6 +55,7 @@ impl OverwriteAction { commit_uuid: None, key_metadata: None, snapshot_properties: HashMap::default(), + extra_requirements: vec![], added_data_files: vec![], deleted_data_files: vec![], } @@ -94,6 +96,25 @@ impl OverwriteAction { self.snapshot_properties = snapshot_properties; self } + + /// Assert additional table requirements before this commit is attempted, alongside the + /// ones this action derives automatically from the table it's applied to (currently: + /// `UuidMatch`, and a `RefSnapshotIdMatch` for the branch this commit targets). + /// + /// A supplied requirement of the same kind as one this action would otherwise derive + /// automatically replaces the derived one, rather than producing a duplicate — e.g. to + /// assert a pinned base snapshot id that differs from whatever `Transaction::commit`'s + /// internal retry loop has rebased this action's `Table` onto by the time it runs, supply + /// your own `TableRequirement::RefSnapshotIdMatch` here. Any other requirement kind (schema id, + /// partition spec id, sort order id, ...) is added alongside the derived ones. Calling + /// this more than once accumulates requirements rather than replacing the whole set. + pub fn assert_requirements( + mut self, + requirements: impl IntoIterator, + ) -> Self { + self.extra_requirements.extend(requirements); + self + } } #[async_trait] @@ -112,6 +133,7 @@ impl TransactionAction for OverwriteAction { self.snapshot_properties.clone(), self.added_data_files.clone(), self.deleted_data_files.clone(), + self.extra_requirements.clone(), ); snapshot_producer.validate_added_data_files()?; @@ -480,6 +502,130 @@ mod tests { assert_eq!(data_file, *manifest.entries()[0].data_file()); } + #[tokio::test] + async fn test_overwrite_assert_requirements_overrides_ref_snapshot_id_match() { + let table = make_v2_minimal_table(); + + let data_file = test_data_file( + "test/override.parquet", + table.metadata().default_partition_spec_id(), + ); + + // Deliberately different from the table's actual current snapshot id, to prove this + // is what gets asserted, not what `SnapshotProducer` would derive on its own. + let overridden_base = table + .metadata() + .current_snapshot_id() + .map(|id| id + 1) + .unwrap_or(1); + + let tx = Transaction::new(&table); + let action = tx + .overwrite() + .add_data_files(vec![data_file]) + .assert_requirements(vec![TableRequirement::RefSnapshotIdMatch { + r#ref: MAIN_BRANCH.to_string(), + snapshot_id: Some(overridden_base), + }]); + let mut action_commit = Arc::new(action).commit(&table).await.unwrap(); + let mut requirements = action_commit.take_requirements(); + requirements.extend(action_commit.take_unchecked_requirements()); + + // Exactly one ref-snapshot-id requirement for `main`, and it's the override -- not + // duplicated, and not the table's real current snapshot id. + let ref_requirements: Vec<_> = requirements + .iter() + .filter(|r| matches!(r, TableRequirement::RefSnapshotIdMatch { r#ref, .. } if r#ref == MAIN_BRANCH)) + .collect(); + assert_eq!(ref_requirements.len(), 1); + assert_eq!(ref_requirements[0], &TableRequirement::RefSnapshotIdMatch { + r#ref: MAIN_BRANCH.to_string(), + snapshot_id: Some(overridden_base), + }); + + // The other auto-derived requirement is untouched. + assert!( + requirements + .iter() + .any(|r| matches!(r, TableRequirement::UuidMatch { uuid } if *uuid == table.metadata().uuid())) + ); + assert_eq!(requirements.len(), 2); + } + + #[tokio::test] + async fn test_overwrite_assert_requirements_adds_new_requirement_kind() { + let table = make_v2_minimal_table(); + + let data_file = test_data_file( + "test/extra.parquet", + table.metadata().default_partition_spec_id(), + ); + + let tx = Transaction::new(&table); + let action = tx + .overwrite() + .add_data_files(vec![data_file]) + .assert_requirements(vec![TableRequirement::CurrentSchemaIdMatch { + current_schema_id: table.metadata().current_schema_id(), + }]); + let mut action_commit = Arc::new(action).commit(&table).await.unwrap(); + let mut requirements = action_commit.take_requirements(); + requirements.extend(action_commit.take_unchecked_requirements()); + + // Both auto-derived requirements are still present, plus the supplied one -- nothing + // replaced, since `CurrentSchemaIdMatch` isn't one of the kinds this action derives. + assert_eq!(requirements.len(), 3); + assert!( + requirements + .iter() + .any(|r| matches!(r, TableRequirement::UuidMatch { .. })) + ); + assert!(requirements.iter().any( + |r| matches!(r, TableRequirement::RefSnapshotIdMatch { r#ref, .. } if r#ref == MAIN_BRANCH) + )); + assert!(requirements.iter().any(|r| matches!( + r, + TableRequirement::CurrentSchemaIdMatch { current_schema_id } if *current_schema_id == table.metadata().current_schema_id() + ))); + } + + #[tokio::test] + async fn test_overwrite_assert_requirements_accumulates_across_calls() { + let table = make_v2_minimal_table(); + + let data_file = test_data_file( + "test/accumulate.parquet", + table.metadata().default_partition_spec_id(), + ); + + let tx = Transaction::new(&table); + let action = tx + .overwrite() + .add_data_files(vec![data_file]) + .assert_requirements(vec![TableRequirement::CurrentSchemaIdMatch { + current_schema_id: table.metadata().current_schema_id(), + }]) + .assert_requirements(vec![TableRequirement::LastAssignedFieldIdMatch { + last_assigned_field_id: table.metadata().last_column_id(), + }]); + let mut action_commit = Arc::new(action).commit(&table).await.unwrap(); + let mut requirements = action_commit.take_requirements(); + requirements.extend(action_commit.take_unchecked_requirements()); + + // Both calls' requirements survive, alongside the two auto-derived ones: 4 total. + assert_eq!(requirements.len(), 4); + assert!( + requirements + .iter() + .any(|r| matches!(r, TableRequirement::CurrentSchemaIdMatch { .. })) + ); + assert!( + requirements + .iter() + .any(|r| matches!(r, TableRequirement::LastAssignedFieldIdMatch { .. })) + ); + } + #[tokio::test] async fn test_overwrite_with_deleted_files() { use crate::memory::tests::new_memory_catalog; diff --git a/crates/iceberg/src/transaction/snapshot.rs b/crates/iceberg/src/transaction/snapshot.rs index 1a1f87424a..6162b99127 100644 --- a/crates/iceberg/src/transaction/snapshot.rs +++ b/crates/iceberg/src/transaction/snapshot.rs @@ -123,6 +123,7 @@ pub(crate) struct SnapshotProducer<'a> { snapshot_id: i64, commit_uuid: Uuid, snapshot_properties: HashMap, + extra_requirements: Vec, added_data_files: Vec, deleted_data_files: Vec, // A counter used to generate unique manifest file names. @@ -146,12 +147,14 @@ impl<'a> SnapshotProducer<'a> { snapshot_properties: HashMap, added_data_files: Vec, deleted_data_files: Vec, + extra_requirements: Vec, ) -> Self { Self { table, snapshot_id: Self::generate_unique_snapshot_id(table), commit_uuid, snapshot_properties, + extra_requirements, added_data_files, deleted_data_files, manifest_counter: (0..), @@ -581,16 +584,36 @@ impl<'a> SnapshotProducer<'a> { ]] .concat(); - let requirements = vec![ - TableRequirement::UuidMatch { + // Caller-supplied requirements (via `assert_requirements`) are intentionally *not* + // locally checked against `self.table`: the whole reason to supply one is that it may + // assert a base other than whatever this in-memory `Table` currently reports (e.g. a + // pinned base that a retry must keep asserting even though a fresher load of the table + // has since moved past it). Only the auto-derived fallbacks below — which are always + // read straight off `self.table` and so can never disagree with it — are locally + // checked; that check is a self-consistency safeguard across multiple actions within one + // transaction, not a substitute for the catalog's own requirement check. + let extra_requirements = self.extra_requirements; + let mut auto_requirements = vec![]; + + if !extra_requirements + .iter() + .any(|r| matches!(r, TableRequirement::UuidMatch { .. })) + { + auto_requirements.push(TableRequirement::UuidMatch { uuid: self.table.metadata().uuid(), - }, - TableRequirement::RefSnapshotIdMatch { + }); + } + + if !extra_requirements.iter().any( + |r| matches!(r, TableRequirement::RefSnapshotIdMatch { r#ref, .. } if r#ref == MAIN_BRANCH), + ) { + auto_requirements.push(TableRequirement::RefSnapshotIdMatch { r#ref: MAIN_BRANCH.to_string(), snapshot_id: self.table.metadata().current_snapshot_id(), - }, - ]; + }); + } - Ok(ActionCommit::new(updates, requirements)) + Ok(ActionCommit::new(updates, auto_requirements) + .with_unchecked_requirements(extra_requirements)) } }