From c597fc14492ad2ceb133877c5f5907c64a0dcaf6 Mon Sep 17 00:00:00 2001
From: Sebastian Galkin <code@amisdelabc.com>
Date: Wed, 9 Oct 2024 22:25:30 -0300
Subject: [PATCH] Fix `list_prefix`

This seems to be the behavior zarr-python expects.

Closes #168
---
 icechunk/src/zarr.rs | 59 ++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 57 insertions(+), 2 deletions(-)

diff --git a/icechunk/src/zarr.rs b/icechunk/src/zarr.rs
index 88dc837d..0d41007c 100644
--- a/icechunk/src/zarr.rs
+++ b/icechunk/src/zarr.rs
@@ -823,8 +823,20 @@ impl Store {
             for node in repository.list_nodes().await? {
                 // TODO: handle non-utf8?
                 let meta_key = Key::Metadata { node_path: node.path }.to_string();
-                if meta_key.starts_with(prefix) {
-                    yield meta_key;
+                    match meta_key.strip_prefix(prefix) {
+                        None => {}
+                        Some(rest) => {
+                            // we have a few cases
+                            if prefix.is_empty()   // if prefix was empty anything matches
+                               || rest.is_empty()  // if stripping prefix left empty we have a match
+                               || rest.starts_with('/') // next component so we match
+                               // what we don't include is other matches,
+                               // we want to catch prefix/foo but not prefix-foo
+                            {
+                                yield meta_key;
+                            }
+
+                    }
                 }
             }
         };
@@ -2023,6 +2035,49 @@ mod tests {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn test_list_dir_with_prefix() -> Result<(), Box<dyn std::error::Error>> {
+        let storage: Arc<dyn Storage + Send + Sync> =
+            Arc::new(ObjectStorage::new_in_memory_store(Some("prefix".into())));
+        let ds = Repository::init(Arc::clone(&storage), false).await?.build();
+        let mut store = Store::from_repository(
+            ds,
+            AccessMode::ReadWrite,
+            Some("main".to_string()),
+            None,
+        );
+
+        store
+            .borrow_mut()
+            .set(
+                "zarr.json",
+                Bytes::copy_from_slice(br#"{"zarr_format":3, "node_type":"group"}"#),
+            )
+            .await?;
+
+        store
+            .borrow_mut()
+            .set(
+                "group/zarr.json",
+                Bytes::copy_from_slice(br#"{"zarr_format":3, "node_type":"group"}"#),
+            )
+            .await?;
+
+        store
+            .borrow_mut()
+            .set(
+                "group-suffix/zarr.json",
+                Bytes::copy_from_slice(br#"{"zarr_format":3, "node_type":"group"}"#),
+            )
+            .await?;
+
+        assert_eq!(
+            store.list_dir("group/").await?.try_collect::<Vec<_>>().await?,
+            vec!["zarr.json"]
+        );
+        Ok(())
+    }
+
     #[tokio::test]
     async fn test_get_partial_values() -> Result<(), Box<dyn std::error::Error>> {
         let storage: Arc<dyn Storage + Send + Sync> =