fix: reject tag refs for main snapshot ref

src/iceberg/table_metadata.cc

@@ -1067,7 +1067,6 @@ Status TableMetadataBuilder::Impl::AddSnapshot(std::shared_ptr<Snapshot> snapsho
       "Cannot add snapshot with sequence number {} older than last sequence number {}",
       snapshot->sequence_number, metadata_.last_sequence_number);
 
-  metadata_.last_updated_ms = snapshot->timestamp_ms;
   metadata_.last_sequence_number = snapshot->sequence_number;
   metadata_.snapshots.push_back(snapshot);
   snapshots_by_id_.emplace(snapshot->snapshot_id, snapshot);
@@ -1156,22 +1155,26 @@ Status TableMetadataBuilder::Impl::SetRef(const std::string& name,
                 "Cannot set {} to unknown snapshot: {}", name, snapshot_id);
   const auto& snapshot = snapshot_it->second;
 
-  // If snapshot was added in this set of changes, update last_updated_ms
-  if (std::ranges::any_of(changes_, [snapshot_id](const auto& change) {
-        return change->kind() == TableUpdate::Kind::kAddSnapshot &&
-               internal::checked_cast<const table::AddSnapshot&>(*change)
-                       .snapshot()
-                       ->snapshot_id == snapshot_id;
-      })) {
-    metadata_.last_updated_ms = snapshot->timestamp_ms;
-  }
+  ICEBERG_CHECK(
+      name != SnapshotRef::kMainBranch || ref->type() == SnapshotRefType::kBranch,
+      "Cannot set {} to a tag, it must be a branch", SnapshotRef::kMainBranch);
 
   if (name == SnapshotRef::kMainBranch) {
+    const bool is_added_snapshot =
+        std::ranges::any_of(changes_, [snapshot_id](const auto& change) {
+          return change->kind() == TableUpdate::Kind::kAddSnapshot &&
+                 internal::checked_cast<const table::AddSnapshot&>(*change)
+                         .snapshot()
+                         ->snapshot_id == snapshot_id;
+        });
     metadata_.current_snapshot_id = ref->snapshot_id;
     if (metadata_.last_updated_ms == kInvalidLastUpdatedMs) {
       metadata_.last_updated_ms = CurrentTimePointMs();
     }
-    metadata_.snapshot_log.emplace_back(metadata_.last_updated_ms, ref->snapshot_id);
+
+    auto time_of_change =
+        is_added_snapshot ? snapshot->timestamp_ms : metadata_.last_updated_ms;
+    metadata_.snapshot_log.emplace_back(time_of_change, ref->snapshot_id);
   }
 
   changes_.push_back(std::make_unique<table::SetSnapshotRef>(name, *ref));

src/iceberg/test/snapshot_manager_test.cc

@@ -421,6 +421,25 @@ TEST_F(SnapshotManagerMinimalTableTest, CreateBranchOnEmptyTable) {
   EXPECT_EQ(it->second->type(), SnapshotRefType::kBranch);
 }
 
+TEST_F(SnapshotManagerMinimalTableTest, CreateTagNamedMainFails) {
+  ICEBERG_UNWRAP_OR_FAIL(auto manager, table_->NewSnapshotManager());
+  manager->CreateBranch("branch1");
+  ExpectCommitOk(manager->Commit());
+
+  auto metadata = ReloadMetadata();
+  auto branch_it = metadata->refs.find("branch1");
+  ASSERT_NE(branch_it, metadata->refs.end());
+  ASSERT_EQ(branch_it->second->type(), SnapshotRefType::kBranch);
+
+  ICEBERG_UNWRAP_OR_FAIL(auto table_with_branch, catalog_->LoadTable(table_ident_));
+  ICEBERG_UNWRAP_OR_FAIL(auto new_manager, table_with_branch->NewSnapshotManager());
+  new_manager->CreateTag(std::string(SnapshotRef::kMainBranch),
+                         branch_it->second->snapshot_id);
+  ExpectCommitError(new_manager->Commit(), ErrorKind::kValidationFailed,
+                    "Cannot set main to a tag, it must be a branch");
+  ExpectNoRef(std::string(SnapshotRef::kMainBranch));
+}
+
 TEST_F(SnapshotManagerMinimalTableTest,
        CreateBranchOnEmptyTableFailsWhenRefAlreadyExists) {
   ICEBERG_UNWRAP_OR_FAIL(auto manager, table_->NewSnapshotManager());

src/iceberg/test/table_metadata_builder_test.cc

@@ -36,6 +36,7 @@
 #include "iceberg/test/matchers.h"
 #include "iceberg/transform.h"
 #include "iceberg/type.h"
+#include "iceberg/util/timepoint.h"
 #include "iceberg/util/uuid.h"
 
 namespace iceberg {
@@ -1159,6 +1160,39 @@ TEST(TableMetadataBuilderTest, RemoveSnapshotRef) {
   EXPECT_TRUE(metadata->refs.contains("ref1"));
 }
 
+TEST(TableMetadataBuilderTest, SetRefRejectsTagForMainBranch) {
+  auto base = CreateBaseMetadata();
+  auto builder = TableMetadataBuilder::BuildFrom(base.get());
+
+  builder->AddSnapshot(std::make_shared<Snapshot>(Snapshot{.snapshot_id = 1}));
+  ICEBERG_UNWRAP_OR_FAIL(auto main_tag, SnapshotRef::MakeTag(1));
+
+  builder->SetRef(std::string(SnapshotRef::kMainBranch), std::move(main_tag));
+
+  auto result = builder->Build();
+  ASSERT_THAT(result, IsError(ErrorKind::kValidationFailed));
+  EXPECT_THAT(result, HasErrorMessage("Cannot set main to a tag, it must be a branch"));
+}
+
+TEST(TableMetadataBuilderTest, SetMainRefToAddedSnapshotUsesSnapshotTimestampForLog) {
+  auto base = CreateBaseMetadata();
+  auto builder = TableMetadataBuilder::BuildFrom(base.get());
+
+  auto snapshot_time = TimePointMsFromUnixMs(123456789);
+  builder->AddSnapshot(std::make_shared<Snapshot>(
+      Snapshot{.snapshot_id = 1, .sequence_number = 1, .timestamp_ms = snapshot_time}));
+  ICEBERG_UNWRAP_OR_FAIL(auto main_branch, SnapshotRef::MakeBranch(1));
+
+  builder->SetRef(std::string(SnapshotRef::kMainBranch), std::move(main_branch));
+
+  ICEBERG_UNWRAP_OR_FAIL(auto metadata, builder->Build());
+  EXPECT_EQ(metadata->current_snapshot_id, 1);
+  EXPECT_NE(metadata->last_updated_ms, snapshot_time);
+  ASSERT_FALSE(metadata->snapshot_log.empty());
+  EXPECT_EQ(metadata->snapshot_log.back().snapshot_id, 1);
+  EXPECT_EQ(metadata->snapshot_log.back().timestamp_ms, snapshot_time);
+}
+
 TEST(TableMetadataBuilderTest, RemoveSnapshot) {
   auto base = CreateBaseMetadata();
   auto builder = TableMetadataBuilder::BuildFrom(base.get());

src/iceberg/test/table_update_test.cc

@@ -440,4 +440,24 @@ TEST(TableUpdateTest, SetSnapshotRefApplyUpdate) {
   }
 }
 
+TEST(TableUpdateTest, SetSnapshotRefRejectsTagForMainBranch) {
+  auto base = CreateBaseMetadata();
+  auto builder = TableMetadataBuilder::BuildFrom(base.get());
+
+  auto snapshot = std::make_shared<Snapshot>(
+      Snapshot{.snapshot_id = 987654321,
+               .sequence_number = 1,
+               .timestamp_ms = TimePointMsFromUnixMs(2000000),
+               .manifest_list = "s3://bucket/manifest-list.avro"});
+  builder->AddSnapshot(snapshot);
+
+  table::SetSnapshotRef update(std::string(SnapshotRef::kMainBranch), 987654321,
+                               SnapshotRefType::kTag);
+  update.ApplyTo(*builder);
+
+  auto result = builder->Build();
+  ASSERT_THAT(result, IsError(ErrorKind::kValidationFailed));
+  EXPECT_THAT(result, HasErrorMessage("Cannot set main to a tag, it must be a branch"));
+}
+
 }  // namespace iceberg

src/iceberg/transaction.cc

@@ -209,13 +209,15 @@ Status Transaction::ApplyExpireSnapshots(ExpireSnapshots& update) {
   if (!result.schema_ids_to_remove.empty()) {
     ctx_->metadata_builder->RemoveSchemas(std::move(result.schema_ids_to_remove));
   }
+  ICEBERG_RETURN_UNEXPECTED(ctx_->metadata_builder->CheckErrors());
   return {};
 }
 
 Status Transaction::ApplySetSnapshot(SetSnapshot& update) {
   ICEBERG_ASSIGN_OR_RAISE(auto snapshot_id, update.Apply());
   ctx_->metadata_builder->SetBranchSnapshot(snapshot_id,
                                             std::string(SnapshotRef::kMainBranch));
+  ICEBERG_RETURN_UNEXPECTED(ctx_->metadata_builder->CheckErrors());
   return {};
 }
 
@@ -232,6 +234,7 @@ Status Transaction::ApplyUpdatePartitionSpec(UpdatePartitionSpec& update) {
   } else {
     ctx_->metadata_builder->AddPartitionSpec(std::move(result.spec));
   }
+  ICEBERG_RETURN_UNEXPECTED(ctx_->metadata_builder->CheckErrors());
   return {};
 }
 
@@ -246,6 +249,7 @@ Status Transaction::ApplyUpdateProperties(UpdateProperties& update) {
   if (result.format_version.has_value()) {
     ctx_->metadata_builder->UpgradeFormatVersion(result.format_version.value());
   }
+  ICEBERG_RETURN_UNEXPECTED(ctx_->metadata_builder->CheckErrors());
   return {};
 }
 
@@ -256,6 +260,7 @@ Status Transaction::ApplyUpdateSchema(UpdateSchema& update) {
   if (!result.updated_props.empty()) {
     ctx_->metadata_builder->SetProperties(result.updated_props);
   }
+  ICEBERG_RETURN_UNEXPECTED(ctx_->metadata_builder->CheckErrors());
 
   return {};
 }
@@ -275,6 +280,7 @@ Status Transaction::ApplyUpdateSnapshot(SnapshotUpdate& update) {
   } else {
     temp_update->SetBranchSnapshot(std::move(result.snapshot), result.target_branch);
   }
+  ICEBERG_RETURN_UNEXPECTED(temp_update->CheckErrors());
 
   if (temp_update->changes().empty()) {
     // Do not commit if the metadata has not changed. for example, this may happen
@@ -293,6 +299,7 @@ Status Transaction::ApplyUpdateSnapshot(SnapshotUpdate& update) {
   if (base.table_uuid.empty()) {
     ctx_->metadata_builder->AssignUUID();
   }
+  ICEBERG_RETURN_UNEXPECTED(ctx_->metadata_builder->CheckErrors());
   return {};
 }
 
@@ -304,12 +311,14 @@ Status Transaction::ApplyUpdateSnapshotReference(UpdateSnapshotReference& update
   for (auto&& [name, ref] : result.to_set) {
     ctx_->metadata_builder->SetRef(std::move(name), std::move(ref));
   }
+  ICEBERG_RETURN_UNEXPECTED(ctx_->metadata_builder->CheckErrors());
   return {};
 }
 
 Status Transaction::ApplyUpdateSortOrder(UpdateSortOrder& update) {
   ICEBERG_ASSIGN_OR_RAISE(auto sort_order, update.Apply());
   ctx_->metadata_builder->SetDefaultSortOrder(std::move(sort_order));
+  ICEBERG_RETURN_UNEXPECTED(ctx_->metadata_builder->CheckErrors());
   return {};
 }
 
@@ -321,6 +330,7 @@ Status Transaction::ApplyUpdateStatistics(UpdateStatistics& update) {
   for (const auto& snapshot_id : result.to_remove) {
     ctx_->metadata_builder->RemoveStatistics(snapshot_id);
   }
+  ICEBERG_RETURN_UNEXPECTED(ctx_->metadata_builder->CheckErrors());
   return {};
 }
 
@@ -332,6 +342,7 @@ Status Transaction::ApplyUpdatePartitionStatistics(UpdatePartitionStatistics& up
   for (const auto& snapshot_id : result.to_remove) {
     ctx_->metadata_builder->RemovePartitionStatistics(snapshot_id);
   }
+  ICEBERG_RETURN_UNEXPECTED(ctx_->metadata_builder->CheckErrors());
   return {};
 }