feat: implement set snapshot

src/iceberg/CMakeLists.txt

@@ -84,6 +84,7 @@ set(ICEBERG_SOURCES
     update/expire_snapshots.cc
     update/pending_update.cc
     update/snapshot_update.cc
+    update/set_snapshot.cc
     update/update_partition_spec.cc
     update/update_properties.cc
     update/update_schema.cc

src/iceberg/meson.build

@@ -104,6 +104,7 @@ iceberg_sources = files(
     'type.cc',
     'update/expire_snapshots.cc',
     'update/pending_update.cc',
+    'update/set_snapshot.cc',
     'update/snapshot_update.cc',
     'update/update_partition_spec.cc',
     'update/update_properties.cc',

src/iceberg/table.cc

@@ -192,6 +192,13 @@ Result<std::shared_ptr<ExpireSnapshots>> Table::NewExpireSnapshots() {
   return transaction->NewExpireSnapshots();
 }
 
+Result<std::shared_ptr<SetSnapshot>> Table::NewSetSnapshot() {
+  ICEBERG_ASSIGN_OR_RAISE(
+      auto transaction, Transaction::Make(shared_from_this(), Transaction::Kind::kUpdate,
+                                          /*auto_commit=*/true));
+  return transaction->NewSetSnapshot();
+}
+
 Result<std::shared_ptr<StagedTable>> StagedTable::Make(
     TableIdentifier identifier, std::shared_ptr<TableMetadata> metadata,
     std::string metadata_location, std::shared_ptr<FileIO> io,

src/iceberg/table.h

@@ -151,6 +151,10 @@ class ICEBERG_EXPORT Table : public std::enable_shared_from_this<Table> {
   /// changes.
   virtual Result<std::shared_ptr<ExpireSnapshots>> NewExpireSnapshots();
 
+  /// \brief Create a new SetSnapshot to set the current snapshot or rollback to a
+  /// previous snapshot and commit the changes.
+  virtual Result<std::shared_ptr<SetSnapshot>> NewSetSnapshot();
+
  protected:
   Table(TableIdentifier identifier, std::shared_ptr<TableMetadata> metadata,
         std::string metadata_location, std::shared_ptr<FileIO> io,

src/iceberg/test/CMakeLists.txt

@@ -171,6 +171,7 @@ if(ICEBERG_BUILD_BUNDLE)
                    USE_BUNDLE
                    SOURCES
                    expire_snapshots_test.cc
+                   set_snapshot_test.cc
                    transaction_test.cc
                    update_partition_spec_test.cc
                    update_properties_test.cc

src/iceberg/test/set_snapshot_test.cc

@@ -0,0 +1,156 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#include "iceberg/update/set_snapshot.h"
+
+#include <memory>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+
+#include "iceberg/result.h"
+#include "iceberg/snapshot.h"
+#include "iceberg/test/matchers.h"
+#include "iceberg/test/update_test_base.h"
+#include "iceberg/transaction.h"
+
+namespace iceberg {
+
+// Test fixture for SetSnapshot tests
+class SetSnapshotTest : public UpdateTestBase {
+ protected:
+  // Snapshot IDs from TableMetadataV2Valid.json
+  static constexpr int64_t kOldestSnapshotId = 3051729675574597004;
+  static constexpr int64_t kCurrentSnapshotId = 3055729675574597004;
+
+  // Timestamps from TableMetadataV2Valid.json
+  static constexpr int64_t kOldestSnapshotTimestamp = 1515100955770;
+  static constexpr int64_t kCurrentSnapshotTimestamp = 1555100955770;
+};
+
+TEST_F(SetSnapshotTest, SetCurrentSnapshotValid) {
+  ICEBERG_UNWRAP_OR_FAIL(auto set_snapshot, table_->NewSetSnapshot());
+  EXPECT_EQ(set_snapshot->kind(), PendingUpdate::Kind::kSetSnapshot);
+
+  set_snapshot->SetCurrentSnapshot(kOldestSnapshotId);
+
+  ICEBERG_UNWRAP_OR_FAIL(auto snapshot_id, set_snapshot->Apply());
+  EXPECT_EQ(snapshot_id, kOldestSnapshotId);
+
+  // Commit and verify the change was persisted
+  EXPECT_THAT(set_snapshot->Commit(), IsOk());
+  ICEBERG_UNWRAP_OR_FAIL(auto reloaded, catalog_->LoadTable(table_ident_));
+  ICEBERG_UNWRAP_OR_FAIL(auto current_snapshot, reloaded->current_snapshot());
+  EXPECT_EQ(current_snapshot->snapshot_id, kOldestSnapshotId);
+}
+
+TEST_F(SetSnapshotTest, SetCurrentSnapshotToCurrentSnapshot) {
+  ICEBERG_UNWRAP_OR_FAIL(auto set_snapshot, table_->NewSetSnapshot());
+  set_snapshot->SetCurrentSnapshot(kCurrentSnapshotId);
+
+  ICEBERG_UNWRAP_OR_FAIL(auto snapshot_id, set_snapshot->Apply());
+  EXPECT_EQ(snapshot_id, kCurrentSnapshotId);
+}
+
+TEST_F(SetSnapshotTest, SetCurrentSnapshotInvalid) {
+  ICEBERG_UNWRAP_OR_FAIL(auto set_snapshot, table_->NewSetSnapshot());
+  // Try to set to a non-existent snapshot
+  int64_t invalid_snapshot_id = 9999999999999999;
+  set_snapshot->SetCurrentSnapshot(invalid_snapshot_id);
+
+  // Should fail during Apply
+  auto result = set_snapshot->Apply();
+  EXPECT_THAT(result, IsError(ErrorKind::kValidationFailed));
+  EXPECT_THAT(result, HasErrorMessage("is not found"));
+}
+
+TEST_F(SetSnapshotTest, RollbackToValid) {
+  ICEBERG_UNWRAP_OR_FAIL(auto set_snapshot, table_->NewSetSnapshot());
+  // Rollback to the oldest snapshot (which is an ancestor)
+  set_snapshot->RollbackTo(kOldestSnapshotId);
+
+  // Apply and verify
+  ICEBERG_UNWRAP_OR_FAIL(auto snapshot_id, set_snapshot->Apply());
+  EXPECT_EQ(snapshot_id, kOldestSnapshotId);
+}
+
+TEST_F(SetSnapshotTest, RollbackToInvalidSnapshot) {
+  ICEBERG_UNWRAP_OR_FAIL(auto set_snapshot, table_->NewSetSnapshot());
+  // Try to rollback to a non-existent snapshot
+  int64_t invalid_snapshot_id = 9999999999999999;
+  set_snapshot->RollbackTo(invalid_snapshot_id);
+
+  // Should fail during Apply
+  auto result = set_snapshot->Apply();
+  EXPECT_THAT(result, IsError(ErrorKind::kValidationFailed));
+  EXPECT_THAT(result, HasErrorMessage("unknown snapshot id"));
+}
+
+TEST_F(SetSnapshotTest, RollbackToTimeValidOldestSnapshot) {
+  ICEBERG_UNWRAP_OR_FAIL(auto set_snapshot, table_->NewSetSnapshot());
+  // Rollback to a time between the two snapshots
+  // This should select the oldest snapshot
+  int64_t time_between = (kOldestSnapshotTimestamp + kCurrentSnapshotTimestamp) / 2;
+  set_snapshot->RollbackToTime(time_between);
+
+  // Apply and verify
+  ICEBERG_UNWRAP_OR_FAIL(auto snapshot_id, set_snapshot->Apply());
+  EXPECT_EQ(snapshot_id, kOldestSnapshotId);
+}
+
+TEST_F(SetSnapshotTest, RollbackToTimeBeforeAnySnapshot) {
+  ICEBERG_UNWRAP_OR_FAIL(auto set_snapshot, table_->NewSetSnapshot());
+  // Try to rollback to a time before any snapshot
+  int64_t time_before_all = kOldestSnapshotTimestamp - 1000000;
+  set_snapshot->RollbackToTime(time_before_all);
+
+  // Should fail - no snapshot older than the specified time
+  auto result = set_snapshot->Apply();
+  EXPECT_THAT(result, IsError(ErrorKind::kValidationFailed));
+  EXPECT_THAT(result, HasErrorMessage("no valid snapshot older than"));
+}
+
+TEST_F(SetSnapshotTest, RollbackToTimeExactMatch) {
+  ICEBERG_UNWRAP_OR_FAIL(auto set_snapshot, table_->NewSetSnapshot());
+  // Rollback to a timestamp just after the oldest snapshot
+  // This should return the oldest snapshot (the latest one before this time)
+  int64_t time_just_after_oldest = kOldestSnapshotTimestamp + 1;
+  set_snapshot->RollbackToTime(time_just_after_oldest);
+
+  // Apply and verify - should return the oldest snapshot
+  ICEBERG_UNWRAP_OR_FAIL(auto snapshot_id, set_snapshot->Apply());
+  EXPECT_EQ(snapshot_id, kOldestSnapshotId);
+}
+
+TEST_F(SetSnapshotTest, ApplyWithoutChanges) {
+  ICEBERG_UNWRAP_OR_FAIL(auto set_snapshot, table_->NewSetSnapshot());
+  // Apply without making any changes (NOOP)
+  ICEBERG_UNWRAP_OR_FAIL(auto snapshot_id, set_snapshot->Apply());
+
+  // Should return current snapshot
+  EXPECT_EQ(snapshot_id, kCurrentSnapshotId);
+
+  // Commit NOOP and verify nothing changed
+  EXPECT_THAT(set_snapshot->Commit(), IsOk());
+  ICEBERG_UNWRAP_OR_FAIL(auto reloaded, catalog_->LoadTable(table_ident_));
+  ICEBERG_UNWRAP_OR_FAIL(auto current_snapshot, reloaded->current_snapshot());
+  EXPECT_EQ(current_snapshot->snapshot_id, kCurrentSnapshotId);
+}
+
+}  // namespace iceberg

src/iceberg/transaction.cc

@@ -24,6 +24,7 @@
 
 #include "iceberg/catalog.h"
 #include "iceberg/schema.h"
+#include "iceberg/snapshot.h"
 #include "iceberg/table.h"
 #include "iceberg/table_metadata.h"
 #include "iceberg/table_properties.h"
@@ -32,6 +33,7 @@
 #include "iceberg/table_update.h"
 #include "iceberg/update/expire_snapshots.h"
 #include "iceberg/update/pending_update.h"
+#include "iceberg/update/set_snapshot.h"
 #include "iceberg/update/snapshot_update.h"
 #include "iceberg/update/update_partition_spec.h"
 #include "iceberg/update/update_properties.h"
@@ -183,6 +185,12 @@ Status Transaction::Apply(PendingUpdate& update) {
         metadata_builder_->RemoveSchemas(std::move(result.schema_ids_to_remove));
       }
     } break;
+    case PendingUpdate::Kind::kSetSnapshot: {
+      auto& set_snapshot = internal::checked_cast<SetSnapshot&>(update);
+      ICEBERG_ASSIGN_OR_RAISE(auto snapshot_id, set_snapshot.Apply());
+      metadata_builder_->SetBranchSnapshot(snapshot_id,
+                                           std::string(SnapshotRef::kMainBranch));
+    } break;
     default:
       return NotSupported("Unsupported pending update: {}",
                           static_cast<int32_t>(update.kind()));
@@ -280,4 +288,11 @@ Result<std::shared_ptr<ExpireSnapshots>> Transaction::NewExpireSnapshots() {
   return expire_snapshots;
 }
 
+Result<std::shared_ptr<SetSnapshot>> Transaction::NewSetSnapshot() {
+  ICEBERG_ASSIGN_OR_RAISE(std::shared_ptr<SetSnapshot> set_snapshot,
+                          SetSnapshot::Make(shared_from_this()));
+  ICEBERG_RETURN_UNEXPECTED(AddUpdate(set_snapshot));
+  return set_snapshot;
+}
+
 }  // namespace iceberg

src/iceberg/transaction.h

@@ -26,6 +26,7 @@
 #include "iceberg/iceberg_export.h"
 #include "iceberg/result.h"
 #include "iceberg/type_fwd.h"
+#include "iceberg/update/set_snapshot.h"
 
 namespace iceberg {
 
@@ -82,6 +83,10 @@ class ICEBERG_EXPORT Transaction : public std::enable_shared_from_this<Transacti
   /// changes.
   Result<std::shared_ptr<ExpireSnapshots>> NewExpireSnapshots();
 
+  /// \brief Create a new SetSnapshot to set the current snapshot or rollback to a
+  /// previous snapshot and commit the changes.
+  Result<std::shared_ptr<SetSnapshot>> NewSetSnapshot();
+
  private:
   Transaction(std::shared_ptr<Table> table, Kind kind, bool auto_commit,
               std::unique_ptr<TableMetadataBuilder> metadata_builder);

src/iceberg/type_fwd.h

@@ -194,6 +194,7 @@ class UpdateProperties;
 class UpdateSchema;
 class UpdateSortOrder;
 class ExpireSnapshots;
+class SetSnapshot;
 
 /// ----------------------------------------------------------------------------
 /// TODO: Forward declarations below are not added yet.

src/iceberg/update/meson.build

@@ -18,6 +18,7 @@
 install_headers(
     [
         'pending_update.h',
+        'set_snapshot.h',
         'snapshot_update.h',
         'update_partition_spec.h',
         'update_schema.h',

src/iceberg/update/pending_update.h

@@ -43,6 +43,7 @@ class ICEBERG_EXPORT PendingUpdate : public ErrorCollector {
  public:
   enum class Kind : uint8_t {
     kExpireSnapshots,
+    kSetSnapshot,
     kUpdatePartitionSpec,
     kUpdateProperties,
     kUpdateSchema,