Allow setting row limits in Gaia queries

CHANGES.rst

@@ -8,11 +8,19 @@ New Tools and Services
 Service fixes and enhancements
 ------------------------------
 
+gaia
+^^^^
+
+- Allow for setting row limits in query submissions through class
+  attribute. [#1641]
+
 
 Infrastructure, Utility and Other Changes and Additions
 -------------------------------------------------------
 
 
+
+
 0.4 (2020-01-24)
 ================
 

astroquery/gaia/__init__.py

@@ -30,6 +30,9 @@ class Conf(_config.ConfigNamespace):
                                             "Name of RA parameter in table")
     MAIN_GAIA_TABLE_DEC = _config.ConfigItem("dec",
                                              "Name of Dec parameter in table")
+    ROW_LIMIT = _config.ConfigItem(50,
+                                   "Number of rows to return from database "
+                                   "query (set to -1 for unlimited).")
 
 
 conf = Conf()

astroquery/gaia/core.py

@@ -32,6 +32,7 @@ class GaiaClass(TapPlus):
     MAIN_GAIA_TABLE = conf.MAIN_GAIA_TABLE
     MAIN_GAIA_TABLE_RA = conf.MAIN_GAIA_TABLE_RA
     MAIN_GAIA_TABLE_DEC = conf.MAIN_GAIA_TABLE_DEC
+    ROW_LIMIT = conf.ROW_LIMIT
 
     def __init__(self, tap_plus_conn_handler=None, datalink_handler=None):
         super(GaiaClass, self).__init__(url="https://gea.esac.esa.int/",
@@ -182,6 +183,7 @@ def __query_object(self, coordinate, radius=None, width=None, height=None,
 
             query = """
                     SELECT
+                      {row_limit}
                       DISTANCE(
                         POINT('ICRS', {ra_column}, {dec_column}),
                         POINT('ICRS', {ra}, {dec})
@@ -203,6 +205,7 @@ def __query_object(self, coordinate, radius=None, width=None, height=None,
                     ORDER BY
                       dist ASC
                     """.format(**{
+                        'row_limit': "TOP {0}".format(self.ROW_LIMIT) if self.ROW_LIMIT > 0 else "",
                         'ra_column': self.MAIN_GAIA_TABLE_RA,
                         'dec_column': self.MAIN_GAIA_TABLE_DEC,
                         'columns': columns, 'table_name': self.MAIN_GAIA_TABLE,
@@ -329,6 +332,7 @@ def __cone_search(self, coordinate, radius, table_name=MAIN_GAIA_TABLE,
 
         query = """
                 SELECT
+                  {row_limit}
                   {columns},
                   DISTANCE(
                     POINT('ICRS', {ra_column}, {dec_column}),
@@ -345,6 +349,7 @@ def __cone_search(self, coordinate, radius, table_name=MAIN_GAIA_TABLE,
                   dist ASC
                 """.format(**{
                     'ra_column': ra_column_name,
+                    'row_limit': "TOP {0}".format(self.ROW_LIMIT) if self.ROW_LIMIT > 0 else "",
                     'dec_column': dec_column_name,
                     'columns': columns, 'ra': ra,
                     'dec': dec, 'radius': radiusDeg,

astroquery/gaia/tests/test_gaia_remote.py

@@ -0,0 +1,50 @@
+import astropy.units as u
+import pytest
+from astropy.coordinates import SkyCoord
+from astropy.tests.helper import remote_data
+
+from .. import GaiaClass
+
+
+@remote_data
+def test_query_object_row_limit():
+    Gaia = GaiaClass()
+    coord = SkyCoord(ra=280, dec=-60, unit=(u.degree, u.degree), frame='icrs')
+    width = u.Quantity(0.1, u.deg)
+    height = u.Quantity(0.1, u.deg)
+    r = Gaia.query_object_async(coordinate=coord, width=width, height=height)
+
+    assert len(r) == Gaia.ROW_LIMIT
+
+    Gaia.ROW_LIMIT = 10
+    r = Gaia.query_object_async(coordinate=coord, width=width, height=height)
+
+    assert len(r) == 10 == Gaia.ROW_LIMIT
+
+    Gaia.ROW_LIMIT = -1
+    r = Gaia.query_object_async(coordinate=coord, width=width, height=height)
+
+    assert len(r) == 176
+
+
+@remote_data
+def test_cone_search_row_limit():
+    Gaia = GaiaClass()
+    coord = SkyCoord(ra=280, dec=-60, unit=(u.degree, u.degree), frame='icrs')
+    radius = u.Quantity(0.1, u.deg)
+    j = Gaia.cone_search_async(coord, radius)
+    r = j.get_results()
+
+    assert len(r) == Gaia.ROW_LIMIT
+
+    Gaia.ROW_LIMIT = 10
+    j = Gaia.cone_search_async(coord, radius)
+    r = j.get_results()
+
+    assert len(r) == 10 == Gaia.ROW_LIMIT
+
+    Gaia.ROW_LIMIT = -1
+    j = Gaia.cone_search_async(coord, radius)
+    r = j.get_results()
+
+    assert len(r) == 1188