fix: handle auto commit logic inside SnapshotManager

zhjwpku · zhjwpku · commit 6a644a0ea1bd · 2026-02-13T20:19:38.000+08:00
diff --git a/src/iceberg/table.cc b/src/iceberg/table.cc
@@ -226,7 +226,7 @@ Result<std::shared_ptr<UpdatePartitionStatistics>> Table::NewUpdatePartitionStat
 Result<std::shared_ptr<SnapshotManager>> Table::NewSnapshotManager() {
   ICEBERG_ASSIGN_OR_RAISE(
       auto transaction, Transaction::Make(shared_from_this(), Transaction::Kind::kUpdate,
-                                          /*auto_commit=*/false));
+                                          /*auto_commit=*/true));
   return SnapshotManager::Make(std::move(transaction));
 }
 
diff --git a/src/iceberg/test/snapshot_manager_test.cc b/src/iceberg/test/snapshot_manager_test.cc
@@ -50,6 +50,8 @@ class SnapshotManagerTest : public UpdateTestBase {
   int64_t oldest_snapshot_id_{};
 };
 
+class SnapshotManagerMinimalTableTest : public MinimalUpdateTestBase {};
+
 TEST_F(SnapshotManagerTest, CreateBranch) {
   ICEBERG_UNWRAP_OR_FAIL(auto manager, table_->NewSnapshotManager());
   manager->CreateBranch("branch1", current_snapshot_id_);
@@ -76,12 +78,12 @@ TEST_F(SnapshotManagerTest, CreateBranchWithoutSnapshotId) {
   EXPECT_EQ(ref->snapshot_id, current_snapshot_id_);
 }
 
-TEST_F(SnapshotManagerTest, CreateBranchOnEmptyTable) {
+TEST_F(SnapshotManagerMinimalTableTest, CreateBranchOnEmptyTable) {
   ICEBERG_UNWRAP_OR_FAIL(auto manager, minimal_table_->NewSnapshotManager());
   manager->CreateBranch("branch1");
   EXPECT_THAT(manager->Commit(), IsOk());
 
-  ICEBERG_UNWRAP_OR_FAIL(auto reloaded, catalog_->LoadTable(minimal_table_ident_));
+  ICEBERG_UNWRAP_OR_FAIL(auto reloaded, catalog_->LoadTable(table_ident_));
   EXPECT_FALSE(
       reloaded->metadata()->refs.contains(std::string(SnapshotRef::kMainBranch)));
   auto it = reloaded->metadata()->refs.find("branch1");
@@ -90,13 +92,13 @@ TEST_F(SnapshotManagerTest, CreateBranchOnEmptyTable) {
   EXPECT_EQ(ref->type(), SnapshotRefType::kBranch);
 }
 
-TEST_F(SnapshotManagerTest, CreateBranchOnEmptyTableFailsWhenRefAlreadyExists) {
+TEST_F(SnapshotManagerMinimalTableTest,
+       CreateBranchOnEmptyTableFailsWhenRefAlreadyExists) {
   ICEBERG_UNWRAP_OR_FAIL(auto manager, minimal_table_->NewSnapshotManager());
   manager->CreateBranch("branch1");
   EXPECT_THAT(manager->Commit(), IsOk());
 
-  ICEBERG_UNWRAP_OR_FAIL(auto table_with_branch,
-                         catalog_->LoadTable(minimal_table_ident_));
+  ICEBERG_UNWRAP_OR_FAIL(auto table_with_branch, catalog_->LoadTable(table_ident_));
   ICEBERG_UNWRAP_OR_FAIL(auto manager2, table_with_branch->NewSnapshotManager());
   manager2->CreateBranch("branch1");
   auto result = manager2->Commit();
@@ -509,4 +511,32 @@ TEST_F(SnapshotManagerTest, SnapshotManagerThroughTransaction) {
   EXPECT_EQ(current_snapshot->snapshot_id, oldest_snapshot_id_);
 }
 
+TEST_F(SnapshotManagerTest, SnapshotManagerFromTableAllowsMultipleSnapshotOperations) {
+  ICEBERG_UNWRAP_OR_FAIL(auto manager, table_->NewSnapshotManager());
+
+  manager->SetCurrentSnapshot(oldest_snapshot_id_);
+  manager->SetCurrentSnapshot(current_snapshot_id_);
+  manager->RollbackTo(oldest_snapshot_id_);
+  EXPECT_THAT(manager->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, oldest_snapshot_id_);
+}
+
+TEST_F(SnapshotManagerTest,
+       SnapshotManagerFromTransactionAllowsMultipleSnapshotOperations) {
+  ICEBERG_UNWRAP_OR_FAIL(auto txn, table_->NewTransaction());
+  ICEBERG_UNWRAP_OR_FAIL(auto manager, SnapshotManager::Make(txn));
+
+  manager->SetCurrentSnapshot(oldest_snapshot_id_);
+  manager->SetCurrentSnapshot(current_snapshot_id_);
+  manager->RollbackTo(oldest_snapshot_id_);
+  EXPECT_THAT(txn->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, oldest_snapshot_id_);
+}
+
 }  // namespace iceberg
diff --git a/src/iceberg/test/update_test_base.h b/src/iceberg/test/update_test_base.h
@@ -43,7 +43,6 @@ class UpdateTestBase : public ::testing::Test {
   void SetUp() override {
     InitializeFileIO();
     RegisterTableFromResource("TableMetadataV2Valid.json");
-    RegisterMinimalTableFromResource("TableMetadataV2ValidMinimal.json");
   }
 
   /// \brief Initialize file IO and create necessary directories.
@@ -57,7 +56,7 @@ class UpdateTestBase : public ::testing::Test {
         static_cast<arrow::ArrowFileSystemFileIO&>(*file_io_).fs());
     ASSERT_TRUE(arrow_fs != nullptr);
     ASSERT_TRUE(arrow_fs->CreateDir(table_location_ + "/metadata").ok());
-    ASSERT_TRUE(arrow_fs->CreateDir(minimal_table_location_ + "/metadata").ok());
+    ASSERT_TRUE(arrow_fs->CreateDir(table_location_ + "/metadata").ok());
   }
 
   /// \brief Register a table from a metadata resource file.
@@ -80,34 +79,58 @@ class UpdateTestBase : public ::testing::Test {
                            catalog_->RegisterTable(table_ident_, metadata_location));
   }
 
+  const TableIdentifier table_ident_{.name = "test_table"};
+  const std::string table_location_{"/warehouse/test_table"};
+  std::shared_ptr<FileIO> file_io_;
+  std::shared_ptr<InMemoryCatalog> catalog_;
+  std::shared_ptr<Table> table_;
+};
+
+// Base test fixture for table update operations on minimal table metadata.
+class MinimalUpdateTestBase : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    InitializeFileIO();
+    RegisterMinimalTableFromResource("TableMetadataV2ValidMinimal.json");
+  }
+
+  /// \brief Initialize file IO and create necessary directories.
+  void InitializeFileIO() {
+    file_io_ = arrow::ArrowFileSystemFileIO::MakeMockFileIO();
+    catalog_ =
+        InMemoryCatalog::Make("test_catalog", file_io_, "/warehouse/", /*properties=*/{});
+
+    // Arrow MockFS cannot automatically create directories.
+    auto arrow_fs = std::dynamic_pointer_cast<::arrow::fs::internal::MockFileSystem>(
+        static_cast<arrow::ArrowFileSystemFileIO&>(*file_io_).fs());
+    ASSERT_TRUE(arrow_fs != nullptr);
+    ASSERT_TRUE(arrow_fs->CreateDir(table_location_ + "/metadata").ok());
+  }
+
   /// \brief Register a minimal table from a metadata resource file.
   ///
   /// \param resource_name The name of the metadata resource file
   void RegisterMinimalTableFromResource(const std::string& resource_name) {
     // Drop existing table if it exists
-    std::ignore = catalog_->DropTable(minimal_table_ident_, /*purge=*/false);
+    std::ignore = catalog_->DropTable(table_ident_, /*purge=*/false);
 
     // Write table metadata to the table location.
-    auto metadata_location =
-        std::format("{}/metadata/00001-{}.metadata.json", minimal_table_location_,
-                    Uuid::GenerateV7().ToString());
+    auto metadata_location = std::format("{}/metadata/00001-{}.metadata.json",
+                                         table_location_, Uuid::GenerateV7().ToString());
     ICEBERG_UNWRAP_OR_FAIL(auto metadata, ReadTableMetadataFromResource(resource_name));
-    metadata->location = minimal_table_location_;
+    metadata->location = table_location_;
     ASSERT_THAT(TableMetadataUtil::Write(*file_io_, metadata_location, *metadata),
                 IsOk());
 
     // Register the table in the catalog.
-    ICEBERG_UNWRAP_OR_FAIL(
-        minimal_table_, catalog_->RegisterTable(minimal_table_ident_, metadata_location));
+    ICEBERG_UNWRAP_OR_FAIL(minimal_table_,
+                           catalog_->RegisterTable(table_ident_, metadata_location));
   }
 
-  const TableIdentifier table_ident_{.name = "test_table"};
-  const std::string table_location_{"/warehouse/test_table"};
-  const TableIdentifier minimal_table_ident_{.name = "minimal_table"};
-  const std::string minimal_table_location_{"/warehouse/minimal_table"};
+  const TableIdentifier table_ident_{.name = "minimal_table"};
+  const std::string table_location_{"/warehouse/minimal_table"};
   std::shared_ptr<FileIO> file_io_;
   std::shared_ptr<InMemoryCatalog> catalog_;
-  std::shared_ptr<Table> table_;
   std::shared_ptr<Table> minimal_table_;
 };
 
diff --git a/src/iceberg/update/snapshot_manager.cc b/src/iceberg/update/snapshot_manager.cc
@@ -34,14 +34,39 @@
 
 namespace iceberg {
 
+namespace {
+
+class AutoCommitGuard {
+ public:
+  AutoCommitGuard(std::shared_ptr<Transaction> transaction, bool auto_commit)
+      : transaction_(std::move(transaction)), auto_commit_(auto_commit) {}
+
+  ~AutoCommitGuard() {
+    if (auto_commit_) {
+      transaction_->EnableAutoCommit();
+    } else {
+      transaction_->DisableAutoCommit();
+    }
+  }
+
+ private:
+  std::shared_ptr<Transaction> transaction_;
+  bool auto_commit_;
+};
+
+}  // namespace
+
 Result<std::shared_ptr<SnapshotManager>> SnapshotManager::Make(
     std::shared_ptr<Transaction> transaction) {
   ICEBERG_PRECHECK(transaction != nullptr, "Invalid input transaction: null");
   return std::shared_ptr<SnapshotManager>(new SnapshotManager(std::move(transaction)));
 }
 
 SnapshotManager::SnapshotManager(std::shared_ptr<Transaction> transaction)
-    : PendingUpdate(std::move(transaction)) {}
+    : PendingUpdate(std::move(transaction)),
+      original_auto_commit_(transaction_->auto_commit_) {
+  transaction_->DisableAutoCommit();
+}
 
 SnapshotManager::~SnapshotManager() = default;
 
@@ -172,6 +197,7 @@ SnapshotManager& SnapshotManager::SetMaxRefAgeMs(const std::string& name,
 }
 
 Status SnapshotManager::Commit() {
+  AutoCommitGuard auto_commit_guard(transaction_, original_auto_commit_);
   transaction_->EnableAutoCommit();
   ICEBERG_RETURN_UNEXPECTED(CheckErrors());
   ICEBERG_RETURN_UNEXPECTED(CommitIfRefUpdatesExist());
diff --git a/src/iceberg/update/snapshot_manager.h b/src/iceberg/update/snapshot_manager.h
@@ -186,6 +186,7 @@ class ICEBERG_EXPORT SnapshotManager : public PendingUpdate {
   /// \brief Commit any pending reference updates if they exist.
   Status CommitIfRefUpdatesExist();
 
+  bool original_auto_commit_;
   std::shared_ptr<UpdateSnapshotReference> update_snap_refs_;
 };
 

Original file line number	Diff line number	Diff line change
`@@ -226,7 +226,7 @@ Result<std::shared_ptr<UpdatePartitionStatistics>> Table::NewUpdatePartitionStat`
`226`	`226`	`Result<std::shared_ptr<SnapshotManager>> Table::NewSnapshotManager() {`
`227`	`227`	`ICEBERG_ASSIGN_OR_RAISE(`
`228`	`228`	`auto transaction, Transaction::Make(shared_from_this(), Transaction::Kind::kUpdate,`
`229`		`- /auto_commit=/false));`
	`229`	`+ /auto_commit=/true));`
`230`	`230`	`return SnapshotManager::Make(std::move(transaction));`
`231`	`231`	`}`
`232`	`232`