fix(tensorstore): Specify path explicitly

openclimatefix · Nov 11, 2024 · 9f55084 · 9f55084
1 parent 7bc387f
commit 9f55084
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 9 deletions.
diff --git a/src/nwp_consumer/internal/entities/tensorstore.py b/src/nwp_consumer/internal/entities/tensorstore.py
@@ -129,12 +129,12 @@ def initialize_empty_store(
 
         zarrdir = os.getenv("ZARRDIR", f"~/.local/cache/nwp/{repository}/{model}/data")
         store: zarr.storage.Store
+        path: str
         try:
-            store = zarr.storage.DirectoryStore(
-                pathlib.Path(
-                    "/".join((zarrdir, TensorStore.gen_store_filename(coords=coords))),
-                ).expanduser().as_posix(),
-            )
+            path = pathlib.Path(
+                "/".join((zarrdir, TensorStore.gen_store_filename(coords=coords))),
+            ).expanduser().as_posix()
+            store = zarr.storage.DirectoryStore(path)
             if zarrdir.startswith("s3"):
                 import s3fs
                 log.debug("Attempting AWS connection using credential discovery")
@@ -145,7 +145,8 @@ def initialize_empty_store(
                             "region_name": os.getenv("AWS_REGION", "eu-west-1"),
                         },
                     )
-                    store = s3fs.mapping.S3Map(zarrdir, fs, check=True, create=False)
+                    store = s3fs.mapping.S3Map(zarrdir, fs, check=True, create=True)
+                    path = zarrdir
                 except Exception as e:
                     return Failure(OSError(
                         f"Unable to create file mapping for ZARRDIR '{zarrdir}'. "
@@ -200,7 +201,7 @@ def initialize_empty_store(
         return Success(
             cls(
                 name=model,
-                path=store.path,
+                path=path,
                 coordinate_map=coordinate_map_result.unwrap(),
                 size_kb=0,
                 encoding=encoding,

diff --git a/src/nwp_consumer/internal/handlers/cli.py b/src/nwp_consumer/internal/handlers/cli.py
@@ -89,7 +89,7 @@ def run(self) -> int:
                         log.error(f"Failed to consume NWP data: {e}")
                         return 1
                     case Success(path):
-                        log.info(f"Successfully consumed NWP data to '{path.as_posix()}'")
+                        log.info(f"Successfully consumed NWP data to '{path}'")
                         return 0
 
             case "archive":
@@ -100,7 +100,7 @@ def run(self) -> int:
                         log.error(f"Failed to archive NWP data: {e}")
                         return 1
                     case Success(path):
-                        log.info(f"Successfully archived NWP data to '{path.as_posix()}'")
+                        log.info(f"Successfully archived NWP data to '{path}'")
                         return 0
 
             case "info":