speed up get_cloud_uris()

Author: snbianco

CHANGES.rst

@@ -171,6 +171,9 @@ mast
 
 - Modify ``mast.Observations.get_cloud_uris`` to also accept query criteria and data product filters. [#3064]
 
+- Increased the speed of ``mast.Observations.get_cloud_uris`` by obtaining multiple
+  URIs from MAST at once. [#3064]
+
 
 0.4.7 (2024-03-08)
 ==================

astroquery/mast/cloud.py

@@ -14,7 +14,7 @@
 from astropy.utils.console import ProgressBarOrSpinner
 from astropy.utils.exceptions import AstropyDeprecationWarning
 
-from ..exceptions import NoResultsWarning, InvalidQueryError
+from ..exceptions import NoResultsWarning
 
 from . import utils
 
@@ -109,32 +109,14 @@ def get_cloud_uri(self, data_product, include_bucket=True, full_url=False):
             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/")
-        elif '/ps1/' in path:
-            path = path.replace("/ps1/", "panstarrs/ps1/public/")
+        # Making sure we got at least 1 URI from the query above.
+        if uri_list[0] is 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):
         """
@@ -158,8 +140,33 @@ def get_cloud_uri_list(self, data_products, include_bucket=True, full_url=False)
             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, verbose=True):
         """

astroquery/mast/observations.py

@@ -784,7 +784,7 @@ def get_cloud_uris(self, data_products=None, *, include_bucket=True, full_url=Fa
         ----------
         data_products : `~astropy.table.Table`
             Table containing products to be converted into cloud data uris. If provided, this will supercede
-            page_size, page, or any arguments passed in as **criteria.
+            page_size, page, or any keyword arguments passed in as criteria.
         include_bucket : bool
             Default True. When False, returns the path of the file relative to the
             top level cloud storage location.

astroquery/mast/utils.py

@@ -158,22 +158,54 @@ def parse_input_location(coordinates=None, objectname=None):
 
 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/")
+            elif '/ps1/' in path:
+                path = path.replace("/ps1/", "panstarrs/ps1/public/")
+            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]