use a simpler approach to support internal txn in the snapshot manager

Author: wgtmac

src/iceberg/table.cc

@@ -224,10 +224,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=*/true));
-  return SnapshotManager::Make(std::move(transaction));
+  return SnapshotManager::Make(shared_from_this());
 }
 
 Result<std::shared_ptr<StagedTable>> StagedTable::Make(

src/iceberg/test/fast_append_test.cc

@@ -39,10 +39,12 @@ class FastAppendTest : public UpdateTestBase {
  protected:
   static void SetUpTestSuite() { avro::RegisterAll(); }
 
+  std::string MetadataResource() const override {
+    return "TableMetadataV2ValidMinimal.json";
+  }
+
   void SetUp() override {
-    InitializeFileIO();
-    // Use minimal metadata for FastAppend tests
-    RegisterTableFromResource("TableMetadataV2ValidMinimal.json");
+    UpdateTestBase::SetUp();
 
     // Get partition spec and schema from the base table
     ICEBERG_UNWRAP_OR_FAIL(spec_, table_->spec());

src/iceberg/test/snapshot_manager_test.cc

@@ -28,6 +28,8 @@
 
 #include "iceberg/arrow/arrow_fs_file_io_internal.h"
 #include "iceberg/catalog/memory/in_memory_catalog.h"
+#include "iceberg/result.h"
+#include "iceberg/snapshot.h"
 #include "iceberg/table.h"
 #include "iceberg/table_identifier.h"
 #include "iceberg/table_metadata.h"
@@ -37,12 +39,18 @@
 
 namespace iceberg {
 
-// Base test fixture for table update operations
+/// \brief Base test fixture for table update operations.
 class UpdateTestBase : public ::testing::Test {
  protected:
+  virtual std::string MetadataResource() const { return "TableMetadataV2Valid.json"; }
+  virtual std::string TableName() const { return "test_table"; }
+
   void SetUp() override {
+    table_ident_ = TableIdentifier{.name = TableName()};
+    table_location_ = "/warehouse/" + TableName();
+
     InitializeFileIO();
-    RegisterTableFromResource("TableMetadataV2Valid.json");
+    RegisterTableFromResource(MetadataResource());
   }
 
   /// \brief Initialize file IO and create necessary directories.
@@ -56,82 +64,91 @@ 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(table_location_ + "/metadata").ok());
   }
 
   /// \brief Register a table from a metadata resource file.
   ///
   /// \param resource_name The name of the metadata resource file
   void RegisterTableFromResource(const std::string& resource_name) {
-    // Drop existing table if it exists
     std::ignore = catalog_->DropTable(table_ident_, /*purge=*/false);
 
-    // Write table metadata to the table location.
     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 = table_location_;
     ASSERT_THAT(TableMetadataUtil::Write(*file_io_, metadata_location, *metadata),
                 IsOk());
 
-    // Register the table in the catalog.
     ICEBERG_UNWRAP_OR_FAIL(table_,
                            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_;
-};
+  /// \brief Reload the table from catalog and return its metadata.
+  std::shared_ptr<TableMetadata> ReloadMetadata() {
+    auto result = catalog_->LoadTable(table_ident_);
+    EXPECT_TRUE(result.has_value()) << "Failed to reload table";
+    return result.value()->metadata();
+  }
 
-// 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 Assert that a ref exists with the given type and snapshot id.
+  void ExpectRef(const std::string& name, SnapshotRefType type, int64_t snapshot_id) {
+    auto metadata = ReloadMetadata();
+    auto it = metadata->refs.find(name);
+    ASSERT_NE(it, metadata->refs.end()) << "Ref not found: " << name;
+    EXPECT_EQ(it->second->type(), type);
+    EXPECT_EQ(it->second->snapshot_id, snapshot_id);
   }
 
-  /// \brief Initialize file IO and create necessary directories.
-  void InitializeFileIO() {
-    file_io_ = arrow::ArrowFileSystemFileIO::MakeMockFileIO();
-    catalog_ =
-        InMemoryCatalog::Make("test_catalog", file_io_, "/warehouse/", /*properties=*/{});
+  void ExpectBranch(const std::string& name, int64_t snapshot_id) {
+    ExpectRef(name, SnapshotRefType::kBranch, snapshot_id);
+  }
 
-    // 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());
+  void ExpectTag(const std::string& name, int64_t snapshot_id) {
+    ExpectRef(name, SnapshotRefType::kTag, snapshot_id);
   }
 
-  /// \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(table_ident_, /*purge=*/false);
+  /// \brief Assert that a ref does not exist.
+  void ExpectNoRef(const std::string& name) {
+    auto metadata = ReloadMetadata();
+    EXPECT_FALSE(metadata->refs.contains(name)) << "Ref should not exist: " << name;
+  }
 
-    // Write table metadata to the table location.
-    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 = table_location_;
-    ASSERT_THAT(TableMetadataUtil::Write(*file_io_, metadata_location, *metadata),
-                IsOk());
+  /// \brief Assert the current snapshot id after reloading.
+  void ExpectCurrentSnapshot(int64_t snapshot_id) {
+    auto result = catalog_->LoadTable(table_ident_);
+    ASSERT_TRUE(result.has_value());
+    auto snap_result = result.value()->current_snapshot();
+    ASSERT_TRUE(snap_result.has_value());
+    EXPECT_EQ(snap_result.value()->snapshot_id, snapshot_id);
+  }
 
-    // Register the table in the catalog.
-    ICEBERG_UNWRAP_OR_FAIL(minimal_table_,
-                           catalog_->RegisterTable(table_ident_, metadata_location));
+  /// \brief Assert that a commit succeeded.
+  template <typename T>
+  void ExpectCommitOk(const T& result) {
+    EXPECT_THAT(result, IsOk());
+  }
+
+  /// \brief Assert that a commit failed with the given error kind and message substring.
+  template <typename T>
+  void ExpectCommitError(const T& result, ErrorKind kind, const std::string& message) {
+    EXPECT_THAT(result, IsError(kind));
+    EXPECT_THAT(result, HasErrorMessage(message));
   }
 
-  const TableIdentifier table_ident_{.name = "minimal_table"};
-  const std::string table_location_{"/warehouse/minimal_table"};
+  TableIdentifier table_ident_;
+  std::string table_location_;
   std::shared_ptr<FileIO> file_io_;
   std::shared_ptr<InMemoryCatalog> catalog_;
-  std::shared_ptr<Table> minimal_table_;
+  std::shared_ptr<Table> table_;
+};
+
+/// \brief Test fixture for table update operations on minimal table metadata.
+class MinimalUpdateTestBase : public UpdateTestBase {
+ protected:
+  std::string MetadataResource() const override {
+    return "TableMetadataV2ValidMinimal.json";
+  }
+  std::string TableName() const override { return "minimal_table"; }
 };
 
 }  // namespace iceberg

src/iceberg/test/update_test_base.h

@@ -92,8 +92,6 @@ std::string Transaction::MetadataFileLocation(std::string_view filename) const {
 }
 
 Status Transaction::AddUpdate(const std::shared_ptr<PendingUpdate>& update) {
-  ICEBERG_PRECHECK(update->kind() != PendingUpdate::Kind::kSnapshotManager,
-                   "SnapshotManager should not be added to the transaction");
   ICEBERG_CHECK(last_update_committed_,
                 "Cannot add update when previous update is not committed");
 

src/iceberg/transaction.cc

@@ -55,22 +55,13 @@ class ICEBERG_EXPORT Transaction : public std::enable_shared_from_this<Transacti
   /// \return the location of the metadata file
   std::string MetadataFileLocation(std::string_view filename) const;
 
-  /// \brief Enable auto-commit for this transaction.
-  void EnableAutoCommit() { auto_commit_ = true; }
-
-  /// \brief Disable auto-commit for this transaction.
-  void DisableAutoCommit() { auto_commit_ = false; }
-
   /// \brief Apply the pending changes from all actions and commit.
   ///
   /// \return Updated table if the transaction was committed successfully, or an error:
   /// - ValidationFailed: if any update cannot be applied to the current table metadata.
   /// - CommitFailed: if the updates cannot be committed due to conflicts.
   Result<std::shared_ptr<Table>> Commit();
 
-  /// \brief Returns true if this transaction has been committed.
-  bool is_committed() const { return committed_; }
-
   /// \brief Create a new UpdatePartitionSpec to update the partition spec of this table
   /// and commit the changes.
   Result<std::shared_ptr<UpdatePartitionSpec>> NewUpdatePartitionSpec();
@@ -109,12 +100,6 @@ class ICEBERG_EXPORT Transaction : public std::enable_shared_from_this<Transacti
   /// \brief Create a new SnapshotManager to manage snapshots.
   Result<std::shared_ptr<SnapshotManager>> NewSnapshotManager();
 
- private:
-  Transaction(std::shared_ptr<Table> table, Kind kind, bool auto_commit,
-              std::unique_ptr<TableMetadataBuilder> metadata_builder);
-
-  Status AddUpdate(const std::shared_ptr<PendingUpdate>& update);
-
   /// \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();
@@ -123,6 +108,12 @@ class ICEBERG_EXPORT Transaction : public std::enable_shared_from_this<Transacti
   /// and tags) and commit the changes.
   Result<std::shared_ptr<UpdateSnapshotReference>> NewUpdateSnapshotReference();
 
+ private:
+  Transaction(std::shared_ptr<Table> table, Kind kind, bool auto_commit,
+              std::unique_ptr<TableMetadataBuilder> metadata_builder);
+
+  Status AddUpdate(const std::shared_ptr<PendingUpdate>& update);
+
   /// \brief Apply the pending changes to current table.
   Status Apply(PendingUpdate& updates);
 
@@ -141,7 +132,6 @@ class ICEBERG_EXPORT Transaction : public std::enable_shared_from_this<Transacti
 
  private:
   friend class PendingUpdate;
-  friend class SnapshotManager;
 
   // The table that this transaction will update.
   std::shared_ptr<Table> table_;

src/iceberg/transaction.h

@@ -44,7 +44,6 @@ class ICEBERG_EXPORT PendingUpdate : public ErrorCollector {
   enum class Kind : uint8_t {
     kExpireSnapshots,
     kSetSnapshot,
-    kSnapshotManager,
     kUpdateLocation,
     kUpdatePartitionSpec,
     kUpdatePartitionStatistics,

src/iceberg/update/pending_update.h

@@ -19,11 +19,7 @@
 
 #include "iceberg/update/snapshot_manager.h"
 
-#include <memory>
-#include <string>
-
 #include "iceberg/result.h"
-#include "iceberg/snapshot.h"
 #include "iceberg/table.h"
 #include "iceberg/table_metadata.h"
 #include "iceberg/transaction.h"
@@ -34,39 +30,27 @@
 
 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<Table> table) {
+  ICEBERG_PRECHECK(table != nullptr, "Invalid input table: null");
+  ICEBERG_ASSIGN_OR_RAISE(auto transaction,
+                          Transaction::Make(std::move(table), Transaction::Kind::kUpdate,
+                                            /*auto_commit=*/false));
+  return std::shared_ptr<SnapshotManager>(
+      new SnapshotManager(std::move(transaction), /*is_external_transaction=*/false));
+}
 
 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)));
+  return std::shared_ptr<SnapshotManager>(
+      new SnapshotManager(std::move(transaction), /*is_external_transaction=*/true));
 }
 
-SnapshotManager::SnapshotManager(std::shared_ptr<Transaction> transaction)
-    : PendingUpdate(std::move(transaction)),
-      original_auto_commit_(transaction_->auto_commit_) {
-  transaction_->DisableAutoCommit();
-}
+SnapshotManager::SnapshotManager(std::shared_ptr<Transaction> transaction,
+                                 bool is_external_transaction)
+    : transaction_(std::move(transaction)),
+      is_external_transaction_(is_external_transaction) {}
 
 SnapshotManager::~SnapshotManager() = default;
 
@@ -102,13 +86,13 @@ SnapshotManager& SnapshotManager::RollbackTo(int64_t snapshot_id) {
 }
 
 SnapshotManager& SnapshotManager::CreateBranch(const std::string& name) {
-  if (base().current_snapshot_id != kInvalidSnapshotId) {
-    ICEBERG_BUILDER_ASSIGN_OR_RETURN(auto current_snapshot, base().Snapshot());
+  const auto& base = transaction_->current();
+  if (base.current_snapshot_id != kInvalidSnapshotId) {
+    ICEBERG_BUILDER_ASSIGN_OR_RETURN(auto current_snapshot, base.Snapshot());
     ICEBERG_DCHECK(current_snapshot != nullptr, "Current snapshot should not be null");
     return CreateBranch(name, current_snapshot->snapshot_id);
   }
-  const auto& current_refs = base().refs;
-  ICEBERG_BUILDER_CHECK(!base().refs.contains(name), "Ref {} already exists", name);
+  ICEBERG_BUILDER_CHECK(!base.refs.contains(name), "Ref {} already exists", name);
   ICEBERG_BUILDER_ASSIGN_OR_RETURN(auto fast_append, transaction_->NewFastAppend());
   ICEBERG_BUILDER_RETURN_IF_ERROR(fast_append->SetTargetBranch(name).Commit());
   return *this;
@@ -197,11 +181,9 @@ 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());
-  if (!transaction_->is_committed()) {
+  if (!is_external_transaction_) {
     ICEBERG_RETURN_UNEXPECTED(transaction_->Commit());
   }
   return {};