Obstore workflow cache

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

@@ -10,6 +10,7 @@
 from botocore.exceptions import ClientError
 from google.cloud.exceptions import NotFound
 from google.cloud.storage import Blob, Bucket
+from obstore.exceptions import GenericError
 from obstore.store import ObjectStore
 
 
@@ -100,11 +101,14 @@ def __getitem__(self, key: str) -> bytes:
         try:
             entry = self.store.get(str(self.prefix / key))
             return bytes(entry.bytes())
-        except OSError:
+        except (OSError, GenericError):
+            # GenericError is raised if the key contains separator characters, but one of the parents is a file
+            # instead of a directory
             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) + "/")