Speed up mast.Observations.get_cloud_uris()

CHANGES.rst

@@ -242,6 +242,12 @@ xmatch
 - Minor internal change to use VOTable as the response format that include
  units, too. [#1375]
 
+mast
+^^^^
+
+- Increased the speed of ``Observations.get_cloud_uris`` by obtaining multiple
+ URIs from MAST at once. [#2145]
+
 
 Infrastructure, Utility and Other Changes and Additions
 -------------------------------------------------------

astroquery/mast/cloud.py

@@ -109,30 +109,14 @@
  found in the cloud, None is returned.
  """
 
- s3_client = self.boto3.client('s3', config=self.config)
-
- path = utils.mast_relative_path(data_product["dataURI"])
- if path is None:
- raise InvalidQueryError("Malformed data uri {}".format(data_product['dataURI']))
+ uri_list = self.get_cloud_uri_list(data_product, include_bucket=include_bucket, full_url=full_url)
 
- if 'galex' in path:
- path = path.lstrip("/mast/")
+ # Making sure we got at least 1 URI from the query above.
+ if uri_list[0] == None:
+ warnings.warn("Unable to locate file {}.".format(data_product), NoResultsWarning)
  else:
- path = path.lstrip("/")
-
- try:
- s3_client.head_object(Bucket=self.pubdata_bucket, Key=path)
- if include_bucket:
- path = "s3://{}/{}".format(self.pubdata_bucket, path)
- elif full_url:
- path = "http://s3.amazonaws.com/{}/{}".format(self.pubdata_bucket, path)
- return path
- except self.botocore.exceptions.ClientError as e:
- if e.response['Error']['Code'] != "404":
- raise
-
- warnings.warn("Unable to locate file {}.".format(data_product['productFilename']), NoResultsWarning)
- return None
+ # Output from ``get_cloud_uri_list`` is always a list even when it's only 1 URI
+ return uri_list[0]
 
  def get_cloud_uri_list(self, data_products, include_bucket=True, full_url=False):
  """
@@ -156,8 +140,33 @@
  List of URIs generated from the data products, list way contain entries that are None
  if data_products includes products not found in the cloud.
  """
+ s3_client = self.boto3.client('s3', config=self.config)
 
- return [self.get_cloud_uri(product, include_bucket, full_url) for product in data_products]
+ paths = utils.mast_relative_path(data_products["dataURI"])
+ if isinstance(paths, str): # Handle the case where only one product was requested
+ paths = [paths]
+
+ uri_list = []
+ for path in paths:
+ if path is None:
+ uri_list.append(None)
+ else:
+ try:
+ # Use `head_object` to verify that the product is available on S3 (not all products are)
+ s3_client.head_object(Bucket=self.pubdata_bucket, Key=path)
+ if include_bucket:
+ s3_path = "s3://{}/{}".format(self.pubdata_bucket, path)
+ uri_list.append(s3_path)
+ elif full_url:
+ path = "http://s3.amazonaws.com/{}/{}".format(self._pubdata_bucket, path)
+ uri_list.append(path)
+ except self.botocore.exceptions.ClientError as e:
+ if e.response['Error']['Code'] != "404":
+ raise
+ warnings.warn("Unable to locate file {}.".format(path), NoResultsWarning)
+ uri_list.append(None)
+
+ return uri_list
 
  def download_file(self, data_product, local_path, cache=True):
  """

astroquery/mast/tests/test_mast_remote.py

@@ -403,14 +403,15 @@ def test_get_cloud_uri(self):
 
  assert len(uri) > 0, f'Product for OBSID {test_obs_id} was not found in the cloud.'
 
- def test_get_cloud_uris(self):
+ @pytest.mark.parametrize("test_obs_id", ["25568122", "31411"])
+ def test_get_cloud_uris(self, test_obs_id):
  pytest.importorskip("boto3")
- test_obs_id = '25568122'
 
  # get a product list
- products = mast.Observations.get_product_list(test_obs_id)[24:]
+ index = 24 if test_obs_id == "25568122" else 0
+ products = mast.Observations.get_product_list(test_obs_id)[index:]
 
- assert len(products) > 0, (f'No products found for OBSID {test_obs_id}. '
+ assert len(products) > 0, (f'No products found for OBSID {test_obs_id}.'
  'Unable to move forward with getting URIs from the cloud.')
 
  # enable access to public AWS S3 bucket
@@ -427,6 +428,7 @@ def test_get_cloud_uris(self):
 
  # query functions
  def test_catalogs_query_region_async(self):
+
  responses = mast.Catalogs.query_region_async("158.47924 -7.30962", catalog="Galex")
  assert isinstance(responses, list)
 

astroquery/mast/utils.py

@@ -155,22 +155,52 @@
 
 def mast_relative_path(mast_uri):
  """
- Given a MAST dataURI, return the associated relative path.
+ Given one or more MAST dataURI(s), return the associated relative path(s).
 
  Parameters
  ----------
- mast_uri : str
- The MAST uri.
+ mast_uri : str, list of str
+ The MAST uri(s).
 
  Returns
  -------
- response : str
- The associated relative path.
+ response : str, list of str
+ The associated relative path(s).
  """
-
- response = _simple_request("https://mast.stsci.edu/api/v0.1/path_lookup/",
- {"uri": mast_uri})
- result = response.json()
- uri_result = result.get(mast_uri)
-
- return uri_result["path"]
+ if isinstance(mast_uri, str):
+ uri_list = [("uri", mast_uri)]
+ else: # mast_uri parameter is a list
+ uri_list = [("uri", uri) for uri in mast_uri]
+
+ # Split the list into chunks of 50 URIs; this is necessary
+ # to avoid "414 Client Error: Request-URI Too Large".
+ uri_list_chunks = list(_split_list_into_chunks(uri_list, chunk_size=50))
+
+ result = []
+ for chunk in uri_list_chunks:
+ response = _simple_request("https://mast.stsci.edu/api/v0.1/path_lookup/",
+ {"uri": chunk})
+ json_response = response.json()
+
+ for uri in chunk:
+ # Chunk is a list of tuples where the tuple is
+ # ("uri", "/path/to/product")
+ # so we index for path (index=1)
+ path = json_response.get(uri[1])["path"]
+ if 'galex' in path:
+ path = path.lstrip("/mast/")
+ else:
+ path = path.lstrip("/")
+ result.append(path)
+
+ # If the input was a single URI string, we return a single string
+ if isinstance(mast_uri, str):
+ return result[0]
+ # Else, return a list of paths
+ return result
+
+
+def _split_list_into_chunks(input_list, chunk_size):
+ """Helper function for `mast_relative_path`."""
+ for idx in range(0, len(input_list), chunk_size):
+ yield input_list[idx:idx + chunk_size]