From b4b5de867bc6fa047f05be445d14ec7fe65d7acf Mon Sep 17 00:00:00 2001
From: Ciaran Sweet <9111975+ciaransweet@users.noreply.github.com>
Date: Fri, 26 Jul 2024 11:17:12 +0100
Subject: [PATCH] Add ability to provide AWS Session Token and Error on no
 files found (#7)

---
 src/hazard/docs_store.py           |  7 ++++++-
 src/hazard/sources/osc_zarr.py     | 16 ++++++++++++----
 src/hazard/sources/ukcp18_rcp85.py |  6 ++++++
 3 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/src/hazard/docs_store.py b/src/hazard/docs_store.py
index 2249322..6207291 100644
--- 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
index 3e0cb3f..752995b 100644
--- 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
index 015524a..a2cc1dd 100644
--- 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)