diff --git a/icechunk-python/python/icechunk/repository.py b/icechunk-python/python/icechunk/repository.py index 8603f888f..0b10b5add 100644 --- a/icechunk-python/python/icechunk/repository.py +++ b/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, diff --git a/icechunk-python/tests/test_store.py b/icechunk-python/tests/test_store.py index 154d2f9a0..8590a71ae 100644 --- a/icechunk-python/tests/test_store.py +++ b/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"