Obstore workflow cache

Author: lukasbindreiter

tilebox-workflows/tests/test_cache.py

@@ -8,8 +8,9 @@
 from _pytest.fixtures import SubRequest
 from moto import mock_aws
 from mypy_boto3_s3 import S3Client
+from obstore.store import LocalStore, MemoryStore
 
-from tilebox.workflows.cache import AmazonS3Cache, InMemoryCache, JobCache, LocalFileSystemCache
+from tilebox.workflows.cache import AmazonS3Cache, InMemoryCache, JobCache, LocalFileSystemCache, ObstoreCache
 
 
 @pytest.fixture
@@ -30,7 +31,7 @@ def aws(_aws_credentials: None) -> Iterator[S3Client]:
         yield boto3.client("s3", region_name="us-east-1")
 
 
-caches = ["LocalFileSystem", "InMemory", "AmazonS3", "AmazonS3_no_prefix"]
+caches = ["LocalFileSystem", "InMemory", "AmazonS3", "AmazonS3_no_prefix", "ObstoreMemory", "ObstoreLocal"]
 
 
 @pytest.fixture
@@ -48,6 +49,10 @@ def cache(request: SubRequest, tmp_path: Path, aws: S3Client) -> JobCache:
             bucket = "bucket1"
             aws.create_bucket(Bucket=bucket)
             return AmazonS3Cache(bucket, prefix="")
+        case "ObstoreMemory":
+            return ObstoreCache(MemoryStore())
+        case "ObstoreLocal":
+            return ObstoreCache(LocalStore(tmp_path))
         case _:
             raise ValueError("Invalid cache type")
 

tilebox-workflows/tilebox/workflows/cache.py

@@ -104,7 +104,8 @@ def __getitem__(self, key: str) -> bytes:
             raise KeyError(f"{key} is not cached!") from None
 
     def __iter__(self) -> Iterator[str]:
-        for obj in self.store.list_with_delimiter(str(self.prefix))["objects"]:
+        prefix = "" if self.prefix == ObjectPath(".") else str(self.prefix)
+        for obj in self.store.list_with_delimiter(prefix)["objects"]:
             path: str = obj["path"]
             yield path.removeprefix(str(self.prefix) + "/")