From 971fee1dbfa606c4b4104f1197925c13bfeae46a Mon Sep 17 00:00:00 2001
From: Carlos Tejo <carlos.tejo@landportal.info>
Date: Fri, 12 Jul 2019 18:42:15 +0200
Subject: [PATCH] initial support for 429 TooManyRequests HTTP code

---
 SPARQLWrapper/SPARQLExceptions.py | 13 +++++++++++++
 SPARQLWrapper/Wrapper.py          | 10 +++++++++-
 test/wrapper_test.py              | 19 +++++++++++++++----
 3 files changed, 37 insertions(+), 5 deletions(-)

diff --git a/SPARQLWrapper/SPARQLExceptions.py b/SPARQLWrapper/SPARQLExceptions.py
index 1de7b2d..824ead5 100644
--- a/SPARQLWrapper/SPARQLExceptions.py
+++ b/SPARQLWrapper/SPARQLExceptions.py
@@ -65,3 +65,16 @@ class URITooLong(SPARQLWrapperException):
 
     msg = "the URI requested by the client is longer than the server is willing to interpret. Check if the request was sent using GET method instead of POST method."
 
+class TooManyRequests(SPARQLWrapperException):
+    """
+    The user has sent too many requests in a given amount of time ("rate limiting"). Usually HTTP response status code 429.
+    @since: 1.8.5
+    """
+
+    def __init__(self, response=None, retry_after=None):
+        self.msg = "The user has sent too many requests in a given amount of time ('rate limiting'). Retry-After header (if present) indicates how long to wait before making a new request (possible values are <http-date> and <delay-seconds>)."
+        if retry_after:
+            self.msg += "Retry-After header value="+retry_after
+
+        super(TooManyRequests, self).__init__(response)
+        self.retry_after = retry_after
diff --git a/SPARQLWrapper/Wrapper.py b/SPARQLWrapper/Wrapper.py
index 988abc9..ca68026 100644
--- a/SPARQLWrapper/Wrapper.py
+++ b/SPARQLWrapper/Wrapper.py
@@ -51,7 +51,7 @@
 
 import json
 from KeyCaseInsensitiveDict import KeyCaseInsensitiveDict
-from SPARQLExceptions import QueryBadFormed, EndPointNotFound, EndPointInternalError, Unauthorized, URITooLong
+from SPARQLExceptions import QueryBadFormed, EndPointNotFound, EndPointInternalError, Unauthorized, URITooLong, TooManyRequests
 from SPARQLWrapper import __agent__
 
 #  From <https://www.w3.org/TR/sparql11-protocol/#query-success>
@@ -940,7 +940,9 @@ def _query(self):
         @raise Unauthorized: If the C{HTTP return code} is C{401}.
         @raise EndPointNotFound: If the C{HTTP return code} is C{404}.
         @raise URITooLong: If the C{HTTP return code} is C{414}.
+        @raise TooManyRequests: If the C{HTTP return code} is C{429}.
         @raise EndPointInternalError: If the C{HTTP return code} is C{500}.
+        @raise urllib2.HTTPError: If the C{HTTP return code} is different to C{400}, C{401}, C{404}, C{414}, C{429}, C{500}.
         """
         request = self._createRequest()
 
@@ -959,6 +961,12 @@ def _query(self):
                 raise Unauthorized(e.read())
             elif e.code == 414:
                 raise URITooLong(e.read())
+            elif e.code == 429:
+                hdrs = KeyCaseInsensitiveDict(e.headers) # Case insensitive headers
+                if "retry-after" in hdrs:
+                    raise TooManyRequests(e.read(), hdrs["retry-after"])
+                else:
+                    raise TooManyRequests(e.read())
             elif e.code == 500:
                 raise EndPointInternalError(e.read())
             else:
diff --git a/test/wrapper_test.py b/test/wrapper_test.py
index ed9bbca..93e1456 100644
--- a/test/wrapper_test.py
+++ b/test/wrapper_test.py
@@ -37,7 +37,7 @@
 from SPARQLWrapper import XML, GET, POST, JSON, JSONLD, N3, TURTLE, RDF, SELECT, INSERT, RDFXML, CSV, TSV
 from SPARQLWrapper import URLENCODED, POSTDIRECTLY
 from SPARQLWrapper import BASIC, DIGEST
-from SPARQLWrapper.Wrapper import QueryResult, QueryBadFormed, EndPointNotFound, EndPointInternalError, Unauthorized, URITooLong
+from SPARQLWrapper.Wrapper import QueryResult, QueryBadFormed, EndPointNotFound, EndPointInternalError, Unauthorized, URITooLong, TooManyRequests
 
 
 class FakeResult(object):
@@ -49,13 +49,12 @@ def urlopener(request):
     return FakeResult(request)
 
 
-def urlopener_error_generator(code):
+def urlopener_error_generator(code, headers={}):
     def urlopener_error(request):
-        raise HTTPError(request.get_full_url, code, '', {}, StringIO(u''))
+        raise HTTPError(request.get_full_url, code, '', headers, StringIO(u''))
 
     return urlopener_error
 
-
 def urlopener_check_data_encoding(request):
     if sys.version < '3':  # have to write it like this, for 2to3 compatibility
         if isinstance(request.data, unicode):
@@ -546,6 +545,18 @@ def testQuery(self):
         except:
             self.fail('got wrong exception')
 
+        RETRY_AFTER_VALUE = "100"
+        _victim.urlopener = urlopener_error_generator(429, {"Retry-After": RETRY_AFTER_VALUE})
+        try:
+            self.wrapper.query()
+            self.fail('should have raised exception')
+        except TooManyRequests as e:
+            #  TODO: check exception-format
+            self.assertEquals(e.retry_after, RETRY_AFTER_VALUE)
+            pass
+        except:
+            self.fail('got wrong exception')
+
         _victim.urlopener = urlopener_error_generator(500)
         try:
             self.wrapper.query()