1

Author: HeartLinked

src/iceberg/table_requirements.cc

@@ -97,6 +97,11 @@ void TableUpdateContext::RequireNoBranchesChanged() {
   }
 }
 
+bool TableUpdateContext::AddChangedRef(const std::string& ref_name) {
+  auto [it, inserted] = changed_refs_.insert(ref_name);
+  return inserted;
+}
+
 Result<std::vector<std::unique_ptr<TableRequirement>>> TableRequirements::ForCreateTable(
     const std::vector<std::unique_ptr<TableUpdate>>& table_updates) {
   TableUpdateContext context(nullptr, false);

src/iceberg/table_requirements.h

@@ -27,6 +27,8 @@
 /// for optimistic concurrency control when committing table changes.
 
 #include <memory>
+#include <string>
+#include <unordered_set>
 #include <vector>
 
 #include "iceberg/iceberg_export.h"
@@ -82,6 +84,11 @@ class ICEBERG_EXPORT TableUpdateContext {
   /// \brief Require that no branches have been changed
   void RequireNoBranchesChanged();
 
+  /// \brief Track a changed ref and return whether it was newly added
+  /// \param ref_name The name of the ref being changed
+  /// \return true if this is the first time the ref is being changed
+  bool AddChangedRef(const std::string& ref_name);
+
  private:
   const TableMetadata* base_;
   const bool is_replace_;
@@ -94,6 +101,9 @@ class ICEBERG_EXPORT TableUpdateContext {
   bool added_last_assigned_partition_id_ = false;
   bool added_default_spec_id_ = false;
   bool added_default_sort_order_id_ = false;
+
+  // Track refs that have been changed to avoid duplicate requirements
+  std::unordered_set<std::string> changed_refs_;
 };
 
 /// \brief Factory class for generating table requirements

src/iceberg/table_update.cc

@@ -142,7 +142,7 @@ void AddSnapshot::GenerateRequirements(TableUpdateContext& context) const {
 void RemoveSnapshots::ApplyTo(TableMetadataBuilder& builder) const {}
 
 void RemoveSnapshots::GenerateRequirements(TableUpdateContext& context) const {
-  throw IcebergError("RemoveTableSnapshots::GenerateRequirements not implemented");
+  // RemoveSnapshots doesn't generate any requirements
 }
 
 // RemoveSnapshotRef
@@ -162,7 +162,16 @@ void SetSnapshotRef::ApplyTo(TableMetadataBuilder& builder) const {
 }
 
 void SetSnapshotRef::GenerateRequirements(TableUpdateContext& context) const {
-  throw NotImplemented("SetTableSnapshotRef::GenerateRequirements not implemented");
+  bool added = context.AddChangedRef(ref_name_);
+  if (added && context.base() != nullptr && !context.is_replace()) {
+    const auto& refs = context.base()->refs;
+    auto it = refs.find(ref_name_);
+    // Require that the ref does not exist (nullopt) or is the same as the base snapshot
+    std::optional<int64_t> base_snapshot_id =
+        (it != refs.end()) ? std::make_optional(it->second->snapshot_id) : std::nullopt;
+    context.AddRequirement(
+        std::make_unique<table::AssertRefSnapshotID>(ref_name_, base_snapshot_id));
+  }
 }
 
 // SetProperties

src/iceberg/test/table_requirements_test.cc

@@ -20,6 +20,7 @@
 #include "iceberg/table_requirements.h"
 
 #include <memory>
+#include <ranges>
 #include <string>
 #include <vector>
 
@@ -33,6 +34,7 @@
 #include "iceberg/table_requirement.h"
 #include "iceberg/table_update.h"
 #include "iceberg/test/matchers.h"
+#include "iceberg/type.h"
 
 namespace iceberg {
 
@@ -60,21 +62,17 @@ std::unique_ptr<TableMetadata> CreateBaseMetadata(
 // Helper function to create a simple schema for tests
 std::shared_ptr<Schema> CreateTestSchema(int32_t schema_id = 0) {
   std::vector<SchemaField> fields;
-  fields.emplace_back(SchemaField::MakeRequired(1, "id", std::make_shared<IntType>()));
+  fields.emplace_back(SchemaField::MakeRequired(1, "id", int32()));
   return std::make_shared<Schema>(std::move(fields), schema_id);
 }
 
 // Helper function to count requirements of a specific type
 template <typename T>
 int CountRequirementsOfType(
     const std::vector<std::unique_ptr<TableRequirement>>& requirements) {
-  int count = 0;
-  for (const auto& req : requirements) {
-    if (dynamic_cast<T*>(req.get()) != nullptr) {
-      count++;
-    }
-  }
-  return count;
+  return std::ranges::count_if(requirements, [](const auto& req) {
+    return dynamic_cast<T*>(req.get()) != nullptr;
+  });
 }
 
 // Helper function to add a branch to metadata
@@ -815,47 +813,98 @@ TEST(TableRequirementsTest, SetDefaultSortOrderFailure) {
 
 TEST(TableRequirementsTest, AddSnapshot) {
   auto metadata = CreateBaseMetadata();
-  std::vector<std::unique_ptr<TableUpdate>> updates;
 
+  std::vector<std::unique_ptr<TableUpdate>> updates;
   auto snapshot = std::make_shared<Snapshot>();
   snapshot->snapshot_id = 1;
   snapshot->sequence_number = 1;
   snapshot->timestamp_ms = TimePointMs{std::chrono::milliseconds(1000)};
   snapshot->manifest_list = "s3://bucket/manifest_list";
-
   updates.push_back(std::make_unique<table::AddSnapshot>(snapshot));
 
   auto result = TableRequirements::ForUpdateTable(*metadata, updates);
   ASSERT_THAT(result, IsOk());
 
   auto& requirements = result.value();
-  // AddSnapshot doesn't add additional requirements
   ASSERT_EQ(requirements.size(), 1);
   EXPECT_EQ(CountRequirementsOfType<table::AssertUUID>(requirements), 1);
 
-  // Validate against base metadata
   for (const auto& req : requirements) {
     EXPECT_THAT(req->Validate(metadata.get()), IsOk());
   }
 }
 
+// RemoveSnapshots Tests
+
+TEST(TableRequirementsTest, RemoveSnapshots) {
+  auto metadata = CreateBaseMetadata();
+
+  std::vector<std::unique_ptr<TableUpdate>> updates;
+  updates.push_back(std::make_unique<table::RemoveSnapshots>(std::vector<int64_t>{0}));
+
+  auto result = TableRequirements::ForUpdateTable(*metadata, updates);
+  ASSERT_THAT(result, IsOk());
+
+  auto& requirements = result.value();
+  ASSERT_EQ(requirements.size(), 1);
+  EXPECT_EQ(CountRequirementsOfType<table::AssertUUID>(requirements), 1);
+
+  for (const auto& req : requirements) {
+    EXPECT_THAT(req->Validate(metadata.get()), IsOk());
+  }
+}
+
+// SetSnapshotRef Tests
+
+TEST(TableRequirementsTest, SetSnapshotRef) {
+  constexpr int64_t kSnapshotId = 14;
+  const std::string kRefName = "branch";
+
+  auto metadata = CreateBaseMetadata();
+  AddBranch(*metadata, kRefName, kSnapshotId);
+
+  // Multiple updates to same ref should deduplicate
+  std::vector<std::unique_ptr<TableUpdate>> updates;
+  updates.push_back(std::make_unique<table::SetSnapshotRef>(kRefName, kSnapshotId,
+                                                            SnapshotRefType::kBranch));
+  updates.push_back(std::make_unique<table::SetSnapshotRef>(kRefName, kSnapshotId + 1,
+                                                            SnapshotRefType::kBranch));
+  updates.push_back(std::make_unique<table::SetSnapshotRef>(kRefName, kSnapshotId + 2,
+                                                            SnapshotRefType::kBranch));
+
+  auto result = TableRequirements::ForUpdateTable(*metadata, updates);
+  ASSERT_THAT(result, IsOk());
+
+  auto& requirements = result.value();
+  for (const auto& req : requirements) {
+    EXPECT_THAT(req->Validate(metadata.get()), IsOk());
+  }
+
+  ASSERT_EQ(requirements.size(), 2);
+  EXPECT_EQ(CountRequirementsOfType<table::AssertUUID>(requirements), 1);
+  EXPECT_EQ(CountRequirementsOfType<table::AssertRefSnapshotID>(requirements), 1);
+
+  auto* assert_ref = dynamic_cast<table::AssertRefSnapshotID*>(requirements[1].get());
+  ASSERT_NE(assert_ref, nullptr);
+  EXPECT_EQ(assert_ref->snapshot_id(), kSnapshotId);
+  EXPECT_EQ(assert_ref->ref_name(), kRefName);
+}
+
 // RemoveSnapshotRef Tests
 
 TEST(TableRequirementsTest, RemoveSnapshotRef) {
   auto metadata = CreateBaseMetadata();
-  std::vector<std::unique_ptr<TableUpdate>> updates;
 
+  std::vector<std::unique_ptr<TableUpdate>> updates;
   updates.push_back(std::make_unique<table::RemoveSnapshotRef>("branch"));
 
   auto result = TableRequirements::ForUpdateTable(*metadata, updates);
   ASSERT_THAT(result, IsOk());
 
   auto& requirements = result.value();
-  // RemoveSnapshotRef doesn't add additional requirements
   ASSERT_EQ(requirements.size(), 1);
   EXPECT_EQ(CountRequirementsOfType<table::AssertUUID>(requirements), 1);
 
-  // Validate against base metadata
   for (const auto& req : requirements) {
     EXPECT_THAT(req->Validate(metadata.get()), IsOk());
   }