Add zarr_transaction that yields a Zarr Store

icechunk-python/python/icechunk/repository.py

@@ -1,4 +1,5 @@
 import datetime
+import warnings
 from collections.abc import AsyncIterator, Iterator
 from contextlib import contextmanager
 from typing import Any, Self, cast
@@ -583,7 +584,7 @@ def writable_session(self, branch: str) -> Session:
         return Session(self._repository.writable_session(branch))
 
     @contextmanager
-    def transaction(
+    def zarr_transaction(
         self,
         branch: str,
         *,
@@ -626,6 +627,57 @@ def transaction(
             rebase_tries=rebase_tries,
         )
 
+    @contextmanager
+    def transaction(
+        self,
+        branch: str,
+        *,
+        message: str,
+        metadata: dict[str, Any] | None = None,
+        rebase_with: ConflictSolver | None = None,
+        rebase_tries: int = 1_000,
+    ) -> Iterator[IcechunkStore]:
+        """
+        Create a transaction on a branch.
+
+        This is a context manager that creates a writable session on the specified branch.
+        When the context is exited, the session will be committed to the branch
+        using the specified message.
+
+        Parameters
+        ----------
+        branch : str
+            The branch to create the transaction on.
+        message : str
+            The commit message to use when committing the session.
+        metadata : dict[str, Any] | None, optional
+            Additional metadata to store with the commit snapshot.
+        rebase_with : ConflictSolver | None, optional
+            If other session committed while the current session was writing, use Session.rebase with this solver.
+        rebase_tries : int, optional
+            If other session committed while the current session was writing, use Session.rebase up to this many times in a loop.
+
+        Yields
+        -------
+        store : IcechunkStore
+            A Zarr Store which can be used to interact with the data in the repository.
+        """
+        warnings.warn(
+            "Repository.transaction will return a writable session instead of a Zarr store "
+            "in the next Icechunk release. "
+            "Please use Repository.zarr_transaction instead. ",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        with self.zarr_transaction(
+            branch=branch,
+            message=message,
+            metadata=metadata,
+            rebase_with=rebase_with,
+            rebase_tries=rebase_tries,
+        ) as store:
+            yield store
+
     def expire_snapshots(
         self,
         older_than: datetime.datetime,

icechunk-python/tests/test_store.py

@@ -1,6 +1,7 @@
 import json
 
 import numpy as np
+import pytest
 
 import zarr
 from tests.conftest import parse_repo
@@ -101,10 +102,18 @@ async def test_transaction() -> None:
     repo = parse_repo("memory", "test")
     cid1 = repo.lookup_branch("main")
     # TODO: test metadata, rebase_with, and rebase_tries kwargs
-    with repo.transaction("main", message="initialize group") as store:
+    with pytest.warns(DeprecationWarning, match="zarr_transaction"):
+        with repo.transaction("main", message="initialize group") as store:
+            assert not store.read_only
+            root = zarr.group(store=store)
+            root.attrs["foo"] = "bar"
+    cid2 = repo.lookup_branch("main")
+    assert cid1 != cid2, "Transaction did not commit changes"
+
+    with repo.zarr_transaction("main", message="initialize group") as store:
         assert not store.read_only
         root = zarr.group(store=store)
-        root.attrs["foo"] = "bar"
+        root.attrs["foo"] = "bar2"
     cid2 = repo.lookup_branch("main")
     assert cid1 != cid2, "Transaction did not commit changes"
 
@@ -113,7 +122,7 @@ async def test_transaction_failed_no_commit() -> None:
     repo = parse_repo("memory", "test")
     cid1 = repo.lookup_branch("main")
     try:
-        with repo.transaction("main", message="initialize group") as store:
+        with repo.zarr_transaction("main", message="initialize group") as store:
             assert not store.read_only
             root = zarr.group(store=store)
             root.attrs["foo"] = "bar"