breaking: daac downloader accepts only features collection

Author: wphyojpl

cumulus_lambda_functions/stage_in_out/download_granules_abstract.py

@@ -3,6 +3,8 @@
 import os
 from abc import ABC, abstractmethod
 
+from pystac import ItemCollection
+
 from cumulus_lambda_functions.lib.utils.file_utils import FileUtils
 
 LOGGER = logging.getLogger(__name__)
@@ -11,11 +13,13 @@
 class DownloadGranulesAbstract(ABC):
     STAC_JSON = 'STAC_JSON'
     DOWNLOAD_DIR_KEY = 'DOWNLOAD_DIR'
+    DOWNLOADING_KEYS = 'DOWNLOADING_KEYS'
 
     def __init__(self) -> None:
         super().__init__()
-        self._granules_json = []
+        self._granules_json: ItemCollection = {}
         self._download_dir = '/tmp'
+        self._downloading_keys = set([k.strip() for k in os.environ.get(self.DOWNLOADING_KEYS, 'data').strip().split(',')])
 
     def _setup_download_dir(self):
         self._download_dir = os.environ.get(self.DOWNLOAD_DIR_KEY)
@@ -27,17 +31,18 @@ def _setup_download_dir(self):
     def _retrieve_stac_json(self):
         raw_stac_json = os.environ.get(self.STAC_JSON)
         try:
-            self._granules_json = json.loads(raw_stac_json)
+            self._granules_json = ItemCollection.from_dict(json.loads(raw_stac_json))
             return self
         except:
             LOGGER.debug(f'raw_stac_json is not STAC_JSON: {raw_stac_json}. trying to see if file exists')
         if not FileUtils.file_exist(raw_stac_json):
             raise ValueError(f'missing file or not JSON: {raw_stac_json}')
-        self._granules_json = FileUtils.read_json(raw_stac_json)
-        if self._granules_json is None:
+        json_stac = FileUtils.read_json(raw_stac_json)
+        if json_stac is None:
             raise ValueError(f'{raw_stac_json} is not JSON')
+        self._granules_json = ItemCollection.from_dict(json_stac)
         return self
 
     @abstractmethod
-    def download(self, **kwargs) -> list:
+    def download(self, **kwargs) -> dict:
         raise NotImplementedError()

cumulus_lambda_functions/stage_in_out/download_granules_daac.py

@@ -1,6 +1,8 @@
 import requests
+from pystac import ItemCollection, Item, Asset
 
 from cumulus_lambda_functions.lib.earthdata_login.urs_token_retriever import URSTokenRetriever
+from cumulus_lambda_functions.lib.utils.file_utils import FileUtils
 from cumulus_lambda_functions.stage_in_out.download_granules_abstract import DownloadGranulesAbstract
 import json
 import logging
@@ -25,76 +27,50 @@ def __set_props_from_env(self):
         self.__edl_token = URSTokenRetriever().start()
         return self
 
-    def __get_downloading_urls(self, granules_result: list):
-        if len(granules_result) < 1:
-            LOGGER.warning(f'cannot find any granules')
-            return []
-        downloading_urls = [k['assets'] for k in granules_result]
-        return downloading_urls
-
-    def __download_one_granule(self, assets: dict):
-        """
-        sample assets
-          {
-            "data": {
-              "href": "s3://am-uds-dev-cumulus-internal/ATMS_SCIENCE_Group___1/P1570515ATMSSCIENCEAAT16017044853900.PDS",
-              "title": "P1570515ATMSSCIENCEAAT16017044853900.PDS",
-              "description": "P1570515ATMSSCIENCEAAT16017044853900.PDS"
-            },
-            "metadata__data": {
-              "href": "s3://am-uds-dev-cumulus-internal/ATMS_SCIENCE_Group___1/P1570515ATMSSCIENCEAAT16017044853901.PDS",
-              "title": "P1570515ATMSSCIENCEAAT16017044853901.PDS",
-              "description": "P1570515ATMSSCIENCEAAT16017044853901.PDS"
-            },
-            "metadata__xml": {
-              "href": "s3://am-uds-dev-cumulus-internal/ATMS_SCIENCE_Group___1/P1570515ATMSSCIENCEAAT16017044853901.PDS.xml",
-              "title": "P1570515ATMSSCIENCEAAT16017044853901.PDS.xml",
-              "description": "P1570515ATMSSCIENCEAAT16017044853901.PDS.xml"
-            },
-            "metadata__cmr": {
-              "href": "s3://am-uds-dev-cumulus-internal/ATMS_SCIENCE_Group___1/P1570515ATMSSCIENCEAAT16017044853900.PDS.cmr.xml",
-              "title": "P1570515ATMSSCIENCEAAT16017044853900.PDS.cmr.xml",
-              "description": "P1570515ATMSSCIENCEAAT16017044853900.PDS.cmr.xml"
-            }
-          }
-        :param assets:
-        :return:
-        """
-        error_log = []
+    def __download_one_granule_item(self, granule_item: Item):
         headers = {
             'Authorization': f'Bearer {self.__edl_token}'
         }
-        local_item = {}
-        for k, v in assets.items():
-            local_item[k] = v
-            try:
-                LOGGER.debug(f'downloading: {v["href"]}')
-                r = requests.get(v['href'], headers=headers)
-                if r.status_code >= 400:
-                    raise RuntimeError(f'wrong response status: {r.status_code}. details: {r.content}')
-                # TODO. how to correctly check redirecting to login page
-                local_file_path = os.path.join(self._download_dir, os.path.basename(v["href"]))
-                with open(local_file_path, 'wb') as fd:
-                    fd.write(r.content)
-                local_item[k]['href'] = local_file_path
-            except Exception as e:
-                LOGGER.exception(f'failed to download {v}')
-                local_item[k]['description'] = f'download failed. {str(e)}'
-                error_log.append(v)
-        return local_item, error_log
+        new_asset_dict = {}
+        for name, value_dict in granule_item.assets.items():
+            if name not in self._downloading_keys:
+                LOGGER.debug(f'skipping {name}. Not in downloading keys')
+                continue
+            value_dict: Asset = value_dict
+            downloading_url = value_dict.href
+            LOGGER.debug(f'downloading: {downloading_url}')
+            r = requests.get(downloading_url, headers=headers)
+            r.raise_for_status()
+            local_file_path = os.path.join(self._download_dir, os.path.basename(downloading_url))
+            with open(local_file_path, 'wb') as fd:
+                fd.write(r.content)
+            value_dict.href = os.path.join('.', os.path.basename(downloading_url))
+            new_asset_dict[name] = value_dict
+        granule_item.assets = new_asset_dict
+        return granule_item
 
-    def download(self, **kwargs) -> list:
+    def download(self, **kwargs) -> dict:
         self.__set_props_from_env()
         LOGGER.debug(f'creating download dir: {self._download_dir}')
-        downloading_urls = self.__get_downloading_urls(self._granules_json)
-        error_list = []
+        if len(self._granules_json.items) < 1:
+            LOGGER.warning(f'cannot find any granules')
+            return self._granules_json.to_dict(False)
         local_items = []
-        for each in downloading_urls:
-            LOGGER.debug(f'working on {each}')
-            local_item, current_error_list = self.__download_one_granule(each)
-            error_list.extend(current_error_list)
-            local_items.append({'assets': local_item})
+        error_list = []
+        for each_item in self._granules_json.items:
+            try:
+                local_item = self.__download_one_granule_item(each_item)
+                local_items.append(local_item)
+            except Exception as e:
+                LOGGER.exception(f'error downloading granule: {each_item.id}')
+                error_list.append({'error': str(e), 'id': each_item.id, })
+        LOGGER.debug(f'finished downloading all granules')
+        self._granules_json.items = local_items
+        LOGGER.debug(f'writing features collection json to downloading directory')
+        granules_json_dict = self._granules_json.to_dict(False)
+        FileUtils.write_json(os.path.join(self._download_dir, 'downloaded_feature_collection.json'), granules_json_dict, overwrite=True, prettify=True)
+        LOGGER.debug(f'writing errors if any')
         if len(error_list) > 0:
             with open(f'{self._download_dir}/error.log', 'w') as error_file:
                 error_file.write(json.dumps(error_list, indent=4))
-        return local_items
+        return granules_json_dict