Add ability to provide AWS Session Token and Error on no files found (#7

)
developmentseed · Jul 26, 2024 · b4b5de8 · b4b5de8
1 parent 341911f
commit b4b5de8
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 5 deletions.
diff --git a/src/hazard/docs_store.py b/src/hazard/docs_store.py
@@ -16,6 +16,7 @@ class DocStore:
     # environment variable names:
     __access_key = "OSC_S3_ACCESS_KEY_DEV"
     __secret_key = "OSC_S3_SECRET_KEY_DEV"
+    __token = "OSC_S3_TOKEN_DEV"
     __S3_bucket = "OSC_S3_BUCKET_DEV"  # e.g. redhat-osc-physical-landing-647521352890
 
     def __init__(
@@ -42,7 +43,11 @@ def __init__(
         if fs is None:
             access_key = os.environ.get(self.__access_key, None)
             secret_key = os.environ.get(self.__secret_key, None)
-            fs = s3fs.S3FileSystem(key=access_key, secret=secret_key)
+            token = os.environ.get(self.__token, None)
+            if token:
+                fs = s3fs.S3FileSystem(key=access_key, secret=secret_key, token=token)
+            else:
+                fs = s3fs.S3FileSystem(key=access_key, secret=secret_key)
 
         self._fs = fs
         if type(self._fs) == s3fs.S3FileSystem:  # noqa: E721 # use isinstance?

diff --git a/src/hazard/sources/osc_zarr.py b/src/hazard/sources/osc_zarr.py
@@ -16,6 +16,10 @@
 
 
 class OscZarr(ReadWriteDataArray):
+    __access_key = "OSC_S3_ACCESS_KEY_DEV"
+    __secret_key = "OSC_S3_SECRET_KEY_DEV"
+    __token = "OSC_S3_TOKEN_DEV"
+
     def __init__(
         self,
         bucket: str = default_dev_bucket,
@@ -37,10 +41,14 @@ def __init__(
         if store is None:
             if s3 is None:
                 # zarr_utilities.load_dotenv() # to load environment variables
-                s3 = s3fs.S3FileSystem(
-                    key=os.environ.get("OSC_S3_ACCESS_KEY_DEV", None),
-                    secret=os.environ.get("OSC_S3_SECRET_KEY_DEV", None),
-                )
+                access_key = os.environ.get(self.__access_key, None)
+                secret_key = os.environ.get(self.__secret_key, None)
+                token = os.environ.get(self.__token, None)
+                if token:
+                    s3 = s3fs.S3FileSystem(key=access_key, secret=secret_key, token=token)
+                else:
+                    s3 = s3fs.S3FileSystem(key=access_key, secret=secret_key)
+
             group_path = str(PurePosixPath(bucket, prefix, "hazard.zarr"))
             store = s3fs.S3Map(root=group_path, s3=s3, check=False)
 

diff --git a/src/hazard/sources/ukcp18_rcp85.py b/src/hazard/sources/ukcp18_rcp85.py
@@ -50,6 +50,12 @@ def open_dataset_year(
         files_available_for_quantity: List[str] = self._get_files_available_for_quantity_and_year(
             gcm, scenario, quantity, year
         )
+
+        if not files_available_for_quantity:
+            raise Exception(
+                f"No UKCP18 files available for: gcm:{gcm}, scenario:{scenario}, quantity:{quantity}, year:{year}"
+            )
+
         all_data_from_files = self._combine_all_files_data(files_available_for_quantity)
         only_data_for_year = all_data_from_files.sel(time=str(year))
         reprojected = self._reproject_quantity(only_data_for_year, quantity)