From 621b7afb6fd1e6e02c21540e85edc90a4f0457a4 Mon Sep 17 00:00:00 2001
From: Amogh Jahagirdar <amoghj@apache.org>
Date: Wed, 18 Dec 2024 18:30:52 -0700
Subject: [PATCH] Core: Reset to main as part of replace shouldn't remove main
 ref

---
 core/src/main/java/org/apache/iceberg/MetadataUpdate.java | 6 +++++-
 core/src/main/java/org/apache/iceberg/TableMetadata.java  | 2 +-
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/core/src/main/java/org/apache/iceberg/MetadataUpdate.java b/core/src/main/java/org/apache/iceberg/MetadataUpdate.java
index ba038c196e43..65282302b5f9 100644
--- a/core/src/main/java/org/apache/iceberg/MetadataUpdate.java
+++ b/core/src/main/java/org/apache/iceberg/MetadataUpdate.java
@@ -338,7 +338,11 @@ public String name() {
 
     @Override
     public void applyTo(TableMetadata.Builder metadataBuilder) {
-      metadataBuilder.removeRef(refName);
+      if (refName.equals(SnapshotRef.MAIN_BRANCH)) {
+        metadataBuilder.resetMainBranch();
+      } else {
+        metadataBuilder.removeRef(refName);
+      }
     }
   }
 
diff --git a/core/src/main/java/org/apache/iceberg/TableMetadata.java b/core/src/main/java/org/apache/iceberg/TableMetadata.java
index 9f6ffbcc8714..77bbfaff64e0 100644
--- a/core/src/main/java/org/apache/iceberg/TableMetadata.java
+++ b/core/src/main/java/org/apache/iceberg/TableMetadata.java
@@ -1297,7 +1297,7 @@ public Builder removeRef(String name) {
       return this;
     }
 
-    private Builder resetMainBranch() {
+    protected Builder resetMainBranch() {
       this.currentSnapshotId = -1;
       SnapshotRef ref = refs.remove(SnapshotRef.MAIN_BRANCH);
       if (ref != null) {