ver1

Author: HeartLinked

src/iceberg/CMakeLists.txt

@@ -81,6 +81,7 @@ set(ICEBERG_SOURCES
     transform_function.cc
     type.cc
     update/pending_update.cc
+    update/set_snapshot.cc
     update/update_partition_spec.cc
     update/update_properties.cc
     update/update_schema.cc

src/iceberg/table_metadata.cc

@@ -591,6 +591,8 @@ class TableMetadataBuilder::Impl {
   Status RemoveSchemas(const std::unordered_set<int32_t>& schema_ids);
   Result<int32_t> AddSchema(const Schema& schema, int32_t new_last_column_id);
   void SetLocation(std::string_view location);
+  Status SetBranchSnapshot(int64_t snapshot_id, const std::string& branch);
+  Status SetRef(const std::string& name, std::shared_ptr<SnapshotRef> ref);
 
   Result<std::unique_ptr<TableMetadata>> Build();
 
@@ -613,6 +615,23 @@ class TableMetadataBuilder::Impl {
   /// \return The ID to use for this schema (reused if exists, new otherwise
   int32_t ReuseOrCreateNewSchemaId(const Schema& new_schema) const;
 
+  /// \brief Internal method to set a branch snapshot
+  /// \param snapshot The snapshot to set
+  /// \param branch The branch name
+  Status SetBranchSnapshotInternal(const Snapshot& snapshot, const std::string& branch);
+
+  /// \brief Helper to create a SnapshotRef from an existing ref with a new snapshot ID
+  /// \param ref The existing SnapshotRef
+  /// \param snapshot_id The new snapshot ID
+  /// \return A new SnapshotRef with the same properties but different snapshot ID
+  static std::shared_ptr<SnapshotRef> BuildRefFrom(const SnapshotRef& ref,
+                                                   int64_t snapshot_id);
+
+  /// \brief Helper to create a new branch SnapshotRef
+  /// \param snapshot_id The snapshot ID for the branch
+  /// \return A new branch SnapshotRef
+  static std::shared_ptr<SnapshotRef> BuildBranchRef(int64_t snapshot_id);
+
  private:
   // Base metadata (nullptr for new tables)
   const TableMetadata* base_;
@@ -1077,6 +1096,133 @@ int32_t TableMetadataBuilder::Impl::ReuseOrCreateNewSchemaId(
   return new_schema_id;
 }
 
+std::shared_ptr<SnapshotRef> TableMetadataBuilder::Impl::BuildRefFrom(
+    const SnapshotRef& ref, int64_t snapshot_id) {
+  auto new_ref = std::make_shared<SnapshotRef>();
+  new_ref->snapshot_id = snapshot_id;
+
+  if (std::holds_alternative<SnapshotRef::Branch>(ref.retention)) {
+    const auto& branch = std::get<SnapshotRef::Branch>(ref.retention);
+    new_ref->retention = SnapshotRef::Branch{
+        .min_snapshots_to_keep = branch.min_snapshots_to_keep,
+        .max_snapshot_age_ms = branch.max_snapshot_age_ms,
+        .max_ref_age_ms = branch.max_ref_age_ms,
+    };
+  } else {
+    const auto& tag = std::get<SnapshotRef::Tag>(ref.retention);
+    new_ref->retention = SnapshotRef::Tag{
+        .max_ref_age_ms = tag.max_ref_age_ms,
+    };
+  }
+
+  return new_ref;
+}
+
+std::shared_ptr<SnapshotRef> TableMetadataBuilder::Impl::BuildBranchRef(
+    int64_t snapshot_id) {
+  auto new_ref = std::make_shared<SnapshotRef>();
+  new_ref->snapshot_id = snapshot_id;
+  new_ref->retention = SnapshotRef::Branch{};
+  return new_ref;
+}
+
+Status TableMetadataBuilder::Impl::SetBranchSnapshotInternal(const Snapshot& snapshot,
+                                                             const std::string& branch) {
+  int64_t replacement_snapshot_id = snapshot.snapshot_id;
+
+  // Check if the ref already exists
+  auto ref_it = metadata_.refs.find(branch);
+  if (ref_it != metadata_.refs.end()) {
+    const auto& ref = ref_it->second;
+
+    // Verify that it's a branch, not a tag
+    if (ref->type() != SnapshotRefType::kBranch) {
+      return ValidationFailed("Cannot update branch: {} is a tag", branch);
+    }
+
+    // If the ref already points to the same snapshot, this is a no-op
+    if (ref->snapshot_id == replacement_snapshot_id) {
+      return {};
+    }
+  }
+
+  // For format version 2+, validate sequence number
+  if (metadata_.format_version > 1) {
+    if (snapshot.sequence_number > metadata_.last_sequence_number) {
+      return ValidationFailed(
+          "Last sequence number {} is less than existing snapshot sequence number {}",
+          metadata_.last_sequence_number, snapshot.sequence_number);
+    }
+  }
+
+  // Create the new ref
+  std::shared_ptr<SnapshotRef> new_ref;
+  if (ref_it != metadata_.refs.end()) {
+    // Build from existing ref with new snapshot ID
+    new_ref = BuildRefFrom(*ref_it->second, replacement_snapshot_id);
+  } else {
+    // Create a new branch ref
+    new_ref = BuildBranchRef(replacement_snapshot_id);
+  }
+
+  // Set the ref
+  return SetRef(branch, new_ref);
+}
+
+Status TableMetadataBuilder::Impl::SetBranchSnapshot(int64_t snapshot_id,
+                                                     const std::string& branch) {
+  // Check if this is a no-op
+  auto ref_it = metadata_.refs.find(branch);
+  if (ref_it != metadata_.refs.end() && ref_it->second->snapshot_id == snapshot_id) {
+    return {};
+  }
+
+  // Validate that the snapshot exists
+  auto snapshot_result = metadata_.SnapshotById(snapshot_id);
+  if (!snapshot_result.has_value()) {
+    return ValidationFailed("Cannot set {} to unknown snapshot: {}", branch, snapshot_id);
+  }
+
+  return SetBranchSnapshotInternal(*snapshot_result.value(), branch);
+}
+
+Status TableMetadataBuilder::Impl::SetRef(const std::string& name,
+                                          std::shared_ptr<SnapshotRef> ref) {
+  if (ref == nullptr) {
+    return InvalidArgument("Cannot set ref {} to null", name);
+  }
+
+  // Validate that the snapshot exists
+  auto snapshot_result = metadata_.SnapshotById(ref->snapshot_id);
+  if (!snapshot_result.has_value()) {
+    return ValidationFailed("Cannot set ref {} to unknown snapshot: {}", name,
+                            ref->snapshot_id);
+  }
+
+  // Update the refs map
+  metadata_.refs[name] = ref;
+
+  // If setting the main branch, also update current_snapshot_id
+  if (name == std::string(SnapshotRef::kMainBranch)) {
+    metadata_.current_snapshot_id = ref->snapshot_id;
+  }
+
+  // Record the change
+  changes_.push_back(std::make_unique<table::SetSnapshotRef>(
+      name, ref->snapshot_id, ref->type(),
+      std::holds_alternative<SnapshotRef::Branch>(ref->retention)
+          ? std::get<SnapshotRef::Branch>(ref->retention).min_snapshots_to_keep
+          : std::nullopt,
+      std::holds_alternative<SnapshotRef::Branch>(ref->retention)
+          ? std::get<SnapshotRef::Branch>(ref->retention).max_snapshot_age_ms
+          : std::nullopt,
+      ref->type() == SnapshotRefType::kBranch
+          ? std::get<SnapshotRef::Branch>(ref->retention).max_ref_age_ms
+          : std::get<SnapshotRef::Tag>(ref->retention).max_ref_age_ms));
+
+  return {};
+}
+
 TableMetadataBuilder::TableMetadataBuilder(int8_t format_version)
     : impl_(std::make_unique<Impl>(format_version)) {}
 
@@ -1212,12 +1358,14 @@ TableMetadataBuilder& TableMetadataBuilder::AddSnapshot(
 
 TableMetadataBuilder& TableMetadataBuilder::SetBranchSnapshot(int64_t snapshot_id,
                                                               const std::string& branch) {
-  throw IcebergError(std::format("{} not implemented", __FUNCTION__));
+  ICEBERG_BUILDER_RETURN_IF_ERROR(impl_->SetBranchSnapshot(snapshot_id, branch));
+  return *this;
 }
 
 TableMetadataBuilder& TableMetadataBuilder::SetRef(const std::string& name,
                                                    std::shared_ptr<SnapshotRef> ref) {
-  throw IcebergError(std::format("{} not implemented", __FUNCTION__));
+  ICEBERG_BUILDER_RETURN_IF_ERROR(impl_->SetRef(name, ref));
+  return *this;
 }
 
 TableMetadataBuilder& TableMetadataBuilder::RemoveRef(const std::string& name) {

src/iceberg/table_update.cc

@@ -344,7 +344,23 @@ std::unique_ptr<TableUpdate> RemoveSnapshotRef::Clone() const {
 // SetSnapshotRef
 
 void SetSnapshotRef::ApplyTo(TableMetadataBuilder& builder) const {
-  throw IcebergError(std::format("{} not implemented", __FUNCTION__));
+  // Create a SnapshotRef based on the type
+  auto ref = std::make_shared<SnapshotRef>();
+  ref->snapshot_id = snapshot_id_;
+
+  if (type_ == SnapshotRefType::kBranch) {
+    ref->retention = SnapshotRef::Branch{
+        .min_snapshots_to_keep = min_snapshots_to_keep_,
+        .max_snapshot_age_ms = max_snapshot_age_ms_,
+        .max_ref_age_ms = max_ref_age_ms_,
+    };
+  } else {
+    ref->retention = SnapshotRef::Tag{
+        .max_ref_age_ms = max_ref_age_ms_,
+    };
+  }
+
+  builder.SetRef(ref_name_, ref);
 }
 
 void SetSnapshotRef::GenerateRequirements(TableUpdateContext& context) const {

src/iceberg/test/CMakeLists.txt

@@ -167,6 +167,7 @@ if(ICEBERG_BUILD_BUNDLE)
   add_iceberg_test(table_update_test
                    USE_BUNDLE
                    SOURCES
+                   set_snapshot_test.cc
                    transaction_test.cc
                    update_partition_spec_test.cc
                    update_properties_test.cc