From 108da2ff44b891f04ea06b5e58541121777ed134 Mon Sep 17 00:00:00 2001 From: Sujen Shah Date: Mon, 19 Oct 2026 21:41:27 -0700 Subject: [PATCH] feat(stage_in): prefer direct DAAC S3 access, fall back to https - DAAC s3:// assets are fetched with temporary credentials from maap.aws.earthdata_s3_credentials(); the endpoint, region and allowed buckets come from the collection's CMR DirectDistributionInformation (provider/short name/version parsed from the CMR-STAC collection link) - credentials and CMR lookups are cached per endpoint/collection - on any failure (no CMR info, creds error, AccessDenied, ...) the same file's https asset is used instead; an https asset already fetched via S3 is reused rather than downloaded twice - STAGE_IN_DIRECT_S3=false disables direct S3 - non-MAAP s3:// hrefs no longer go to download_granule (HTTP-only) Refs #25 --- stage_in.py | 145 ++++++++++++++++++++++++++++++++++++++--- tests/test_stage_in.py | 114 +++++++++++++++++++++++++++++++- 2 files changed, 249 insertions(+), 10 deletions(-) diff --git a/stage_in.py b/stage_in.py index b146c0f..1d06049 100644 --- a/stage_in.py +++ b/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,8 +31,98 @@ 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//collections/_. + """ + 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] @@ -27,6 +130,8 @@ def _download_asset(href, destination_path): 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) diff --git a/tests/test_stage_in.py b/tests/test_stage_in.py index b71ba28..6a90f27 100644 --- a/tests/test_stage_in.py +++ b/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)