Support atomic multi-table commits through the REST catalog

mkdocs/docs/api.md

@@ -1408,6 +1408,41 @@ assert tbl.metadata.snapshots[-1].summary["abc"] == "def"
 New snapshot summaries automatically include `engine-name` (`pyiceberg`) and `engine-version`
 (the installed PyIceberg version). These values override same-named entries in `snapshot_properties`.
 
+## Multi-table commits
+
+With the REST catalog, changes to several tables can be committed atomically: either every table
+is updated, or none of them is. Open a transaction on the catalog and stage the changes of each
+table through `tx.table(...)`. The changes are committed together when the block exits, and nothing
+is committed if it raises:
+
+```python
+with catalog.transaction() as tx:
+    tx.table(orders).append(new_orders)
+    tx.table(items).append(new_items)
+```
+
+Don't commit the per-table transactions, or use them as context managers, on their own: that would
+commit that table alone. Without the context manager, stage one transaction per table and pass them
+to `commit_transactions`:
+
+```python
+orders_tx = orders.transaction()
+orders_tx.append(new_orders)
+items_tx = items.transaction()
+items_tx.append(new_items)
+
+catalog.commit_transactions(orders_tx, items_tx)
+```
+
+The commit uses the `POST /v1/{prefix}/transactions/commit` endpoint of the REST catalog
+specification, and each table is refreshed afterwards. A few things differ from committing one
+transaction:
+
+- A conflict raises `CommitFailedException` and isn't retried: nothing was applied, so stage the
+  changes again and commit again.
+- Each table can appear in only one transaction, and creating a table isn't supported.
+- Other catalogs raise `NotImplementedError`.
+
 ## Snapshot Management
 
 Manage snapshots with operations through the `Table` API:

pyiceberg/catalog/__init__.py

@@ -49,9 +49,11 @@
     DOWNCAST_NS_TIMESTAMP_TO_US_ON_WRITE,
     CommitTableResponse,
     CreateTableTransaction,
+    MultiTableTransaction,
     StagedTable,
     Table,
     TableProperties,
+    Transaction,
 )
 from pyiceberg.table.locations import load_location_provider
 from pyiceberg.table.metadata import TableMetadata, TableMetadataV1, new_table_metadata
@@ -619,6 +621,32 @@ def commit_table(
             CommitStateUnknownException: Failed due to an internal exception on the side of the catalog.
         """
 
+    def transaction(self) -> MultiTableTransaction:
+        """Create a transaction to stage changes to several tables and commit them atomically.
+
+        Returns:
+            The transaction object. Used as a context manager, it commits with
+            ``commit_transactions`` when the block exits without an exception.
+        """
+        return MultiTableTransaction(self)
+
+    def commit_transactions(self, *transactions: Transaction) -> None:
+        """Commit the changes of several table transactions atomically.
+
+        Either the changes of every transaction are applied, or none of them. Each table is
+        refreshed after a successful commit.
+
+        Args:
+            transactions: The transactions to commit, at most one per table.
+
+        Raises:
+            NotImplementedError: If the catalog doesn't support multi-table commits.
+            CommitFailedException: A requirement wasn't met, or a concurrent commit conflicted.
+            CommitStateUnknownException: The catalog failed with an internal error; the commit may
+                or may not have been applied.
+        """
+        raise NotImplementedError(f"{type(self).__name__} does not support multi-table commits")
+
     @abstractmethod
     def create_namespace(self, namespace: str | Identifier, properties: Properties = EMPTY_DICT) -> None:
         """Create a namespace in the catalog.

pyiceberg/catalog/rest/__init__.py

@@ -81,12 +81,14 @@
 from pyiceberg.table import (
     CommitTableRequest,
     CommitTableResponse,
+    CommitTransactionRequest,
     CreateTableTransaction,
     FileScanTask,
     StagedTable,
     Table,
     TableIdentifier,
     TableProperties,
+    Transaction,
 )
 from pyiceberg.table.metadata import TableMetadata
 from pyiceberg.table.sorting import UNSORTED_SORT_ORDER, SortOrder, assign_fresh_sort_order_ids
@@ -182,6 +184,7 @@ class Endpoints:
     fetch_planning_result: str = "namespaces/{namespace}/tables/{table}/plan/{plan_id}"
     cancel_planning: str = "namespaces/{namespace}/tables/{table}/plan/{plan_id}"
     fetch_scan_tasks: str = "namespaces/{namespace}/tables/{table}/tasks"
+    commit_transaction: str = "transactions/commit"
 
 
 API_PREFIX = "/v1/{prefix}"
@@ -222,6 +225,7 @@ class Capability:
         http_method=HttpMethod.DELETE, path=f"{API_PREFIX}/namespaces/{{namespace}}/tables/{{table}}/plan/{{plan-id}}"
     )
     V1_TABLE_SCAN_PLAN_TASKS = Endpoint(http_method=HttpMethod.POST, path=f"{API_PREFIX}/{Endpoints.fetch_scan_tasks}")
+    V1_COMMIT_TRANSACTION = Endpoint(http_method=HttpMethod.POST, path=f"{API_PREFIX}/{Endpoints.commit_transaction}")
 
 
 # Default endpoints for backwards compatibility with legacy servers that don't return endpoints
@@ -240,6 +244,7 @@ class Capability:
         Capability.V1_DELETE_TABLE,
         Capability.V1_RENAME_TABLE,
         Capability.V1_REGISTER_TABLE,
+        Capability.V1_COMMIT_TRANSACTION,
     )
 )
 
@@ -1638,6 +1643,103 @@ def commit_table(
             )
         return CommitTableResponse.model_validate_json(response.text)
 
+    @override
+    def commit_transactions(self, *transactions: Transaction) -> None:
+        """Commit the changes of several table transactions atomically.
+
+        Sends the staged changes of every transaction in one request to the
+        ``transactions/commit`` endpoint: either all of them are applied, or none. Each
+        table is refreshed after a successful commit. Unlike ``Transaction.commit_transaction``,
+        a conflict is not retried: re-stage the changes and commit again.
+
+        Args:
+            transactions: The transactions to commit, at most one per table. Transactions
+                without staged changes are skipped. A snapshot operation that staged nothing,
+                such as a delete that matched no rows, is still sent with a requirement that
+                the table has not changed, so a concurrent change to it fails the whole commit.
+
+        Raises:
+            ValueError: If no transaction was given, a transaction is given twice, two
+                transactions change the same table, a transaction belongs to another catalog,
+                or a transaction creates a table.
+            NotImplementedError: If the server does not support multi-table commits.
+            NoSuchTableError: If one of the tables does not exist.
+            CommitFailedException: A requirement wasn't met, or a concurrent commit conflicted.
+                Nothing was applied.
+            CommitStateUnknownException: The catalog failed with an internal error; the commit
+                may or may not have been applied.
+        """
+        if not transactions:
+            raise ValueError("At least one transaction is required")
+        if len({id(transaction) for transaction in transactions}) != len(transactions):
+            raise ValueError("A transaction can be committed only once per call")
+        identifiers: set[Identifier] = set()
+        for transaction in transactions:
+            if isinstance(transaction, CreateTableTransaction):
+                raise ValueError("Creating a table in a multi-table commit is not supported")
+            table = transaction._table  # pylint: disable=W0212
+            if table.catalog is not self:
+                raise ValueError(f"Table {table.name()} belongs to a different catalog")
+            if table.name() in identifiers:
+                raise ValueError(f"More than one transaction changes table {table.name()}")
+            identifiers.add(table.name())
+
+        staged = [
+            transaction
+            for transaction in transactions
+            if transaction._updates or transaction._snapshot_producers  # pylint: disable=W0212
+        ]
+        if not staged:
+            return
+        self._check_endpoint(Capability.V1_COMMIT_TRANSACTION)
+
+        table_changes = []
+        for transaction in staged:
+            requirements, updates = transaction._changes_for_multi_table_commit()  # pylint: disable=W0212
+            identifier = transaction._table.name()  # pylint: disable=W0212
+            table_changes.append(
+                CommitTableRequest(
+                    identifier=TableIdentifier(namespace=identifier[:-1], name=identifier[-1]),
+                    requirements=requirements,
+                    updates=updates,
+                )
+            )
+        request = CommitTransactionRequest(table_changes=table_changes)
+
+        try:
+            self._post_commit_transaction(request)
+        except (CommitFailedException, NoSuchTableError):
+            for transaction in staged:
+                transaction._multi_table_commit_failed(outcome_known=True)  # pylint: disable=W0212
+            raise
+        except Exception:
+            for transaction in staged:
+                transaction._multi_table_commit_failed(outcome_known=False)  # pylint: disable=W0212
+            raise
+
+        for transaction in staged:
+            transaction._multi_table_commit_succeeded()  # pylint: disable=W0212
+
+    @retry(**_RETRY_ARGS)
+    def _post_commit_transaction(self, request: CommitTransactionRequest) -> None:
+        response = self._session.post(
+            self.url(Endpoints.commit_transaction, prefixed=True),
+            data=request.model_dump_json().encode(UTF8),
+        )
+        try:
+            response.raise_for_status()
+        except HTTPError as exc:
+            _handle_non_200_response(
+                exc,
+                {
+                    404: NoSuchTableError,
+                    409: CommitFailedException,
+                    500: CommitStateUnknownException,
+                    502: CommitStateUnknownException,
+                    504: CommitStateUnknownException,
+                },
+            )
+
     @retry(**_RETRY_ARGS)
     @override
     def create_namespace(self, namespace: str | Identifier, properties: Properties = EMPTY_DICT) -> None:

pyiceberg/table/__init__.py

@@ -1208,6 +1208,52 @@ def commit_transaction(self) -> Table:
 
         return self._table
 
+    def _changes_for_multi_table_commit(self) -> tuple[tuple[TableRequirement, ...], tuple[TableUpdate, ...]]:
+        """Return the requirements and updates to send in a multi-table commit.
+
+        Adds the same ``AssertTableUUID`` requirement that ``commit_transaction`` adds, so the
+        commit fails if the table was dropped and recreated under the same name.
+
+        A snapshot operation that staged nothing (for example a delete whose filter matched no
+        rows) is sent as a requirement that the branch still points at the snapshot it started
+        from, with no updates. ``commit_transaction`` validates that case on the client after a
+        refresh; here the catalog checks it in the same atomic commit as the other tables.
+        """
+        if self._failed:
+            raise RuntimeError("This transaction failed to commit and cannot be reused; create a new transaction.")
+        if not any(isinstance(requirement, AssertTableUUID) for requirement in self._requirements):
+            self._requirements += (AssertTableUUID(uuid=self.table_metadata.table_uuid),)
+        if not self._updates and self._snapshot_producers:
+            producer = self._snapshot_producers[0]
+            if producer._target_branch is not None and not any(
+                isinstance(requirement, AssertRefSnapshotId) for requirement in self._requirements
+            ):
+                self._requirements += (
+                    AssertRefSnapshotId(snapshot_id=producer._starting_snapshot_id, ref=producer._target_branch),
+                )
+        return self._requirements, self._updates
+
+    def _multi_table_commit_succeeded(self) -> None:
+        """Clean up after a multi-table commit that included this transaction succeeded."""
+        self._cleanup_uncommitted_manifests()
+        self._snapshot_producers = []
+        self._updates = ()
+        self._requirements = ()
+        self._table.refresh()
+
+    def _multi_table_commit_failed(self, outcome_known: bool) -> None:
+        """Handle a multi-table commit that included this transaction and did not succeed.
+
+        Args:
+            outcome_known: True if the commit is known not to have landed, so the files written
+                for it can be deleted. False if the outcome is unknown: the files are kept,
+                because the catalog may already reference them.
+        """
+        if outcome_known:
+            for producer in self._snapshot_producers:
+                producer._clean_all_uncommitted()
+        self._failed = True
+
     def _cleanup_uncommitted_manifests(self) -> None:
         """Clean up manifests from failed retry attempts after a successful commit."""
         for producer in self._snapshot_producers:
@@ -1291,6 +1337,65 @@ def commit_transaction(self) -> Table:
         return self._table
 
 
+class MultiTableTransaction:
+    """Stage changes to several tables of a catalog and commit them atomically.
+
+    Use it as a context manager: the changes are committed with ``Catalog.commit_transactions``
+    when the block exits without an exception, and nothing is committed otherwise.
+    """
+
+    _catalog: Catalog
+    _transactions: dict[Identifier, Transaction]
+
+    def __init__(self, catalog: Catalog):
+        """Open a transaction over several tables of a catalog.
+
+        Args:
+            catalog: The catalog that holds the tables.
+        """
+        self._catalog = catalog
+        self._transactions = {}
+
+    def __enter__(self) -> MultiTableTransaction:
+        """Start a transaction to update several tables."""
+        return self
+
+    def __exit__(self, exctype: type[BaseException] | None, excinst: BaseException | None, exctb: TracebackType | None) -> None:
+        """Commit the staged changes if no exceptions have been raised."""
+        if exctype is None and excinst is None and exctb is None:
+            self.commit_transaction()
+
+    def table(self, table: Table) -> Transaction:
+        """Return the transaction that stages the changes to a table.
+
+        The same transaction is returned each time for the same table. Don't commit it or use it
+        as a context manager on its own: its changes are committed with the other tables.
+
+        Args:
+            table: The table to change.
+
+        Returns:
+            The transaction for the table.
+        """
+        if table.catalog is not self._catalog:
+            raise ValueError(f"Table {table.name()} belongs to a different catalog")
+        if (transaction := self._transactions.get(table.name())) is None:
+            transaction = self._transactions[table.name()] = table.transaction()
+        return transaction
+
+    def commit_transaction(self) -> None:
+        """Commit the changes of every table atomically.
+
+        Raises:
+            NotImplementedError: If the catalog doesn't support multi-table commits.
+            CommitFailedException: A requirement wasn't met, or a concurrent commit conflicted.
+            CommitStateUnknownException: The catalog failed with an internal error; the commit may
+                or may not have been applied.
+        """
+        if self._transactions:
+            self._catalog.commit_transactions(*self._transactions.values())
+
+
 class Namespace(IcebergRootModel[list[str]]):
     """Reference to one or more levels of a namespace."""
 
@@ -1315,6 +1420,12 @@ class CommitTableRequest(IcebergBaseModel):
     updates: tuple[TableUpdate, ...] = Field(default_factory=tuple)
 
 
+class CommitTransactionRequest(IcebergBaseModel):
+    """A pydantic BaseModel for a multi-table commit request."""
+
+    table_changes: list[CommitTableRequest] = Field(alias="table-changes")
+
+
 class CommitTableResponse(IcebergBaseModel):
     """A pydantic BaseModel for a table commit response."""
 

tests/catalog/test_base.py

@@ -88,3 +88,8 @@ def test_in_memory_catalog_context_manager(self, catalog: InMemoryCatalog) -> No
 
         # InMemoryCatalog inherits close from SqlCatalog, so engine should be disposed
         assert hasattr(cat, "engine")
+
+
+def test_commit_transactions_not_supported(catalog: InMemoryCatalog) -> None:
+    with pytest.raises(NotImplementedError, match="does not support multi-table commits"):
+        catalog.commit_transactions()