feat(stage_in): prefer direct DAAC S3 access, fall back to https

stage_in.py

@@ -4,9 +4,22 @@
 import argparse
 import pathlib
 import boto3
-from datetime import datetime
+import requests
+import urllib.parse
+from datetime import datetime, timedelta, timezone
 from maap.maap import MAAP
 
+CMR_COLLECTIONS_URL = "https://cmr.earthdata.nasa.gov/search/collections.umm_json"
+# Set STAGE_IN_DIRECT_S3=false to always use https for DAAC assets.
+DIRECT_S3_ENABLED = os.environ.get("STAGE_IN_DIRECT_S3", "true").lower() not in ("0", "false", "no")
+_CREDS_REFRESH_MARGIN = timedelta(minutes=5)
+_DEFAULT_CREDS_LIFETIME = timedelta(minutes=30)
+
+# endpoint -> (credentials dict, refresh_after datetime)
+_daac_creds_cache = {}
+# (provider, short_name, version) -> DirectDistributionInformation dict or None
+_daac_info_cache = {}
+
 _maap = None
 
 
@@ -18,15 +31,107 @@ def _get_maap():
     return _maap
 
 
-def _download_asset(href, destination_path):
-    """Route asset download: MAAP S3 buckets use boto3, everything else uses maap.download_granule."""
+def _is_daac_s3(href):
+    """True for s3:// hrefs outside the MAAP buckets (i.e. DAAC-hosted data)."""
+    return href.startswith("s3://") and not href.startswith("s3://maap-")
+
+
+def _cmr_ref_from_item(item):
+    """
+    Derive (provider, short_name, version) from a CMR-STAC item, or None.
+
+    CMR STAC collection links look like .../stac/<PROVIDER>/collections/<SHORT_NAME>_<VERSION>.
+    """
+    for link in item.links:
+        if link.rel not in ("collection", "parent"):
+            continue
+        parts = urllib.parse.urlparse(link.get_href() or "").path.strip("/").split("/")
+        if "stac" in parts and "collections" in parts:
+            i = parts.index("collections")
+            if i >= 1 and i + 1 < len(parts) and "_" in parts[i + 1] and parts[i - 1] != "stac":
+                short_name, version = parts[i + 1].rsplit("_", 1)
+                return parts[i - 1], short_name, version
+    return None
+
+
+def _direct_distribution_info(cmr_ref):
+    """Look up the collection's DirectDistributionInformation (S3 endpoint, region, buckets) in CMR."""
+    if cmr_ref in _daac_info_cache:
+        return _daac_info_cache[cmr_ref]
+    provider, short_name, version = cmr_ref
+    info = None
+    try:
+        response = requests.get(
+            CMR_COLLECTIONS_URL,
+            params={"provider": provider, "short_name": short_name, "version": version, "page_size": 1},
+            timeout=30,
+        )
+        response.raise_for_status()
+        items = response.json().get("items", [])
+        if items:
+            info = items[0]["umm"].get("DirectDistributionInformation")
+    except (requests.exceptions.RequestException, ValueError) as e:
+        print(f"    CMR lookup for {short_name}.{version} failed: {e}")
+    _daac_info_cache[cmr_ref] = info
+    return info
+
+
+def _daac_credentials(endpoint):
+    """Temporary DAAC S3 credentials for an s3credentials endpoint, cached until shortly before expiry."""
+    cached = _daac_creds_cache.get(endpoint)
+    if cached and datetime.now(timezone.utc) < cached[1]:
+        return cached[0]
+    creds = _get_maap().aws.earthdata_s3_credentials(endpoint)
+    try:
+        expires = datetime.fromisoformat(str(creds["expiration"]).replace("Z", "+00:00"))
+        if expires.tzinfo is None:
+            expires = expires.replace(tzinfo=timezone.utc)
+        refresh_after = expires - _CREDS_REFRESH_MARGIN
+    except (KeyError, ValueError):
+        refresh_after = datetime.now(timezone.utc) + _DEFAULT_CREDS_LIFETIME
+    _daac_creds_cache[endpoint] = (creds, refresh_after)
+    return creds
+
+
+def _download_daac_s3(href, destination_path, cmr_ref):
+    """Download a DAAC s3:// object using the DAAC's temporary credentials. Raises if not possible."""
+    if not DIRECT_S3_ENABLED:
+        raise RuntimeError("direct DAAC S3 access disabled (STAGE_IN_DIRECT_S3)")
+    if cmr_ref is None:
+        raise RuntimeError("cannot determine the CMR collection for this item")
+    info = _direct_distribution_info(cmr_ref)
+    endpoint = (info or {}).get("S3CredentialsAPIEndpoint")
+    if not endpoint:
+        raise RuntimeError("collection has no S3CredentialsAPIEndpoint in CMR")
+    bucket, key = href[5:].split("/", 1)
+    prefixes = (info or {}).get("S3BucketAndObjectPrefixNames") or []
+    if prefixes and not any(bucket == p[5:].split("/", 1)[0] for p in prefixes if p.startswith("s3://")):
+        raise RuntimeError(f"bucket '{bucket}' is not one of the collection's direct-access buckets")
+    creds = _daac_credentials(endpoint)
+    client = boto3.client(
+        "s3",
+        aws_access_key_id=creds["accessKeyId"],
+        aws_secret_access_key=creds["secretAccessKey"],
+        aws_session_token=creds["sessionToken"],
+        region_name=info.get("Region", "us-west-2"),
+    )
+    dest = os.path.join(destination_path, os.path.basename(key))
+    client.download_file(bucket, key, dest)
+    return dest
+
+
+def _download_asset(href, destination_path, cmr_ref=None):
+    """Route asset download: MAAP S3 buckets use boto3, DAAC s3:// uses DAAC credentials,
+    everything else uses maap.download_granule."""
     if href.startswith("s3://maap-"):
         parts = href[5:].split("/", 1)
         bucket, key = parts[0], parts[1]
         filename = os.path.basename(key)
         dest = os.path.join(destination_path, filename)
         boto3.client("s3").download_file(bucket, key, dest)
         return dest
+    if _is_daac_s3(href):
+        return _download_daac_s3(href, destination_path, cmr_ref)
     return _get_maap().download_granule(
         online_access_url=href,
         destination_path=destination_path,
@@ -59,15 +164,39 @@ def _download_single_item(item_json_data, output_dir, parent_container):
     local_item.set_self_href(str(pathlib.Path(item_dir).resolve() / f"{item.id}.json"))
 
     print(f"Downloading assets for item: {item.id}")
-    for asset_key, asset in item.assets.items():
+    cmr_ref = _cmr_ref_from_item(item)
+    https_by_name = {}
+    for asset in item.assets.values():
+        if asset.href.startswith(("http://", "https://")):
+            https_by_name.setdefault(os.path.basename(asset.href), asset.href)
+    fetched_via_s3 = {}  # file name -> local relative href, for files already fetched over direct S3
+
+    # DAAC s3:// assets first, so an https asset for the same file can reuse the result
+    ordered = sorted(item.assets.items(), key=lambda kv: not _is_daac_s3(kv[1].href))
+    for asset_key, asset in ordered:
         href = asset.href
-        print(f"  Downloading asset '{asset_key}': {os.path.basename(href)}")
+        name = os.path.basename(href)
+        key_dir = asset_key.replace("/", "_")
+        print(f"  Downloading asset '{asset_key}': {name}")
+        if href.startswith(("http://", "https://")) and name in fetched_via_s3:
+            print("    Already downloaded via direct S3, reusing")
+            local_item.assets[asset_key].href = fetched_via_s3[name]
+            continue
         try:
-            asset_dir = os.path.join(item_dir, asset_key.replace("/", "_"))
+            asset_dir = os.path.join(item_dir, key_dir)
             os.makedirs(asset_dir, exist_ok=True)
-            dest = _download_asset(href, asset_dir)
+            try:
+                dest = _download_asset(href, asset_dir, cmr_ref)
+                if _is_daac_s3(href):
+                    fetched_via_s3[name] = os.path.join(key_dir, os.path.basename(dest))
+            except Exception as e:
+                fallback = https_by_name.get(name) if _is_daac_s3(href) else None
+                if not fallback:
+                    raise
+                print(f"    Direct S3 failed ({e}); falling back to https")
+                dest = _download_asset(fallback, asset_dir, cmr_ref)
             print("Download complete")
-            local_item.assets[asset_key].href = os.path.join(asset_key.replace("/", "_"), os.path.basename(dest))
+            local_item.assets[asset_key].href = os.path.join(key_dir, os.path.basename(dest))
         except Exception as e:
             print(f"    Error downloading {asset_key}: {e}")
             local_item.assets.pop(asset_key, None)

tests/test_stage_in.py

@@ -865,8 +865,8 @@ def test_daac_url_uses_maap_download_granule(self, mock_maap, mock_boto3):
 
     @patch('stage_in.boto3')
     @patch('stage_in._get_maap')
-    def test_non_maap_s3_url_falls_through_to_maap(self, mock_maap, mock_boto3):
-        href = "s3://some-other-bucket/data/file.tif"
+    def test_non_s3_url_uses_maap_download_granule(self, mock_maap, mock_boto3):
+        href = "https://example.com/data/file.tif"
         mock_maap.return_value.download_granule.return_value = "/local/file.tif"
 
         _download_asset(href, "/local")
@@ -912,3 +912,113 @@ def test_download_asset_uses_download_granule(self, mock_get_maap):
         stage_in._download_asset("https://example.com/a.tif", "/tmp/dest")
         mock_get_maap.return_value.download_granule.assert_called_once_with(
             online_access_url="https://example.com/a.tif", destination_path="/tmp/dest", overwrite=True)
+
+
+DAAC_INFO = {
+    "Region": "us-west-2",
+    "S3CredentialsAPIEndpoint": "https://data.example-daac.nasa.gov/s3credentials",
+    "S3BucketAndObjectPrefixNames": ["s3://daac-prod-protected/Coll/"],
+}
+FUTURE = "2999-01-01T00:00:00+00:00"
+CREDS = {"accessKeyId": "AK", "secretAccessKey": "SK", "sessionToken": "ST", "expiration": FUTURE}
+
+
+def _daac_item(tmp_path):
+    """CMR-STAC style item: the same file as an https asset and as a DAAC s3:// asset."""
+    d = {
+        "type": "Feature", "stac_version": "1.0.0", "id": "daac-item",
+        "geometry": None, "bbox": None, "properties": {"datetime": "2024-01-01T00:00:00Z"},
+        "links": [{"rel": "collection", "href": "https://cmr.earthdata.nasa.gov/stac/DAAC_PROV/collections/MYD03_6.1"}],
+        "assets": {
+            "data": {"href": "https://daac.example/files/granule.hdf"},
+            "s3_data": {"href": "s3://daac-prod-protected/Coll/granule.hdf"},
+        },
+    }
+    return d
+
+
+@pytest.fixture(autouse=False)
+def clean_daac_caches():
+    import stage_in
+    stage_in._daac_creds_cache.clear()
+    stage_in._daac_info_cache.clear()
+    yield
+    stage_in._daac_creds_cache.clear()
+    stage_in._daac_info_cache.clear()
+
+
+class TestDaacDirectS3:
+    def test_cmr_ref_from_item_links(self, tmp_path):
+        import stage_in
+        item = pystac.Item.from_dict(_daac_item(tmp_path))
+        assert stage_in._cmr_ref_from_item(item) == ("DAAC_PROV", "MYD03", "6.1")
+
+    def test_cmr_ref_none_without_cmr_collection_link(self):
+        import stage_in
+        d = _daac_item(None)
+        d["links"] = []
+        assert stage_in._cmr_ref_from_item(pystac.Item.from_dict(d)) is None
+
+    @patch('stage_in.boto3')
+    @patch('stage_in._get_maap')
+    @patch('stage_in._direct_distribution_info', return_value=DAAC_INFO)
+    def test_prefers_direct_s3_and_skips_duplicate_https(self, _info, mock_get_maap, mock_boto3, tmp_path, clean_daac_caches):
+        mock_get_maap.return_value.aws.earthdata_s3_credentials.return_value = CREDS
+        local = stage_in_download(tmp_path, _daac_item(tmp_path))
+        mock_boto3.client.return_value.download_file.assert_called_once()
+        assert mock_boto3.client.call_args.kwargs["aws_session_token"] == "ST"
+        assert mock_boto3.client.call_args.kwargs["region_name"] == "us-west-2"
+        mock_get_maap.return_value.download_granule.assert_not_called()
+        assert set(local.assets) == {"data", "s3_data"}
+
+    @patch('stage_in.boto3')
+    @patch('stage_in._get_maap')
+    @patch('stage_in._direct_distribution_info', return_value=DAAC_INFO)
+    def test_falls_back_to_https_when_direct_s3_fails(self, _info, mock_get_maap, mock_boto3, tmp_path, clean_daac_caches):
+        mock_get_maap.return_value.aws.earthdata_s3_credentials.return_value = CREDS
+        mock_boto3.client.return_value.download_file.side_effect = Exception("AccessDenied")
+        mock_get_maap.return_value.download_granule.return_value = str(tmp_path / "daac-item" / "s3_data" / "granule.hdf")
+        local = stage_in_download(tmp_path, _daac_item(tmp_path))
+        # s3 asset falls back to https; the https asset is then downloaded as before
+        assert mock_get_maap.return_value.download_granule.call_count == 2
+        assert set(local.assets) == {"data", "s3_data"}
+
+    @patch('stage_in._get_maap')
+    @patch('stage_in._direct_distribution_info', return_value=None)
+    def test_falls_back_when_collection_has_no_direct_access(self, _info, mock_get_maap, tmp_path, clean_daac_caches):
+        mock_get_maap.return_value.download_granule.return_value = str(tmp_path / "daac-item" / "x" / "granule.hdf")
+        local = stage_in_download(tmp_path, _daac_item(tmp_path))
+        mock_get_maap.return_value.aws.earthdata_s3_credentials.assert_not_called()
+        assert set(local.assets) == {"data", "s3_data"}
+
+    @patch('stage_in._get_maap')
+    @patch('stage_in._direct_distribution_info', return_value=DAAC_INFO)
+    def test_s3_only_asset_without_https_sibling_is_dropped_on_failure(self, _info, mock_get_maap, tmp_path, clean_daac_caches):
+        mock_get_maap.return_value.aws.earthdata_s3_credentials.side_effect = Exception("401")
+        d = _daac_item(tmp_path)
+        del d["assets"]["data"]
+        local = stage_in_download(tmp_path, d)
+        assert "s3_data" not in local.assets
+
+    @patch('stage_in.boto3')
+    @patch('stage_in._get_maap')
+    @patch('stage_in._direct_distribution_info', return_value=DAAC_INFO)
+    def test_credentials_cached_per_endpoint(self, _info, mock_get_maap, mock_boto3, tmp_path, clean_daac_caches):
+        mock_get_maap.return_value.aws.earthdata_s3_credentials.return_value = CREDS
+        d = _daac_item(tmp_path)
+        d["assets"]["s3_two"] = {"href": "s3://daac-prod-protected/Coll/other.hdf"}
+        stage_in_download(tmp_path, d)
+        assert mock_get_maap.return_value.aws.earthdata_s3_credentials.call_count == 1
+
+    @patch('stage_in._get_maap')
+    @patch('stage_in._direct_distribution_info', return_value=DAAC_INFO)
+    def test_bucket_not_in_collection_buckets_is_rejected(self, _info, mock_get_maap, tmp_path, clean_daac_caches):
+        import stage_in
+        with pytest.raises(RuntimeError, match="not one of the collection"):
+            stage_in._download_daac_s3("s3://evil-bucket/x.hdf", str(tmp_path), ("P", "N", "1"))
+
+
+def stage_in_download(tmp_path, item_dict):
+    import stage_in
+    cat = pystac.Catalog(id="c", description="d")
+    return stage_in._download_single_item(item_dict, str(tmp_path), cat)