From d8c4ab15c60af508b95679565bf224072b820433 Mon Sep 17 00:00:00 2001 From: Deepak Cherian Date: Wed, 23 Oct 2024 13:39:48 -0600 Subject: [PATCH] Read-Only mode when checking out snapshot or tag (#313) * Read-Only mode when checking out snapshot or tag Closes #237 * lint * Try better docstring --- .gitignore | 5 +++- icechunk-python/python/icechunk/__init__.py | 18 +++++++++++++ .../python/icechunk/_icechunk_python.pyi | 1 + icechunk-python/src/lib.rs | 12 +++++++++ icechunk-python/tests/test_timetravel.py | 1 + icechunk/src/zarr.rs | 25 ++++++++++++++++--- 6 files changed, 57 insertions(+), 5 deletions(-) diff --git a/.gitignore b/.gitignore index e43a1f53..9e920a5f 100644 --- a/.gitignore +++ b/.gitignore @@ -4,4 +4,7 @@ .DS_Store # Docs build -.docs \ No newline at end of file +.docs + +# hypothesis +.hypothesis diff --git a/icechunk-python/python/icechunk/__init__.py b/icechunk-python/python/icechunk/__init__.py index f88e89a4..175d3c69 100644 --- a/icechunk-python/python/icechunk/__init__.py +++ b/icechunk-python/python/icechunk/__init__.py @@ -159,6 +159,24 @@ def create( store = pyicechunk_store_create(storage, config=config) return cls(store=store, mode=mode, args=args, kwargs=kwargs) + def set_mode(self, mode: AccessModeLiteral) -> None: + """ + Set the mode on this Store. + + Parameters + ---------- + mode: AccessModeLiteral + The new mode to use. + + Returns + ------- + None + + """ + read_only = mode == "r" + self._store.set_mode(read_only) + + def with_mode(self, mode: AccessModeLiteral) -> Self: """ Return a new store of the same type pointing to the same location with a new mode. diff --git a/icechunk-python/python/icechunk/_icechunk_python.pyi b/icechunk-python/python/icechunk/_icechunk_python.pyi index bb2c424d..047aa466 100644 --- a/icechunk-python/python/icechunk/_icechunk_python.pyi +++ b/icechunk-python/python/icechunk/_icechunk_python.pyi @@ -5,6 +5,7 @@ from typing import Any class PyIcechunkStore: def as_bytes(self) -> bytes: ... + def set_mode(self, read_only: bool) -> None: ... def with_mode(self, read_only: bool) -> PyIcechunkStore: ... @property def snapshot_id(self) -> str: ... diff --git a/icechunk-python/src/lib.rs b/icechunk-python/src/lib.rs index 224be2ff..b91c2ec0 100644 --- a/icechunk-python/src/lib.rs +++ b/icechunk-python/src/lib.rs @@ -329,6 +329,18 @@ impl PyIcechunkStore { Ok(Cow::Owned(serialized)) } + fn set_mode(&self, read_only: bool) -> PyResult<()> { + let access_mode = if read_only { + icechunk::zarr::AccessMode::ReadOnly + } else { + icechunk::zarr::AccessMode::ReadWrite + }; + + let mut writeable_store = self.store.blocking_write(); + writeable_store.set_mode(access_mode); + Ok(()) + } + fn with_mode(&self, read_only: bool) -> PyResult { let access_mode = if read_only { icechunk::zarr::AccessMode::ReadOnly diff --git a/icechunk-python/tests/test_timetravel.py b/icechunk-python/tests/test_timetravel.py index 6f149139..0a0801f6 100644 --- a/icechunk-python/tests/test_timetravel.py +++ b/icechunk-python/tests/test_timetravel.py @@ -30,6 +30,7 @@ def test_timetravel(): assert air_temp[200, 6] == 54 store.checkout(branch="main") + store.set_mode("w-") air_temp[:, :] = 76 assert store.has_uncommitted_changes assert store.branch == "main" diff --git a/icechunk/src/zarr.rs b/icechunk/src/zarr.rs index bf32991a..f87af0d0 100644 --- a/icechunk/src/zarr.rs +++ b/icechunk/src/zarr.rs @@ -346,6 +346,10 @@ impl Store { } } + pub fn set_mode(&mut self, mode: AccessMode) { + self.mode = mode; + } + /// Creates a new clone of the store with the given access mode. pub fn with_access_mode(&self, mode: AccessMode) -> Self { Store { @@ -396,8 +400,12 @@ impl Store { /// Checkout a specific version of the repository. This can be a snapshot id, a tag, or a branch tip. /// - /// If the version is a branch tip, the branch will be set as the current branch. If the version - /// is a tag or snapshot id, the current branch will be unset and the store will be in detached state. + /// If the version is a branch tip, the branch will be set as the current branch. + /// The current [`AccessMode`] will be unchanged. + /// + /// If the version is a tag or snapshot id, the current branch will be unset and the store + /// will be in detached state. No commits are allowed in this state. Further the store will + /// be set to [`AccessMode::ReadOnly`] mode so that any attempts to modify the store will fail. /// /// If there are uncommitted changes, this method will return an error. pub async fn checkout(&mut self, version: VersionInfo) -> StoreResult<()> { @@ -412,10 +420,12 @@ impl Store { VersionInfo::SnapshotId(sid) => { self.current_branch = None; repo.set_snapshot_id(sid); + self.mode = AccessMode::ReadOnly; } VersionInfo::TagRef(tag) => { self.current_branch = None; - repo.set_snapshot_from_tag(tag.as_str()).await? + repo.set_snapshot_from_tag(tag.as_str()).await?; + self.mode = AccessMode::ReadOnly; } VersionInfo::BranchTipRef(branch) => { self.current_branch = Some(branch.clone()); @@ -2332,11 +2342,18 @@ mod tests { let new_snapshot_id = store.commit("update").await.unwrap(); store.checkout(VersionInfo::SnapshotId(snapshot_id.clone())).await.unwrap(); + assert_eq!(store.mode, AccessMode::ReadOnly); assert_eq!(store.get("array/c/0/1/0", &ByteRange::ALL).await.unwrap(), data); - store.checkout(VersionInfo::SnapshotId(new_snapshot_id)).await.unwrap(); + store.checkout(VersionInfo::SnapshotId(new_snapshot_id.clone())).await.unwrap(); assert_eq!(store.get("array/c/0/1/0", &ByteRange::ALL).await.unwrap(), new_data); + store.tag("tag_0", &new_snapshot_id).await.unwrap(); + store.checkout(VersionInfo::TagRef("tag_0".to_string())).await.unwrap(); + assert_eq!(store.mode, AccessMode::ReadOnly); + + store.checkout(VersionInfo::BranchTipRef("main".to_string())).await.unwrap(); + store.set_mode(AccessMode::ReadWrite); let _newest_data = Bytes::copy_from_slice(b"earth"); store.set("array/c/0/1/0", data.clone()).await.unwrap(); assert_eq!(store.has_uncommitted_changes().await, true);