No longer show overflow warning if user set maxrec

Author: stvoutsin

CHANGES.rst

@@ -4,6 +4,8 @@
 Enhancements and Fixes
 ----------------------
 
+- No longer show a warning when an overflow status is encountered with a maxrec set by the user [#690]
+
 - Add support for the jobInfo element for UWS jobs [#679]
 
 - Fix job result handling to prioritize standard URL structure [#684]

docs/dal/index.rst

@@ -348,13 +348,23 @@ rows they will return before overflowing:
     20000
 
 To retrieve more rows than that (often conservative) default limit, you
-must override maxrec in the call to ``search``. A warning can be expected if
-you reach the ``maxrec`` limit:
+must override maxrec in the call to ``search``. PyVO will only warn about 
+truncation when it's unexpected. If you request 5 records and get 5 records,
+no warning is issued:
 
 .. doctest-remote-data::
 
-    >>> tap_results = tap_service.search("SELECT * FROM arihip.main", maxrec=5)  # doctest: +SHOW_WARNINGS
-    DALOverflowWarning: Result set limited by user- or server-supplied MAXREC parameter.
+    >>> tap_results = tap_service.search("SELECT * FROM arihip.main", maxrec=5)
+    >>> len(tap_results)
+    5
+
+However, if results are truncated by server limits without you specifying 
+maxrec, you'll receive a helpful warning:
+
+.. doctest-remote-data::
+
+    >>> tap_results = tap_service.search("SELECT * FROM arihip.main")  # doctest: +SHOW_WARNINGS
+    DALOverflowWarning: Results truncated due to server limits. Consider setting a maxrec value.
 
 Services will not let you raise maxrec beyond the hard match limit:
 

pyvo/dal/query.py

@@ -303,7 +303,7 @@ def from_result_url(cls, result_url, *, session=None):
             url=result_url,
             session=session)
 
-    def __init__(self, votable, *, url=None, session=None):
+    def __init__(self, votable, *, url=None, session=None, user_maxrec=None):
         """
         initialize the cursor.  This constructor is not typically called
         by directly applications; rather an instance is obtained from calling
@@ -318,6 +318,8 @@ def __init__(self, votable, *, url=None, session=None):
            the URL that produced the response
         session : object
            optional session to use for network requests
+        user_maxrec: int
+              the maximum number of records that the user requested.
 
         Raises
         ------
@@ -332,14 +334,14 @@ def __init__(self, votable, *, url=None, session=None):
 
         self._url = url
         self._session = use_session(session)
+        self._user_maxrec = user_maxrec
 
         self._status = self._findstatus(votable)
         if self._status[0].lower() not in ("ok", "overflow"):
             raise DALQueryError(self._status[1], self._status[0], url)
 
         if self._status[0].lower() == "overflow":
-            warn("Result set limited by user- or server-supplied MAXREC parameter.",
-                 category=DALOverflowWarning)
+            self._handle_overflow_warning(user_maxrec)
 
         self._resultstable = self._findresultstable(votable)
         if not self._resultstable:
@@ -355,6 +357,54 @@ def __init__(self, votable, *, url=None, session=None):
 
         self._infos = self._findinfos(votable)
 
+    def _handle_overflow_warning(self, user_maxrec=None):
+        """
+        Handle overflow warning - can be overridden by subclasses.
+
+        Default implementation issues a generic overflow warning.
+        Subclasses can override this to customize or suppress warnings.
+
+        Parameters
+        ----------
+        user_maxrec : int, optional
+            The maximum records requested by the user
+        """
+        warn("Result set limited by user- or server-supplied MAXREC "
+             "parameter.", category=DALOverflowWarning)
+
+    def check_overflow_warning(self, user_maxrec=None):
+        """
+        Check for overflow warnings and issue them if appropriate.
+        It will check if the results were truncated due to server limits
+        and issue a warning if the number of records returned is less than
+        the maximum records requested by the user.
+
+        If the results were truncated it distinguishes between expected
+        truncation (where the user requested a maxrec)
+        and unexpected truncation (where the user did not specify a maxrec).
+        If the results were truncated it issues a warning.
+
+        Parameters
+        ----------
+        user_maxrec : int, optional
+            The maximum records explicitly requested by the user
+        """
+        if self._status[0].lower() == "overflow":
+            maxrec_to_check = user_maxrec if user_maxrec is not None else self._user_maxrec
+
+            if (maxrec_to_check is not None
+                    and len(self.resultstable.array) == maxrec_to_check):
+                pass
+            else:
+                if maxrec_to_check is not None:
+                    warn(f"Results truncated at {len(self.resultstable.array)} records by service limits "
+                         f"(you requested maxrec={maxrec_to_check})",
+                         category=DALOverflowWarning)
+                else:
+                    warn("Results truncated due to server limits. Consider "
+                         "setting a maxrec value.",
+                         category=DALOverflowWarning)
+
     def _findresultstable(self, votable):
         # this can be overridden to specialize for a particular DAL protocol
         res = self._findresultsresource(votable)

pyvo/dal/tap.py

@@ -654,6 +654,7 @@ def create(
             uploads=uploads, session=session, **keywords)
         response = tapquery.submit()
         job = cls(response.url, session=session)
+        job._user_maxrec = maxrec
         return job
 
     def __init__(self, url, *, session=None, delete=True):
@@ -672,6 +673,7 @@ def __init__(self, url, *, session=None, delete=True):
         self._url = url
         self._session = use_session(session)
         self._delete_on_exit = delete
+        self._user_maxrec = None
         self._update()
 
     def __enter__(self):
@@ -1047,7 +1049,9 @@ def fetch_result(self):
 
         response.raw.read = partial(
             response.raw.read, decode_content=True)
-        return TAPResults(votableparse(response.raw.read), url=self.result_uri, session=self._session)
+        result = TAPResults(votableparse(response.raw.read), url=self.result_uri, session=self._session)
+        result.check_overflow_warning(self._user_maxrec)
+        return result
 
 
 class TAPQuery(DALQuery):
@@ -1103,6 +1107,8 @@ def __init__(
         self["REQUEST"] = "doQuery"
         self["LANG"] = language
 
+        self._user_maxrec = maxrec
+
         if maxrec:
             self["MAXREC"] = maxrec
 
@@ -1155,7 +1161,14 @@ def execute(self):
         DALFormatError
            for errors parsing the VOTable response
         """
-        return TAPResults(self.execute_votable(), url=self.queryurl, session=self._session)
+        result = TAPResults(
+            self.execute_votable(),
+            url=self.queryurl,
+            session=self._session
+        )
+        result.check_overflow_warning(self._user_maxrec)
+
+        return result
 
     def submit(self, *, post=False):
         """
@@ -1265,6 +1278,15 @@ def getrecord(self, index):
         """
         return TAPRecord(self, index, session=self._session)
 
+    def _handle_overflow_warning(self, user_maxrec=None):
+        """
+        TAP-specific overflow warning handling.
+
+        For TAP results we suppress the default overflow warning during
+        initialization because TAPQuery.execute() will call check_overflow_warning()
+        """
+        pass
+
 
 class TAPRecord(SodaRecordMixin, DatalinkRecordMixin, Record):
     pass

pyvo/dal/tests/test_query.py

@@ -3,6 +3,7 @@
 """
 Tests for pyvo.dal.query
 """
+import warnings
 from functools import partial
 
 from contextlib import ExitStack
@@ -419,6 +420,70 @@ def test_columnaliases(self):
         assert dalresults.fieldname_with_ucd('baz') is None
         assert dalresults.fieldname_with_utype('foobaz') is None
 
+    def test_check_overflow_warning_no_maxrec(self):
+        with pytest.warns(DALOverflowWarning):
+            dalresults = DALResults.from_result_url('http://example.com/query/overflowstatus')
+
+        with pytest.warns(DALOverflowWarning, match="Results truncated due to server limits"):
+            dalresults.check_overflow_warning()
+
+    def test_check_overflow_warning_exact_match(self):
+        with pytest.warns(DALOverflowWarning):
+            dalresults = DALResults.from_result_url('http://example.com/query/overflowstatus')
+
+        with warnings.catch_warnings():
+            warnings.simplefilter("error")
+            dalresults.check_overflow_warning(user_maxrec=3)
+
+    def test_check_overflow_warning_service_truncation(self):
+        with pytest.warns(DALOverflowWarning):
+            dalresults = DALResults.from_result_url('http://example.com/query/overflowstatus')
+
+        with pytest.warns(DALOverflowWarning, match="Results truncated at 3 records by service limits"):
+            dalresults.check_overflow_warning(user_maxrec=1000)
+
+    def test_check_overflow_warning_uses_stored_maxrec(self):
+        with pytest.warns(DALOverflowWarning):
+            dalresults = DALResults.from_result_url('http://example.com/query/overflowstatus')
+
+        dalresults._user_maxrec = 1000
+
+        with pytest.warns(DALOverflowWarning, match="Results truncated at 3 records by service limits"):
+            dalresults.check_overflow_warning()
+
+    def test_check_overflow_warning_parameter_overrides_stored(self):
+        with pytest.warns(DALOverflowWarning):
+            dalresults = DALResults.from_result_url('http://example.com/query/overflowstatus')
+
+        dalresults._user_maxrec = 1000
+
+        with warnings.catch_warnings():
+            warnings.simplefilter("error")
+            dalresults.check_overflow_warning(user_maxrec=3)
+
+    def test_check_overflow_warning_no_overflow_status(self):
+        dalresults = DALResults.from_result_url('http://example.com/query/basic')
+
+        with warnings.catch_warnings():
+            warnings.simplefilter("error")
+            dalresults.check_overflow_warning(user_maxrec=1000)
+
+    def test_handle_overflow_warning_default_behavior(self):
+        votable = votableparse(BytesIO(get_pkg_data_contents('data/query/overflowstatus.xml')))
+
+        with pytest.warns(DALOverflowWarning,
+                          match="Result set limited by user- or server-supplied MAXREC parameter."):
+            _ = DALResults(votable, url='http://test.com')
+
+    def test_stored_user_maxrec_initialization(self):
+        votable = votableparse(BytesIO(get_pkg_data_contents('data/query/basic.xml')))
+
+        result = DALResults(votable, url='http://test.com', user_maxrec=100)
+        assert result._user_maxrec == 100
+
+        result2 = DALResults(votable, url='http://test.com')
+        assert result2._user_maxrec is None
+
 
 @pytest.mark.filterwarnings('ignore::astropy.io.votable.exceptions.W03')
 @pytest.mark.filterwarnings('ignore::astropy.io.votable.exceptions.W06')