fix review

Author: HeartLinked

src/iceberg/table_metadata.cc

@@ -1093,8 +1093,11 @@ Status TableMetadataBuilder::Impl::SetBranchSnapshot(int64_t snapshot_id,
     // change is a noop
     return {};
   }
-  ICEBERG_ASSIGN_OR_RAISE(auto snapshot, metadata_.SnapshotById(snapshot_id));
-  return SetBranchSnapshotInternal(*snapshot, branch);
+
+  auto snapshot_it = snapshots_by_id_.find(snapshot_id);
+  ICEBERG_CHECK(snapshot_it != snapshots_by_id_.end(),
+                "Cannot set {} to unknown snapshot: {}", branch, snapshot_id);
+  return SetBranchSnapshotInternal(*snapshot_it->second, branch);
 }
 
 Status TableMetadataBuilder::Impl::SetBranchSnapshot(std::shared_ptr<Snapshot> snapshot,

src/iceberg/update/pending_update.h

@@ -43,12 +43,12 @@ class ICEBERG_EXPORT PendingUpdate : public ErrorCollector {
  public:
   enum class Kind : uint8_t {
     kExpireSnapshots,
+    kSetSnapshot,
     kUpdatePartitionSpec,
     kUpdateProperties,
     kUpdateSchema,
     kUpdateSnapshot,
     kUpdateSortOrder,
-    kSetSnapshot,
   };
 
   /// \brief Return the kind of this pending update.

src/iceberg/update/set_snapshot.cc

@@ -46,10 +46,8 @@ SetSnapshot::SetSnapshot(std::shared_ptr<Transaction> transaction)
 SetSnapshot::~SetSnapshot() = default;
 
 SetSnapshot& SetSnapshot::SetCurrentSnapshot(int64_t snapshot_id) {
-  const TableMetadata& base_metadata = transaction_->current();
-
   // Validate that the snapshot exists
-  auto snapshot_result = base_metadata.SnapshotById(snapshot_id);
+  auto snapshot_result = base().SnapshotById(snapshot_id);
   ICEBERG_BUILDER_CHECK(snapshot_result.has_value(),
                         "Cannot roll back to unknown snapshot id: {}", snapshot_id);
 
@@ -74,10 +72,8 @@ SetSnapshot& SetSnapshot::RollbackToTime(int64_t timestamp_ms) {
 }
 
 SetSnapshot& SetSnapshot::RollbackTo(int64_t snapshot_id) {
-  const TableMetadata& current = transaction_->current();
-
   // Validate that the snapshot exists
-  auto snapshot_result = current.SnapshotById(snapshot_id);
+  auto snapshot_result = base().SnapshotById(snapshot_id);
   ICEBERG_BUILDER_CHECK(snapshot_result.has_value(),
                         "Cannot roll back to unknown snapshot id: {}", snapshot_id);
 

src/iceberg/update/set_snapshot.h

@@ -34,12 +34,6 @@
 namespace iceberg {
 
 /// \brief Sets the current snapshot directly or by rolling back.
-///
-/// This update is not exposed through the Table API. Instead, it is a package-private
-/// part of the Transaction API intended for use in ManageSnapshots.
-///
-/// When committing, these changes will be applied to the current table metadata.
-/// Commit conflicts will not be resolved and will result in a CommitFailed error.
 class ICEBERG_EXPORT SetSnapshot : public PendingUpdate {
  public:
   static Result<std::shared_ptr<SetSnapshot>> Make(