Merge pull request #3517 from snbianco/subscription-prods

Download products from event subscription JSON files

Author: bsipocz

CHANGES.rst

@@ -74,6 +74,11 @@ mast
 - The ``select_cols`` parameter in ``MastMissions`` query functions now accepts an iterable of column names, a comma-delimited 
   string of column names, or the special values 'all' or '*' to return all available columns. [#3492]
 
+- Improved robustness of product downloads for ``MastMissions``, including support for subscription-service JSON inputs and 
+  clearer validation of MAST URIs and product metadata. [#3517]
+
+- Added full support for the International Ultraviolet Explorer (IUE) mission in ``MastMissions``. [#3517]
+
 jplspec
 ^^^^^^^
 

astroquery/mast/missions.py

@@ -7,6 +7,7 @@
 """
 
 import difflib
+import json
 import warnings
 from collections.abc import Iterable
 from json import JSONDecodeError
@@ -59,7 +60,8 @@ def __init__(self, *, mission='hst', mast_token=None):
             'jwst': 'fileSetName',
             'roman': 'fileSetName',
             'classy': 'Target',
-            'ullyses': 'observation_id'
+            'ullyses': 'observation_id',
+            'iue': 'iue_data_id'
         }
 
         # Service attributes
@@ -616,7 +618,7 @@ def download_file(self, uri, *, local_path=None, cache=True, verbose=True):
         Parameters
         ----------
         uri : str
-            The product dataURI
+            The product filename or URI to be downloaded.
         local_path : str
             Directory or filename to which the file will be downloaded.  Defaults to current working directory.
         cache : bool
@@ -635,14 +637,18 @@ def download_file(self, uri, *, local_path=None, cache=True, verbose=True):
         """
 
         # Construct the full data URL based on mission
-        if self.mission in ['hst', 'jwst', 'roman']:
+        if self.mission in ['hst', 'jwst', 'roman', 'roman_spectra', 'roman_cgi']:
             # HST, JWST, and RST have a dedicated endpoint for retrieving products
             base_url = self._service_api_connection.MISSIONS_DOWNLOAD_URL + self.mission + '/api/v0.1/retrieve_product'
             keyword = 'product_name'
         else:
             # HLSPs use MAST download URL
             base_url = self._service_api_connection.MAST_DOWNLOAD_URL
             keyword = 'uri'
+            # These files require a MAST URI and not just a filename
+            if not uri.startswith('mast:'):
+                raise InvalidQueryError(f'For mission "{self.mission}", a full MAST URI is required for downloading. '
+                                        f'Got "{uri}".')
         data_url = base_url + f'?{keyword}=' + uri
         escaped_url = base_url + f'?{keyword}=' + quote(uri, safe='')
 
@@ -714,13 +720,25 @@ def _download_files(self, products, base_dir, *, flat=False, cache=True, verbose
         base_dir = Path(base_dir)
 
         for data_product in products:
+            col_names = data_product.colnames
             # Determine local path for each file
-            local_path = base_dir / data_product['dataset'] if not flat else base_dir
+            filename = data_product['filename']
+            uri = data_product['uri'] if 'uri' in col_names else filename
+            dataset = None
+            if 'dataset' in col_names:
+                dataset = data_product['dataset']
+            elif 'fileset' in col_names:
+                dataset = data_product['fileset']
+            if not dataset and not flat:
+                raise InvalidQueryError('Data product is missing "dataset" or "fileset" field required for '
+                                        'constructing local download path. Specify `flat=True` to avoid this '
+                                        'requirement.')
+            local_path = base_dir / dataset if not flat else base_dir
             local_path.mkdir(parents=True, exist_ok=True)
-            local_file_path = local_path / Path(data_product['filename']).name
+            local_file_path = local_path / Path(filename).name
 
             # Download files and record status
-            status, msg, url = self.download_file(data_product['uri'],
+            status, msg, url = self.download_file(uri,
                                                   local_path=local_file_path,
                                                   cache=cache,
                                                   verbose=verbose)
@@ -737,9 +755,10 @@ def download_products(self, products, *, download_dir=None, flat=False,
 
         Parameters
         ----------
-        products : str, list, `~astropy.table.Table`
+        products : str, list of str, `~astropy.table.Table`, or list of dict
             Either a single or list of dataset IDs (e.g., as input for `get_product_list`),
-            or a Table of products (e.g., as output from `get_product_list`)
+            a Table of products (e.g., as output from `get_product_list`), or a JSON file or data from
+            the MAST subscription service containing product information.
         download_dir : str or Path, optional
             Directory for file downloads.  Defaults to current directory.
         flat : bool, optional
@@ -764,11 +783,30 @@ def download_products(self, products, *, download_dir=None, flat=False,
         manifest : `~astropy.table.Table`
             A table manifest showing downloaded file locations and statuses.
         """
+        if not products:
+            raise InvalidQueryError('No products specified for download.')
+
         # Ensure `products` is a Table, collecting products if necessary
-        if isinstance(products, (str, list)):
+        if (isinstance(products, str) and products.endswith('.json')) or isinstance(products, Path):
+            # Products coming from local JSON filepath from subscription service
+            try:
+                with open(products, 'r') as f:
+                    json_data = json.load(f)
+            except JSONDecodeError as ex:
+                raise InvalidQueryError(f'Failed to decode JSON file at {products}: {ex}')
+
+            if not isinstance(json_data, (list, tuple)):
+                raise InvalidQueryError(f'Expected a list of product rows in JSON file at {products}.')
+            products = Table(rows=json_data)
+        elif isinstance(products, (list)) and all(isinstance(prod, dict) for prod in products):
+            # Products coming from JSON data from subscription service
+            products = Table(rows=products)
+        elif isinstance(products, (str, list)):
+            # Products given as dataset ID(s)
             products = [products] if isinstance(products, str) else products
             products = vstack([self.get_product_list(oid) for oid in products])
         elif isinstance(products, Row):
+            # Single row of products
             products = Table(products, masked=True)
 
         # Apply filters
@@ -778,7 +816,7 @@ def download_products(self, products, *, download_dir=None, flat=False,
         products = utils.remove_duplicate_products(products, 'filename')
 
         if not len(products):
-            warnings.warn("No products to download.", NoResultsWarning)
+            warnings.warn("No products to download after applying filters.", NoResultsWarning)
             return
 
         # Set up base directory for downloads