Handle Azure SAS tokens and credential refresh

gefs_wave.py

@@ -1,5 +1,6 @@
 import json
 import logging
+import os
 import tarfile
 import threading
 from dataclasses import dataclass
@@ -13,6 +14,7 @@ import requests
 NOMADS_URL = "https://nomads.ncep.noaa.gov/pub/data/nccf/com/gens/prod"
 AZURE_URL = "https://noaagefs.blob.core.windows.net/gefs"
 BASE_URLS = [AZURE_URL, NOMADS_URL]
+AZURE_SAS_ENDPOINT = "https://planetarycomputer.microsoft.com/api/sas/v1/token/noaagefs/gefs"
 FORECAST_HOURS: List[int] = list(range(0, 387, 3))
 STATE_FILENAME = "state.json"
 TARBALL_FILENAME = "gefswave-wave0-latest.tar.gz"
@@ -51,6 +53,8 @@ class WaveDownloader:
         self.data_root = data_root
         self.data_root.mkdir(parents=True, exist_ok=True)
         self.session = session or requests.Session()
+        self._azure_sas_token: Optional[str] = None
+        self._azure_token_checked = False
 
     def ensure_latest_cycle(self) -> Dict:
         """Ensure the latest available control-member dataset is present locally."""
@@ -98,7 +102,11 @@ class WaveDownloader:
     def _cycle_available(self, day: str, cycle: str) -> bool:
         filename = f"gefswave.t{cycle}z.c00.global.0p25.f000.grib2"
         for base_url in BASE_URLS:
-            url = f"{base_url}/gefs.{day}/{cycle}/wave/gridded/{filename}"
+            path = f"gefs.{day}/{cycle}/wave/gridded/{filename}"
+            url = self._build_url(base_url, path)
+            if not url:
+                logger.debug("Skipping %s for cycle %s %sz due to missing credentials.", base_url, day, cycle)
+                continue
             try:
                 response = self.session.head(url, timeout=20)
                 if response.ok:
@@ -106,6 +114,16 @@ class WaveDownloader:
                     return True
                 if response.status_code == 404:
                     logger.debug("Cycle %s %sz missing on %s (404).", day, cycle, base_url)
+                elif response.status_code in (401, 403):
+                    logger.warning(
+                        "Cycle %s %sz request to %s returned %s; will retry after refreshing credentials.",
+                        day,
+                        cycle,
+                        base_url,
+                        response.status_code,
+                    )
+                    if base_url == AZURE_URL:
+                        self._invalidate_azure_token()
                 else:
                     logger.debug(
                         "Cycle %s %sz head request on %s returned HTTP %s.",
@@ -145,19 +163,26 @@ class WaveDownloader:
     def _download_with_fallback(self, cycle: Cycle, filename: str, destination: Path) -> None:
         errors = []
         for base_url in BASE_URLS:
-            url = f"{base_url}/{cycle.directory}/wave/gridded/{filename}"
+            path = f"{cycle.directory}/wave/gridded/{filename}"
+            url = self._build_url(base_url, path)
+            if not url:
+                logger.debug("Skipping download from %s for %s due to missing URL.", base_url, filename)
+                continue
             try:
-                self._stream_to_file(url, destination)
+                self._stream_to_file(url, destination, base_url)
                 return
             except WaveDownloadError as exc:
                 errors.append(str(exc))
                 logger.warning("Failed to download %s from %s (%s).", filename, base_url, exc)
         raise WaveDownloadError(f"All download attempts failed for {filename}: {'; '.join(errors)}")
 
-    def _stream_to_file(self, url: str, destination: Path) -> None:
+    def _stream_to_file(self, url: str, destination: Path, base_url: str) -> None:
         tmp_path = destination.with_suffix(destination.suffix + ".part")
         try:
             with self.session.get(url, stream=True, timeout=120) as response:
+                if response.status_code in (401, 403) and base_url == AZURE_URL:
+                    self._invalidate_azure_token()
+                    response.raise_for_status()
                 response.raise_for_status()
                 with tmp_path.open("wb") as handle:
                     for chunk in response.iter_content(chunk_size=1 << 20):
@@ -171,6 +196,48 @@ class WaveDownloader:
             logger.error("Failed to download %s: %s", url, exc)
             raise WaveDownloadError(f"Failed to download {url}: {exc}") from exc
 
+    def _build_url(self, base_url: str, path: str) -> Optional[str]:
+        if base_url == AZURE_URL:
+            token = self._get_azure_sas_token()
+            if not token:
+                return None
+            token = token if token.startswith("?") else f"?{token}"
+            return f"{base_url}/{path}{token}"
+        return f"{base_url}/{path}"
+
+    def _get_azure_sas_token(self) -> Optional[str]:
+        if self._azure_sas_token:
+            return self._azure_sas_token
+
+        env_token = os.environ.get("GEFS_AZURE_SAS_TOKEN")
+        if env_token:
+            self._azure_sas_token = env_token.strip()
+            logger.info("Using Azure SAS token from environment variable.")
+            return self._azure_sas_token
+
+        if self._azure_token_checked:
+            return None
+
+        self._azure_token_checked = True
+        try:
+            response = self.session.get(AZURE_SAS_ENDPOINT, timeout=20)
+            response.raise_for_status()
+            token = response.json().get("token")
+            if token:
+                self._azure_sas_token = token.strip()
+                logger.info("Fetched Azure SAS token from Planetary Computer API.")
+            else:
+                logger.warning("Azure SAS endpoint response did not include a token.")
+        except requests.RequestException as exc:
+            logger.warning("Failed to fetch Azure SAS token: %s", exc)
+        return self._azure_sas_token
+
+    def _invalidate_azure_token(self) -> None:
+        if self._azure_sas_token:
+            logger.info("Invalidating cached Azure SAS token.")
+        self._azure_sas_token = None
+        self._azure_token_checked = False
+
     def _cycle_complete(self, cycle: Cycle) -> bool:
         target_dir = self.data_root / cycle.date / cycle.cycle
         if not target_dir.exists():