restore get_query_payload to all methods

Author: weaverba137

astroquery/sdss/core.py

@@ -520,7 +520,7 @@ class = 'galaxy' \
 
     def get_spectra_async(self, coordinates=None, radius=2. * u.arcsec,
                           matches=None, plate=None, fiberID=None, mjd=None,
-                          timeout=TIMEOUT,
+                          timeout=TIMEOUT, get_query_payload=False,
                           data_release=conf.default_release, cache=True,
                           show_progress=True):
         """
@@ -559,6 +559,9 @@ def get_spectra_async(self, coordinates=None, radius=2. * u.arcsec,
         timeout : float, optional
             Time limit (in seconds) for establishing successful connection with
             remote server.  Defaults to `SDSSClass.TIMEOUT`.
+        get_query_payload : bool, optional
+            If True, this will return the data the query would have sent out,
+            but does not actually do the query.
         data_release : int, optional
             The data release of the SDSS to use. With the default server, this
             only supports DR8 or later.
@@ -599,12 +602,19 @@ def get_spectra_async(self, coordinates=None, radius=2. * u.arcsec,
             if coordinates is None:
                 matches = self.query_specobj(plate=plate, mjd=mjd, fiberID=fiberID,
                                              fields=['run2d', 'plate', 'mjd', 'fiberID'],
-                                             timeout=timeout, data_release=data_release, cache=cache)
+                                             timeout=timeout, get_query_payload=get_query_payload,
+                                             data_release=data_release, cache=cache)
             else:
-                matches = self.query_crossid(coordinates, radius=radius,
+                matches = self.query_crossid(coordinates, radius=radius, timeout=timeout,
                                              specobj_fields=['run2d', 'plate', 'mjd', 'fiberID'],
-                                             spectro=True,
-                                             timeout=timeout, data_release=data_release, cache=cache)
+                                             spectro=True, get_query_payload=get_query_payload,
+                                             data_release=data_release, cache=cache)
+            if get_query_payload:
+                if coordinates is None:
+                    return matches
+                else:
+                    return matches[0]
+
             if matches is None:
                 warnings.warn("Query returned no results.", NoResultsWarning)
                 return
@@ -640,8 +650,8 @@ def get_spectra_async(self, coordinates=None, radius=2. * u.arcsec,
     @prepend_docstr_nosections(get_spectra_async.__doc__)
     def get_spectra(self, coordinates=None, radius=2. * u.arcsec,
                     matches=None, plate=None, fiberID=None, mjd=None,
-                    timeout=TIMEOUT, cache=True,
-                    data_release=conf.default_release,
+                    timeout=TIMEOUT, get_query_payload=False,
+                    data_release=conf.default_release, cache=True,
                     show_progress=True):
         """
         Returns
@@ -654,9 +664,14 @@ def get_spectra(self, coordinates=None, radius=2. * u.arcsec,
                                                radius=radius, matches=matches,
                                                plate=plate, fiberID=fiberID,
                                                mjd=mjd, timeout=timeout,
+                                               get_query_payload=get_query_payload,
                                                data_release=data_release,
+                                               cache=cache,
                                                show_progress=show_progress)
 
+        if get_query_payload:
+            return readable_objs
+
         if readable_objs is not None:
             if isinstance(readable_objs, dict):
                 return readable_objs
@@ -666,7 +681,7 @@ def get_spectra(self, coordinates=None, radius=2. * u.arcsec,
     def get_images_async(self, coordinates=None, radius=2. * u.arcsec,
                          matches=None, run=None, rerun=301, camcol=None,
                          field=None, band='g', timeout=TIMEOUT,
-                         cache=True,
+                         cache=True, get_query_payload=False,
                          data_release=conf.default_release,
                          show_progress=True):
         """
@@ -714,6 +729,9 @@ def get_images_async(self, coordinates=None, radius=2. * u.arcsec,
         timeout : float, optional
             Time limit (in seconds) for establishing successful connection with
             remote server.  Defaults to `SDSSClass.TIMEOUT`.
+        get_query_payload : bool, optional
+            If True, this will return the data the query would have sent out,
+            but does not actually do the query.
         cache : bool, optional
             Cache the images using astropy's caching system
         data_release : int, optional
@@ -753,12 +771,19 @@ def get_images_async(self, coordinates=None, radius=2. * u.arcsec,
                 matches = self.query_photoobj(run=run, rerun=rerun,
                                               camcol=camcol, field=field,
                                               fields=['run', 'rerun', 'camcol', 'field'],
-                                              timeout=timeout,
+                                              timeout=timeout, get_query_payload=get_query_payload,
                                               data_release=data_release, cache=cache)
             else:
-                matches = self.query_crossid(coordinates, radius=radius,
+                matches = self.query_crossid(coordinates, radius=radius, timeout=timeout,
                                              fields=['run', 'rerun', 'camcol', 'field'],
-                                             timeout=timeout, data_release=data_release, cache=cache)
+                                             get_query_payload=get_query_payload,
+                                             data_release=data_release, cache=cache)
+            if get_query_payload:
+                if coordinates is None:
+                    return matches
+                else:
+                    return matches[0]
+
             if matches is None:
                 warnings.warn("Query returned no results.", NoResultsWarning)
                 return
@@ -798,10 +823,22 @@ def get_images(self, coordinates=None, radius=2. * u.arcsec,
 
         """
 
-        readable_objs = self.get_images_async(
-            coordinates=coordinates, radius=radius, matches=matches, run=run,
-            rerun=rerun, data_release=data_release, camcol=camcol, field=field,
-            band=band, timeout=timeout, show_progress=show_progress)
+        readable_objs = self.get_images_async(coordinates=coordinates,
+                                              radius=radius,
+                                              matches=matches,
+                                              run=run,
+                                              rerun=rerun,
+                                              camcol=camcol,
+                                              field=field,
+                                              band=band,
+                                              timeout=timeout,
+                                              cache=cache,
+                                              get_query_payload=get_query_payload,
+                                              data_release=data_release,
+                                              show_progress=show_progress)
+
+        if get_query_payload:
+            return readable_objs
 
         if readable_objs is not None:
             if isinstance(readable_objs, dict):

astroquery/sdss/tests/test_sdss.py

@@ -454,6 +454,62 @@ def test_photoobj_run_camcol_field_payload(patch_request, dr):
     assert query_payload['format'] == 'csv'
 
 
+@pytest.mark.parametrize("dr", dr_list)
+def test_get_spectra_specobj_payload(patch_request, dr):
+    expect = ("SELECT DISTINCT "
+              "s.run2d, s.plate, s.mjd, s.fiberID "
+              "FROM PhotoObjAll AS p "
+              "JOIN SpecObjAll AS s ON p.objID = s.bestObjID "
+              "WHERE "
+              "(s.plate=751 AND s.mjd=52251)")
+    query_payload = sdss.SDSS.get_spectra_async(plate=751, mjd=52251,
+                                                get_query_payload=True,
+                                                data_release=dr)
+    assert query_payload['cmd'] == expect
+    assert query_payload['format'] == 'csv'
+
+
+@pytest.mark.parametrize("dr", dr_list)
+def test_get_spectra_coordinates_payload(patch_request, dr):
+    expect = ("SELECT\r\n"
+              "s.run2d, s.plate, s.mjd, s.fiberID, s.SpecObjID AS obj_id, dbo.fPhotoTypeN(p.type) AS type "
+              "FROM #upload u JOIN #x x ON x.up_id = u.up_id JOIN PhotoObjAll AS p ON p.objID = x.objID "
+              "JOIN SpecObjAll AS s ON p.objID = s.bestObjID "
+              "ORDER BY x.up_id")
+    query_payload = sdss.SDSS.get_spectra_async(coordinates=coords_column,
+                                                get_query_payload=True,
+                                                data_release=dr)
+    assert query_payload['uquery'] == expect
+    assert query_payload['format'] == 'csv'
+    assert query_payload['photoScope'] == 'nearPrim'
+
+@pytest.mark.parametrize("dr", dr_list)
+def test_get_images_photoobj_payload(patch_request, dr):
+    expect = ("SELECT DISTINCT "
+              "p.run, p.rerun, p.camcol, p.field "
+              "FROM PhotoObjAll AS p WHERE "
+              "(p.run=5714 AND p.camcol=6 AND p.rerun=301)")
+    query_payload = sdss.SDSS.get_images_async(run=5714, camcol=6,
+                                               get_query_payload=True,
+                                               data_release=dr)
+    assert query_payload['cmd'] == expect
+    assert query_payload['format'] == 'csv'
+
+
+@pytest.mark.parametrize("dr", dr_list)
+def test_get_images_coordinates_payload(patch_request, dr):
+    expect = ("SELECT\r\n"
+              "p.run, p.rerun, p.camcol, p.field, dbo.fPhotoTypeN(p.type) AS type "
+              "FROM #upload u JOIN #x x ON x.up_id = u.up_id JOIN PhotoObjAll AS p ON p.objID = x.objID "
+              "ORDER BY x.up_id")
+    query_payload = sdss.SDSS.get_images_async(coordinates=coords_column,
+                                               get_query_payload=True,
+                                               data_release=dr)
+    assert query_payload['uquery'] == expect
+    assert query_payload['format'] == 'csv'
+    assert query_payload['photoScope'] == 'nearPrim'
+
+
 @pytest.mark.parametrize("dr", dr_list)
 def test_spectra_plate_mjd_payload(patch_request, dr):
     expect = ("SELECT DISTINCT "