get_cloud_uris_query() method

snbianco · bsipocz · commit 9bd2c0672898 · 2024-07-15T21:22:31.000-07:00
diff --git a/CHANGES.rst b/CHANGES.rst
@@ -169,6 +169,8 @@ mast
 
 - Fix bug in ``Catalogs.query_criteria()`` to use ``page`` and ``pagesize`` parameters correctly. [#3065]
 
+- Add ``mast.Observations.get_cloud_uris_query`` method so that given a set of query criteria and optional filters,
+  the user receives a list of cloud data URIs for matching data products. [#3064]
 
 
 0.4.7 (2024-03-08)
diff --git a/astroquery/mast/observations.py b/astroquery/mast/observations.py
@@ -841,6 +841,59 @@ def get_cloud_uri(self, data_product, *, include_bucket=True, full_url=False):
         # Query for product URIs
         return self._cloud_connection.get_cloud_uri(data_product, include_bucket, full_url)
 
+    def get_cloud_uris_query(self, *, pagesize=None, page=None, mrp_only=False, extension=None,
+                             filter_products={}, **criteria):
+        """
+        Given a set of criteria and optional filters, get a list of matching data products and return their
+        associated cloud data URIs.
+
+        Parameters
+        ----------
+        pagesize : int, optional
+            Default None. Can be used to override the default pagesize.
+            E.g. when using a slow internet connection.
+        page : int, optional
+            Default None. Can be used to override the default behavior of all results being returned to obtain
+            one specific page of results.
+        mrp_only : bool, optional
+            Default False. When set to True, only "Minimum Recommended Products" will be returned.
+        extension : string or array, optional
+            Default None. Option to filter by file extension.
+        filter_products : dict, optional
+            Filters to be applied.  Valid filters are all products fields listed
+            `here <https://masttest.stsci.edu/api/v0/_productsfields.html>`__.
+            The column name as a string is the key. The corresponding value is one
+            or more acceptable values for that parameter.
+            Filter behavior is AND between the filters and OR within a filter set.
+            For example: {"productType": "SCIENCE", "extension"=["fits","jpg"]}
+        **criteria
+            Criteria to apply. At least one non-positional criteria must be supplied.
+            Valid criteria are coordinates, objectname, radius (as in `query_region` and `query_object`),
+            and all observation fields returned by the ``get_metadata("observations")``.
+            The Column Name is the keyword, with the argument being one or more acceptable values for that parameter,
+            except for fields with a float datatype where the argument should be in the form [minVal, maxVal].
+            For non-float type criteria wildcards maybe used (both * and % are considered wildcards), however
+            only one wildcarded value can be processed per criterion.
+            RA and Dec must be given in decimal degrees, and datetimes in MJD.
+            For example: filters=["FUV","NUV"],proposal_pi="Ost*",t_max=[52264.4586,54452.8914]
+
+        Returns
+        -------
+        response : list
+
+        """
+        # Get table of observations based on query criteria
+        obs = self.query_criteria(pagesize=pagesize, page=page, **criteria)
+
+        # Return list of associated data products
+        prod = self.get_product_list(obs)
+
+        # Filter product list
+        filt = self.filter_products(prod, mrp_only=mrp_only, extension=extension, **filter_products)
+
+        # Return list of cloud URIs
+        return self.get_cloud_uris(filt)
+
     def _remove_duplicate_products(self, data_products):
         """
         Removes duplicate data products that have the same dataURI.
diff --git a/astroquery/mast/tests/test_mast_remote.py b/astroquery/mast/tests/test_mast_remote.py
@@ -544,6 +544,27 @@ def test_get_cloud_uris(self):
 
         assert len(uris) > 0, f'Products for OBSID {test_obs_id} were not found in the cloud.'
 
+    def test_get_cloud_uris_query(self):
+        pytest.importorskip("boto3")
+
+        # enable access to public AWS S3 bucket
+        Observations.enable_cloud_dataset()
+
+        # get uris with other functions
+        obs = Observations.query_criteria(target_name=234295610,
+                                          provenance_name="SPOC",
+                                          sequence_number=[1, 2])
+        prod = Observations.get_product_list(obs)
+        filt = Observations.filter_products(prod, calib_level=[2])
+        s3_uris = Observations.get_cloud_uris(filt)
+
+        # get uris with streamlined function
+        uris = Observations.get_cloud_uris_query(target_name=234295610,
+                                                 provenance_name="SPOC",
+                                                 sequence_number=[1, 2],
+                                                 filter_products={'calib_level': [2]})
+        assert s3_uris == uris
+
     ######################
     # CatalogClass tests #
     ######################
diff --git a/docs/mast/mast_obsquery.rst b/docs/mast/mast_obsquery.rst
@@ -427,14 +427,16 @@ MAST until it is disabled with `~astroquery.mast.ObservationsClass.disable_cloud
 To directly access a list of cloud URIs for a given dataset, use the
 `~astroquery.mast.ObservationsClass.get_cloud_uris`
 function (Python will prompt you to enable cloud access if you haven't already).
+To return a list of cloud URIs based on query criteria and product filters, use the
+`~astroquery.mast.ObservationsClass.get_cloud_uris_query` function.
 
 When cloud access is enabled, the standard download function
 `~astroquery.mast.ObservationsClass.download_products` preferentially pulls files from AWS when they
 are available. When set to `True`, the ``cloud_only`` parameter in
 `~astroquery.mast.ObservationsClass.download_products` skips all data products not available in the cloud.
 
 
-Getting a list of S3 URIs:
+To get a list of S3 URIs, use the following workflow:
 
 .. doctest-skip::
 
@@ -456,10 +458,32 @@ Getting a list of S3 URIs:
    ...                                         productSubGroupDescription='DRZ')
    >>> s3_uris = Observations.get_cloud_uris(filtered)
    >>> print(s3_uris)
-   ['s3://stpubdata/hst/public/jbev/jbeveo010/jbeveo010_drz.fits', 's3://stpubdata/hst/public/jbev/jbeveo010/jbeveo010_drz.fits', 's3://stpubdata/hst/public/jbev/jbevet010/jbevet010_drz.fits', 's3://stpubdata/hst/public/jbev/jbevet010/jbevet010_drz.fits']
+   ['s3://stpubdata/hst/public/jbev/jbeveo010/jbeveo010_drz.fits', 's3://stpubdata/hst/public/jbev/jbevet010/jbevet010_drz.fits']
    ...
    >>> Observations.disable_cloud_dataset()
 
+Alternatively, you can use the streamlined `~astroquery.mast.ObservationsClass.get_cloud_uris_query` function. This approach is recommended
+for code brevity. Query criteria are supplied as keyword arguments, and filters are supplied through the ``filter_products`` parameter.
+
+.. doctest-skip::
+
+   >>> import os
+   >>> from astroquery.mast import Observations
+   ...
+   >>> Observations.enable_cloud_dataset(provider='AWS')
+   INFO: Using the S3 STScI public dataset [astroquery.mast.core]
+   ...
+   >>> # Getting the cloud URIs
+   >>> s3_uris = Observations.get_cloud_uris_query(obs_collection='HST',
+                                                   filters='F606W',
+                                                   instrument_name='ACS/WFC',
+                                                   proposal_id=['12062'],
+                                                   dataRights='PUBLIC',
+                                                   filter_products={'productSubGroupDescription': 'DRZ'})
+   >>> print(s3_uris)
+   ['s3://stpubdata/hst/public/jbev/jbeveo010/jbeveo010_drz.fits', 's3://stpubdata/hst/public/jbev/jbevet010/jbevet010_drz.fits']
+   ...
+   >>> Observations.disable_cloud_dataset()
 
 Downloading data products from S3:
 
