Error when checking out invalid snapshot ID (#317)

* Error when checking out invalid snapshot ID Closes #235 * Migrate to free function. * narrow trait requirement * Revert "narrow trait requirement" This reverts commit 30ad102.
earth-mover · Oct 24, 2024 · c8fcff5 · c8fcff5
1 parent d0b9aad
commit c8fcff5
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 4 deletions.
diff --git a/icechunk/src/repository.rs b/icechunk/src/repository.rs
@@ -142,6 +142,8 @@ impl RepositoryBuilder {
 pub enum RepositoryError {
     #[error("error contacting storage {0}")]
     StorageError(#[from] StorageError),
+    #[error("snapshot not found: `{id}`")]
+    SnapshotNotFound { id: SnapshotId },
     #[error("error in icechunk file")]
     FormatError(#[from] IcechunkFormatError),
     #[error("node not found at `{path}`: {message}")]
@@ -1098,6 +1100,17 @@ async fn all_chunks<'a>(
     Ok(existing_array_chunks.chain(new_array_chunks))
 }
 
+pub async fn raise_if_invalid_snapshot_id(
+    storage: &(dyn Storage + Send + Sync),
+    snapshot_id: &SnapshotId,
+) -> RepositoryResult<()> {
+    storage
+        .fetch_snapshot(snapshot_id)
+        .await
+        .map_err(|_| RepositoryError::SnapshotNotFound { id: snapshot_id.clone() })?;
+    Ok(())
+}
+
 #[cfg(test)]
 #[allow(clippy::panic, clippy::unwrap_used, clippy::expect_used)]
 mod tests {

diff --git a/icechunk/src/zarr.rs b/icechunk/src/zarr.rs
@@ -29,9 +29,10 @@ use crate::{
     },
     refs::{update_branch, BranchVersion, Ref, RefError},
     repository::{
-        get_chunk, ArrayShape, ChunkIndices, ChunkKeyEncoding, ChunkPayload, ChunkShape,
-        Codec, DataType, DimensionNames, FillValue, Path, RepositoryError,
-        RepositoryResult, StorageTransformer, UserAttributes, ZarrArrayMetadata,
+        get_chunk, raise_if_invalid_snapshot_id, ArrayShape, ChunkIndices,
+        ChunkKeyEncoding, ChunkPayload, ChunkShape, Codec, DataType, DimensionNames,
+        FillValue, Path, RepositoryError, RepositoryResult, StorageTransformer,
+        UserAttributes, ZarrArrayMetadata,
     },
     storage::{
         s3::{S3Config, S3Storage},
@@ -419,6 +420,7 @@ impl Store {
         match version {
             VersionInfo::SnapshotId(sid) => {
                 self.current_branch = None;
+                raise_if_invalid_snapshot_id(repo.storage().as_ref(), &sid).await?;
                 repo.set_snapshot_id(sid);
                 self.mode = AccessMode::ReadOnly;
             }
@@ -472,8 +474,9 @@ impl Store {
         match self.current_branch() {
             None => Err(StoreError::NotOnBranch),
             Some(branch) => {
-                let old_snapshot = guard.snapshot_id();
                 let storage = guard.storage();
+                raise_if_invalid_snapshot_id(storage.as_ref(), &to_snapshot).await?;
+                let old_snapshot = guard.snapshot_id();
                 let overwrite = guard.config().unsafe_overwrite_refs;
                 let version = update_branch(
                     storage.as_ref(),
@@ -2341,6 +2344,13 @@ mod tests {
 
         let new_snapshot_id = store.commit("update").await.unwrap();
 
+        let random_id = SnapshotId::random();
+        let res = store.checkout(VersionInfo::SnapshotId(random_id.clone())).await;
+        assert!(matches!(
+            res,
+            Err(StoreError::RepositoryError(RepositoryError::SnapshotNotFound { .. }))
+        ));
+
         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);