Add back safety

earth-mover · Feb 28, 2025 · ed5746e · ed5746e
1 parent 80aecfd
commit ed5746e
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 3 deletions.
diff --git a/icechunk/src/storage/object_store.rs b/icechunk/src/storage/object_store.rs
@@ -1021,6 +1021,8 @@ mod tests {
 
     use tempfile::TempDir;
 
+    use crate::format::{ChunkId, ManifestId, SnapshotId};
+
     use super::ObjectStorage;
 
     #[tokio::test]
@@ -1070,4 +1072,34 @@ mod tests {
             ObjectStorage::new_local_filesystem(PathBuf::from(&rel_path).as_path()).await;
         assert!(store.is_ok());
     }
+
+    #[tokio::test]
+    async fn test_object_store_paths() {
+        let store = ObjectStorage::new_local_filesystem(PathBuf::from(".").as_path())
+            .await
+            .unwrap();
+
+        let ref_key = "ref_key";
+        let ref_path = store.ref_key(ref_key);
+        assert_eq!(ref_path.to_string(), format!("refs/{ref_key}"));
+
+        let snapshot_id = SnapshotId::random();
+        let snapshot_path = store.get_snapshot_path(&snapshot_id);
+        assert_eq!(snapshot_path.to_string(), format!("snapshots/{snapshot_id}"));
+
+        let manifest_id = ManifestId::random();
+        let manifest_path = store.get_manifest_path(&manifest_id);
+        assert_eq!(manifest_path.to_string(), format!("manifests/{manifest_id}"));
+
+        let chunk_id = ChunkId::random();
+        let chunk_path = store.get_chunk_path(&chunk_id);
+        assert_eq!(chunk_path.to_string(), format!("chunks/{chunk_id}"));
+
+        let transaction_id = SnapshotId::random();
+        let transaction_path = store.get_transaction_path(&transaction_id);
+        assert_eq!(
+            transaction_path.to_string(),
+            format!("transactions/{transaction_id}")
+        );
+    }
 }
diff --git a/icechunk/src/storage/s3.rs b/icechunk/src/storage/s3.rs
@@ -2,7 +2,7 @@ use std::{
     fmt,
     future::ready,
     ops::Range,
-    path::Path,
+    path::{Path, PathBuf},
     sync::{
         atomic::{AtomicUsize, Ordering},
         Arc,
@@ -148,7 +148,11 @@ impl S3Storage {
     }
 
     fn get_path_str(&self, file_prefix: &str, id: &str) -> StorageResult<String> {
-        Ok(format!("{}/{}/{}", self.prefix, file_prefix, id).replace("//", "/"))
+        let path = PathBuf::from_iter([self.prefix.as_str(), file_prefix, id]);
+        let path_str =
+            path.into_os_string().into_string().map_err(StorageErrorKind::BadPrefix)?;
+
+        Ok(path_str.replace("\\", "/"))
     }
 
     fn get_path<const SIZE: usize, T: FileTypeTag>(
@@ -181,7 +185,11 @@ impl S3Storage {
     }
 
     fn ref_key(&self, ref_key: &str) -> StorageResult<String> {
-        Ok(format!("{}/{}/{}", self.prefix, REF_PREFIX, ref_key).replace("//", "/"))
+        let path = PathBuf::from_iter([self.prefix.as_str(), REF_PREFIX, ref_key]);
+        let path_str =
+            path.into_os_string().into_string().map_err(StorageErrorKind::BadPrefix)?;
+
+        Ok(path_str.replace("\\", "/"))
     }
 
     async fn get_object_reader(