Repository: earth-mover/icechunk (language: python)

Issue:
The store loses its virtual chunk resolver on checkouts



Discussion / hints:

Fix patch:
diff --git a/icechunk/src/repository.rs b/icechunk/src/repository.rs
index 990be5a4..6e16d171 100644
--- a/icechunk/src/repository.rs
+++ b/icechunk/src/repository.rs
@@ -266,6 +266,28 @@ impl Repository {
         }
     }
 
+    pub(crate) fn set_snapshot_id(&mut self, snapshot_id: SnapshotId) {
+        self.snapshot_id = snapshot_id;
+    }
+
+    pub(crate) async fn set_snapshot_from_tag(
+        &mut self,
+        tag: &str,
+    ) -> RepositoryResult<()> {
+        let ref_data = fetch_tag(self.storage.as_ref(), tag).await?;
+        self.snapshot_id = ref_data.snapshot;
+        Ok(())
+    }
+
+    pub(crate) async fn set_snapshot_from_branch(
+        &mut self,
+        branch: &str,
+    ) -> RepositoryResult<()> {
+        let ref_data = fetch_branch_tip(self.storage.as_ref(), branch).await?;
+        self.snapshot_id = ref_data.snapshot;
+        Ok(())
+    }
+
     /// Returns a pointer to the storage for the repository
     pub fn storage(&self) -> &Arc<dyn Storage + Send + Sync> {
         &self.storage
diff --git a/icechunk/src/zarr.rs b/icechunk/src/zarr.rs
index 5c9ce517..d96818f4 100644
--- a/icechunk/src/zarr.rs
+++ b/icechunk/src/zarr.rs
@@ -391,33 +391,28 @@ impl Store {
     ///
     /// If there are uncommitted changes, this method will return an error.
     pub async fn checkout(&mut self, version: VersionInfo) -> StoreResult<()> {
-        // this needs to be done carefully to avoid deadlocks and race conditions
-        let storage = {
-            let guard = self.repository.read().await;
-            // Checking out is not allowed if there are uncommitted changes
-            if guard.has_uncommitted_changes() {
-                return Err(StoreError::UncommittedChanges);
-            }
-            guard.storage().clone()
-        };
+        let mut repo = self.repository.write().await;
+
+        // Checking out is not allowed if there are uncommitted changes
+        if repo.has_uncommitted_changes() {
+            return Err(StoreError::UncommittedChanges);
+        }
 
-        let repository = match version {
+        match version {
             VersionInfo::SnapshotId(sid) => {
                 self.current_branch = None;
-                Repository::update(storage, sid)
+                repo.set_snapshot_id(sid);
             }
             VersionInfo::TagRef(tag) => {
                 self.current_branch = None;
-                Repository::from_tag(storage, &tag).await?
+                repo.set_snapshot_from_tag(tag.as_str()).await?
             }
             VersionInfo::BranchTipRef(branch) => {
                 self.current_branch = Some(branch.clone());
-                Repository::from_branch_tip(storage, &branch).await?
+                repo.set_snapshot_from_branch(&branch).await?
             }
         }
-        .build();
 
-        self.repository = Arc::new(RwLock::new(repository));
         Ok(())
     }
 


Test patch:
diff --git a/icechunk-python/tests/test_virtual_ref.py b/icechunk-python/tests/test_virtual_ref.py
index 3b92e1ee..87c4696e 100644
--- a/icechunk-python/tests/test_virtual_ref.py
+++ b/icechunk-python/tests/test_virtual_ref.py
@@ -97,7 +97,7 @@ async def test_from_s3_public_virtual_refs(tmpdir):
 
     await store.set_virtual_ref(
         "depth/c/0", 
-        "s3://noaa-nos-ofs-pds/dbofs/netcdf/202410/dbofs.t00z.20241009.regulargrid.f030.nc", 
+        "s3://noaa-nos-ofs-pds/dbofs/netcdf/202410/dbofs.t00z.20241012.regulargrid.f030.nc", 
         offset=42499, 
         length=176
     )
diff --git a/icechunk/tests/test_virtual_refs.rs b/icechunk/tests/test_virtual_refs.rs
index 73e76a97..b4bc60ff 100644
--- a/icechunk/tests/test_virtual_refs.rs
+++ b/icechunk/tests/test_virtual_refs.rs
@@ -434,7 +434,7 @@ mod tests {
 
         let ref2 = VirtualChunkRef {
             location: VirtualChunkLocation::from_absolute_path(
-                "s3://noaa-nos-ofs-pds/dbofs/netcdf/202410/dbofs.t00z.20241009.regulargrid.f030.nc",
+                "s3://noaa-nos-ofs-pds/dbofs/netcdf/202410/dbofs.t00z.20241012.regulargrid.f030.nc",
             )?,
             offset: 42499,
             length: 176,