refactor: simplify snapshot expiration implementation and remove builder struct

ForeverAngry · JanKaul · commit f2ef422e51a3 · 2026-01-07T08:30:01.000+01:00
diff --git a/iceberg-rust/src/lib.rs b/iceberg-rust/src/lib.rs
@@ -46,13 +46,17 @@
 //!     .await?;
 //!
 //! // Expire old snapshots for maintenance
-//! let expired_snapshot_ids = table.expire_snapshots()
-//!     .expire_older_than(chrono::Utc::now().timestamp_millis() - 30 * 24 * 60 * 60 * 1000)
-//!     .retain_last(10)
-//!     .execute()
+//! table
+//!     .new_transaction(None)
+//!     .expire_snapshots(
+//!         Some(chrono::Utc::now().timestamp_millis() - 30 * 24 * 60 * 60 * 1000),
+//!         Some(10),
+//!         true,
+//!         true,
+//!         false,
+//!     )
+//!     .commit()
 //!     .await?;
-//!     
-//! println!("Expired {} snapshots", expired_snapshot_ids.len());
 //! # Ok(())
 //! # }
 //! ```
diff --git a/iceberg-rust/src/table/mod.rs b/iceberg-rust/src/table/mod.rs
@@ -45,86 +45,6 @@ pub mod manifest;
 pub mod manifest_list;
 pub mod transaction;
 
-/// Builder for configuring and executing snapshot expiration operations
-///
-/// This builder provides a fluent API for configuring how snapshots should be expired:
-/// * [`expire_older_than`](ExpireSnapshotsBuilder::expire_older_than) - Remove snapshots older than a timestamp
-/// * [`retain_last`](ExpireSnapshotsBuilder::retain_last) - Keep only the most recent N snapshots
-/// * [`clean_orphan_files`](ExpireSnapshotsBuilder::clean_orphan_files) - Also remove unreferenced data files
-/// * [`dry_run`](ExpireSnapshotsBuilder::dry_run) - Preview what would be deleted without actually deleting
-pub struct ExpireSnapshotsBuilder<'a> {
-    table: &'a mut Table,
-    older_than: Option<i64>,
-    retain_last: Option<usize>,
-    clean_orphan_files: bool,
-    retain_ref_snapshots: bool,
-    dry_run: bool,
-}
-
-impl<'a> ExpireSnapshotsBuilder<'a> {
-    /// Create a new snapshot expiration builder for the given table
-    fn new(table: &'a mut Table) -> Self {
-        Self {
-            table,
-            older_than: None,
-            retain_last: None,
-            clean_orphan_files: false,
-            retain_ref_snapshots: true,
-            dry_run: false,
-        }
-    }
-
-    /// Expire snapshots older than the given timestamp (in milliseconds since Unix epoch)
-    pub fn expire_older_than(mut self, timestamp_ms: i64) -> Self {
-        self.older_than = Some(timestamp_ms);
-        self
-    }
-
-    /// Retain only the most recent N snapshots, expiring all others
-    pub fn retain_last(mut self, count: usize) -> Self {
-        self.retain_last = Some(count);
-        self
-    }
-
-    /// Enable or disable cleanup of orphaned data files
-    pub fn clean_orphan_files(mut self, enabled: bool) -> Self {
-        self.clean_orphan_files = enabled;
-        self
-    }
-
-    /// Control whether snapshots referenced by branches/tags should be preserved
-    pub fn retain_ref_snapshots(mut self, enabled: bool) -> Self {
-        self.retain_ref_snapshots = enabled;
-        self
-    }
-
-    /// Enable dry run mode to preview what would be deleted without actually deleting
-    pub fn dry_run(mut self, enabled: bool) -> Self {
-        self.dry_run = enabled;
-        self
-    }
-
-    /// Execute the snapshot expiration operation
-    pub async fn execute(self) -> Result<Vec<i64>, Error> {
-        let _result = self.table.new_transaction(None)
-            .expire_snapshots(
-                self.older_than,
-                self.retain_last,
-                self.clean_orphan_files,
-                self.retain_ref_snapshots,
-                self.dry_run,
-            )
-            .commit()
-            .await?;
-
-        // Extract the expired snapshot IDs from the commit result
-        // For now, we'll need to return empty vec since the transaction commit
-        // doesn't directly return the expired snapshot IDs
-        // TODO: Enhance transaction result to include operation-specific details
-        Ok(vec![])
-    }
-}
-
 #[derive(Debug, Clone)]
 /// Iceberg table
 pub struct Table {
@@ -378,39 +298,6 @@ impl Table {
         TableTransaction::new(self, branch)
     }
 
-    /// Configures snapshot expiration for this table
-    ///
-    /// Returns a builder that allows configuring snapshot expiration policies:
-    /// * Time-based expiration: Remove snapshots older than a timestamp
-    /// * Count-based retention: Keep only the most recent N snapshots  
-    /// * Orphan file cleanup: Remove data files no longer referenced by any snapshot
-    /// * Reference preservation: Protect snapshots referenced by branches/tags
-    /// * Dry run mode: Preview what would be deleted without actually deleting
-    ///
-    /// The operation is executed through the table's transaction system, ensuring
-    /// atomicity and consistency with other table operations.
-    ///
-    /// # Returns
-    /// * `ExpireSnapshotsBuilder` - A builder for configuring expiration parameters
-    ///
-    /// # Examples
-    /// ```rust,no_run
-    /// # async fn example(table: &mut Table) -> Result<(), Box<dyn std::error::Error>> {
-    /// // Expire snapshots older than 7 days but keep at least 5 snapshots
-    /// let result = table.expire_snapshots()
-    ///     .expire_older_than(chrono::Utc::now().timestamp_millis() - 7 * 24 * 60 * 60 * 1000)
-    ///     .retain_last(5)
-    ///     .clean_orphan_files(true)
-    ///     .execute()
-    ///     .await?;
-    /// 
-    /// println!("Expired {} snapshots", result.len());
-    /// # Ok(())
-    /// # }
-    /// ```
-    pub fn expire_snapshots(&mut self) -> ExpireSnapshotsBuilder<'_> {
-        ExpireSnapshotsBuilder::new(self)
-    }
 }
 
 /// Path of a Manifest file
diff --git a/iceberg-rust/src/table/transaction/operation.rs b/iceberg-rust/src/table/transaction/operation.rs
@@ -941,6 +941,145 @@ impl Operation {
     }
 }
 
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use futures::executor::block_on;
+    use iceberg_rust_spec::spec::schema::SchemaBuilder;
+    use iceberg_rust_spec::spec::table_metadata::TableMetadataBuilder;
+    use iceberg_rust_spec::spec::types::{PrimitiveType, StructField, Type};
+    use object_store::memory::InMemory;
+
+    fn sample_metadata(
+        snapshot_defs: &[(i64, i64)],
+        current_snapshot: Option<i64>,
+        refs: &[(&str, i64)],
+    ) -> TableMetadata {
+        let schema = SchemaBuilder::default()
+            .with_schema_id(0)
+            .with_struct_field(StructField {
+                id: 1,
+                name: "id".to_string(),
+                required: true,
+                field_type: Type::Primitive(PrimitiveType::Long),
+                doc: None,
+            })
+            .build()
+            .unwrap();
+
+        let snapshots = snapshot_defs
+            .iter()
+            .enumerate()
+            .map(|(idx, (snapshot_id, timestamp))| {
+                let snapshot = SnapshotBuilder::default()
+                    .with_snapshot_id(*snapshot_id)
+                    .with_sequence_number((idx + 1) as i64)
+                    .with_timestamp_ms(*timestamp)
+                    .with_manifest_list(format!("manifest-{snapshot_id}.avro"))
+                    .with_summary(Summary {
+                        operation: SnapshotOperation::Append,
+                        other: HashMap::new(),
+                    })
+                    .with_schema_id(0)
+                    .build()
+                    .unwrap();
+                (*snapshot_id, snapshot)
+            })
+            .collect::<HashMap<_, _>>();
+
+        let refs = refs
+            .iter()
+            .map(|(name, snapshot_id)| {
+                (
+                    (*name).to_string(),
+                    SnapshotReference {
+                        snapshot_id: *snapshot_id,
+                        retention: SnapshotRetention::default(),
+                    },
+                )
+            })
+            .collect::<HashMap<_, _>>();
+
+        TableMetadataBuilder::default()
+            .location("s3://tests/table".to_owned())
+            .current_schema_id(0)
+            .schemas(HashMap::from_iter(vec![(0, schema)]))
+            .snapshots(snapshots)
+            .current_snapshot_id(current_snapshot)
+            .last_sequence_number(snapshot_defs.len() as i64)
+            .refs(refs)
+            .build()
+            .unwrap()
+    }
+
+    fn execute_operation(
+        metadata: &TableMetadata,
+        older_than: Option<i64>,
+        retain_last: Option<usize>,
+        retain_refs: bool,
+        dry_run: bool,
+    ) -> Result<Vec<TableUpdate>, Error> {
+        let op = Operation::ExpireSnapshots {
+            older_than,
+            retain_last,
+            clean_orphan_files: false,
+            retain_ref_snapshots: retain_refs,
+            dry_run,
+        };
+        let store = Arc::new(InMemory::new());
+        block_on(op.execute(metadata, store)).map(|(_, updates)| updates)
+    }
+
+    fn collect_snapshot_ids(updates: &[TableUpdate]) -> Vec<i64> {
+        updates
+            .iter()
+            .flat_map(|update| match update {
+                TableUpdate::RemoveSnapshots { snapshot_ids } => snapshot_ids.clone(),
+                _ => Vec::new(),
+            })
+            .collect()
+    }
+
+    #[test]
+    fn snapshot_expiration_requires_policy() {
+        let metadata = sample_metadata(&[(1, 1_000)], Some(1), &[]);
+        let result = execute_operation(&metadata, None, None, true, false);
+        assert!(matches!(result, Err(Error::InvalidFormat(_))));
+    }
+
+    #[test]
+    fn snapshot_expiration_applies_time_and_count_filters() {
+        let metadata = sample_metadata(
+            &[(1, 1_000), (2, 2_000), (3, 3_000), (4, 4_000)],
+            Some(4),
+            &[],
+        );
+        let updates = execute_operation(&metadata, Some(2_500), Some(2), true, false).unwrap();
+        let mut expired = collect_snapshot_ids(&updates);
+        expired.sort();
+        assert_eq!(expired, vec![1, 2]);
+    }
+
+    #[test]
+    fn snapshot_expiration_preserves_current_and_refs() {
+        let metadata = sample_metadata(
+            &[(10, 1_000), (20, 2_000), (30, 3_000)],
+            Some(30),
+            &[ ("branch", 20) ],
+        );
+        let updates = execute_operation(&metadata, Some(1_500), None, true, false).unwrap();
+        // Snapshot 10 is the only candidate because 20 is referenced and 30 is current.
+        assert_eq!(collect_snapshot_ids(&updates), vec![10]);
+    }
+
+    #[test]
+    fn snapshot_expiration_supports_dry_run() {
+        let metadata = sample_metadata(&[(1, 1_000), (2, 900)], Some(1), &[]);
+        let updates = execute_operation(&metadata, Some(950), None, true, true).unwrap();
+        assert!(updates.is_empty());
+    }
+}
+
 pub fn bounding_partition_values<'a>(
     mut iter: impl Iterator<Item = &'a DataFile>,
     partition_column_names: &SmallVec<[&str; 4]>,
