Repository: auth0/auth0-python
Pull request #219 (2020-06-23 22:37:07)

Expose the time at which the Rate Limit will reset
### Changes

When too many requests have been made against the server, the next request will fail with a `429`. In the response headers there will be extra information, such as the time at which this limit is reset.

With this PR: 
* `429` errors are raised as `RateLimitError` instead of `Auth0Error`. 
* They will expose a property `reset_at` with the value of the `X-RateLimit-Reset` header. This value represents the Unix timestamp at which the limit will be reset.
* When this header is not set, a value of -1 will be set as default.

### References

More info https://auth0.com/docs/policies/rate-limits

### Testing

- [x] This change adds unit test coverage
- [ ] This change adds integration test coverage
- [x] This change has been tested on the latest version of the platform/language or why not

### Checklist

- [x] I have read the [Auth0 general contribution guidelines](https://github.com/auth0/open-source-template/blob/master/GENERAL-CONTRIBUTING.md)
- [x] I have read the [Auth0 Code of Conduct](https://github.com/auth0/open-source-template/blob/master/CODE-OF-CONDUCT.md)
- [x] All existing and new tests complete without errors



Patch:
diff --git a/README.rst b/README.rst
index 1df1c82..c80209a 100644
--- a/README.rst
+++ b/README.rst
@@ -213,6 +213,16 @@ With all in place, the next snippets shows how to verify an RS256 signed ID toke
 Provided something goes wrong, a ``TokenValidationError`` will be raised. In this
 scenario, the ID token should be deemed invalid and its contents not be trusted.
 
+==============
+Error Handling
+==============
+
+When consuming methods from the API clients, the requests could fail for a number of reasons:
+- Invalid data sent as part of the request: An ``Auth0Error` is raised with the error code and description.
+- Global or Client Rate Limit reached: A ``RateLimitError`` is raised and the time at which the limit
+resets is exposed in the ``reset_at`` property. When the header is unset, this value will be ``-1``.
+- Network timeouts: Adjustable by passing a ``timeout`` argument to the client. See the `rate limit docs <https://auth0.com/docs/policies/rate-limits>`_ for details.
+
 Available Management Endpoints
 ==============================
 
diff --git a/auth0/v3/authentication/base.py b/auth0/v3/authentication/base.py
index bfb770f..056c4b5 100644
--- a/auth0/v3/authentication/base.py
+++ b/auth0/v3/authentication/base.py
@@ -3,14 +3,12 @@ import json
 import sys
 import platform
 import requests
-from ..exceptions import Auth0Error
-
+from ..exceptions import Auth0Error, RateLimitError
 
 UNKNOWN_ERROR = 'a0.sdk.internal.unknown'
 
 
 class AuthenticationBase(object):
-
     """Base authentication object providing simple REST methods.
 
     Args:
@@ -69,12 +67,19 @@ class AuthenticationBase(object):
 
 
 class Response(object):
-    def __init__(self, status_code, content):
+    def __init__(self, status_code, content, headers):
         self._status_code = status_code
         self._content = content
+        self._headers = headers
 
     def content(self):
         if self._is_error():
+            if self._status_code == 429:
+                reset_at = int(self._headers.get('x-ratelimit-reset', '-1'))
+                raise RateLimitError(error_code=self._error_code(),
+                                     message=self._error_message(),
+                                     reset_at=reset_at)
+
             raise Auth0Error(status_code=self._status_code,
                              error_code=self._error_code(),
                              message=self._error_message())
@@ -95,7 +100,7 @@ class Response(object):
 class JsonResponse(Response):
     def __init__(self, response):
         content = json.loads(response.text)
-        super(JsonResponse, self).__init__(response.status_code, content)
+        super(JsonResponse, self).__init__(response.status_code, content, response.headers)
 
     def _error_code(self):
         if 'error' in self._content:
@@ -111,7 +116,7 @@ class JsonResponse(Response):
 
 class PlainResponse(Response):
     def __init__(self, response):
-        super(PlainResponse, self).__init__(response.status_code, response.text)
+        super(PlainResponse, self).__init__(response.status_code, response.text, response.headers)
 
     def _error_code(self):
         return UNKNOWN_ERROR
@@ -122,7 +127,7 @@ class PlainResponse(Response):
 
 class EmptyResponse(Response):
     def __init__(self, status_code):
-        super(EmptyResponse, self).__init__(status_code, '')
+        super(EmptyResponse, self).__init__(status_code, '', {})
 
     def _error_code(self):
         return UNKNOWN_ERROR
diff --git a/auth0/v3/exceptions.py b/auth0/v3/exceptions.py
index b086af3..789d2a0 100644
--- a/auth0/v3/exceptions.py
+++ b/auth0/v3/exceptions.py
@@ -8,5 +8,11 @@ class Auth0Error(Exception):
         return '{}: {}'.format(self.status_code, self.message)
 
 
+class RateLimitError(Auth0Error):
+    def __init__(self, error_code, message, reset_at):
+        super(RateLimitError, self).__init__(status_code=429, error_code=error_code, message=message)
+        self.reset_at = reset_at
+
+
 class TokenValidationError(Exception):
     pass
diff --git a/auth0/v3/management/rest.py b/auth0/v3/management/rest.py
index 8e6fd20..e8dc1fa 100644
--- a/auth0/v3/management/rest.py
+++ b/auth0/v3/management/rest.py
@@ -1,16 +1,16 @@
-import sys
-import platform
-import json
 import base64
+import json
+import platform
+import sys
+
 import requests
-from ..exceptions import Auth0Error
 
+from ..exceptions import Auth0Error, RateLimitError
 
 UNKNOWN_ERROR = 'a0.sdk.internal.unknown'
 
 
 class RestClient(object):
-
     """Provides simple methods for handling all RESTful api endpoints.
 
     Args:
@@ -21,6 +21,7 @@ class RestClient(object):
             both values separately or a float to set both to it.
             (defaults to 5.0 for both)
     """
+
     def __init__(self, jwt, telemetry=True, timeout=5.0):
         self.jwt = jwt
         self.timeout = timeout
@@ -96,12 +97,19 @@ class RestClient(object):
 
 
 class Response(object):
-    def __init__(self, status_code, content):
+    def __init__(self, status_code, content, headers):
         self._status_code = status_code
         self._content = content
+        self._headers = headers
 
     def content(self):
         if self._is_error():
+            if self._status_code == 429:
+                reset_at = int(self._headers.get('x-ratelimit-reset', '-1'))
+                raise RateLimitError(error_code=self._error_code(),
+                                     message=self._error_message(),
+                                     reset_at=reset_at)
+
             raise Auth0Error(status_code=self._status_code,
                              error_code=self._error_code(),
                              message=self._error_message())
@@ -122,7 +130,7 @@ class Response(object):
 class JsonResponse(Response):
     def __init__(self, response):
         content = json.loads(response.text)
-        super(JsonResponse, self).__init__(response.status_code, content)
+        super(JsonResponse, self).__init__(response.status_code, content, response.headers)
 
     def _error_code(self):
         if 'errorCode' in self._content:
@@ -141,7 +149,7 @@ class JsonResponse(Response):
 
 class PlainResponse(Response):
     def __init__(self, response):
-        super(PlainResponse, self).__init__(response.status_code, response.text)
+        super(PlainResponse, self).__init__(response.status_code, response.text, response.headers)
 
     def _error_code(self):
         return UNKNOWN_ERROR
@@ -152,7 +160,7 @@ class PlainResponse(Response):
 
 class EmptyResponse(Response):
     def __init__(self, status_code):
-        super(EmptyResponse, self).__init__(status_code, '')
+        super(EmptyResponse, self).__init__(status_code, '', {})
 
     def _error_code(self):
         return UNKNOWN_ERROR


Test patch:
diff --git a/auth0/v3/test/authentication/test_base.py b/auth0/v3/test/authentication/test_base.py
index b2426a1..31b6027 100644
--- a/auth0/v3/test/authentication/test_base.py
+++ b/auth0/v3/test/authentication/test_base.py
@@ -7,7 +7,7 @@ import sys
 import requests
 import unittest
 from ...authentication.base import AuthenticationBase
-from ...exceptions import Auth0Error
+from ...exceptions import Auth0Error, RateLimitError
 
 
 class TestBase(unittest.TestCase):
@@ -27,7 +27,7 @@ class TestBase(unittest.TestCase):
                                            sys.version_info.micro)
 
         client_info = {
-            'name': 'auth0-python', 
+            'name': 'auth0-python',
             'version': auth0_version,
             'env': {
                 'python': python_version
@@ -53,10 +53,10 @@ class TestBase(unittest.TestCase):
         data = ab.post('the-url', data={'a': 'b'}, headers={'c': 'd'})
 
         mock_post.assert_called_with(url='the-url', json={'a': 'b'},
-                headers={'c': 'd', 'Content-Type': 'application/json'}, timeout=(10, 2))
+                                     headers={'c': 'd', 'Content-Type': 'application/json'}, timeout=(10, 2))
 
         self.assertEqual(data, {'x': 'y'})
-        
+
     @mock.patch('requests.post')
     def test_post_with_defaults(self, mock_post):
         ab = AuthenticationBase('auth0.com', telemetry=False)
@@ -68,7 +68,7 @@ class TestBase(unittest.TestCase):
         data = ab.post('the-url')
 
         mock_post.assert_called_with(url='the-url', json=None,
-                headers={'Content-Type': 'application/json'}, timeout=5.0)
+                                     headers={'Content-Type': 'application/json'}, timeout=5.0)
 
         self.assertEqual(data, {'x': 'y'})
 
@@ -109,6 +109,48 @@ class TestBase(unittest.TestCase):
             self.assertEqual(context.exception.error_code, 'e0')
             self.assertEqual(context.exception.message, 'desc')
 
+    @mock.patch('requests.post')
+    def test_post_rate_limit_error(self, mock_post):
+        ab = AuthenticationBase('auth0.com', telemetry=False)
+
+        mock_post.return_value.text = '{"statusCode": 429,' \
+                                      ' "error": "e0",' \
+                                      ' "error_description": "desc"}'
+        mock_post.return_value.status_code = 429
+        mock_post.return_value.headers = {
+            'x-ratelimit-limit': '3',
+            'x-ratelimit-remaining': '6',
+            'x-ratelimit-reset': '9',
+        }
+
+        with self.assertRaises(Auth0Error) as context:
+            ab.post('the-url', data={'a': 'b'}, headers={'c': 'd'})
+
+        self.assertEqual(context.exception.status_code, 429)
+        self.assertEqual(context.exception.error_code, 'e0')
+        self.assertEqual(context.exception.message, 'desc')
+        self.assertIsInstance(context.exception, RateLimitError)
+        self.assertEqual(context.exception.reset_at, 9)
+
+    @mock.patch('requests.post')
+    def test_post_rate_limit_error_without_headers(self, mock_post):
+        ab = AuthenticationBase('auth0.com', telemetry=False)
+
+        mock_post.return_value.text = '{"statusCode": 429,' \
+                                      ' "error": "e0",' \
+                                      ' "error_description": "desc"}'
+        mock_post.return_value.status_code = 429
+        mock_post.return_value.headers = {}
+
+        with self.assertRaises(Auth0Error) as context:
+            ab.post('the-url', data={'a': 'b'}, headers={'c': 'd'})
+
+        self.assertEqual(context.exception.status_code, 429)
+        self.assertEqual(context.exception.error_code, 'e0')
+        self.assertEqual(context.exception.message, 'desc')
+        self.assertIsInstance(context.exception, RateLimitError)
+        self.assertEqual(context.exception.reset_at, -1)
+
     @mock.patch('requests.post')
     def test_post_error_with_code_property(self, mock_post):
         ab = AuthenticationBase('auth0.com', telemetry=False)
@@ -184,7 +226,7 @@ class TestBase(unittest.TestCase):
         data = ab.get('the-url', params={'a': 'b'}, headers={'c': 'd'})
 
         mock_get.assert_called_with(url='the-url', params={'a': 'b'},
-                headers={'c': 'd', 'Content-Type': 'application/json'}, timeout=(10, 2))
+                                    headers={'c': 'd', 'Content-Type': 'application/json'}, timeout=(10, 2))
 
         self.assertEqual(data, {'x': 'y'})
 
@@ -199,7 +241,7 @@ class TestBase(unittest.TestCase):
         data = ab.get('the-url')
 
         mock_get.assert_called_with(url='the-url', params=None,
-                headers={'Content-Type': 'application/json'}, timeout=5.0)
+                                    headers={'Content-Type': 'application/json'}, timeout=5.0)
 
         self.assertEqual(data, {'x': 'y'})
 
diff --git a/auth0/v3/test/management/test_rest.py b/auth0/v3/test/management/test_rest.py
index 4d3ac5b..05a88ac 100644
--- a/auth0/v3/test/management/test_rest.py
+++ b/auth0/v3/test/management/test_rest.py
@@ -7,7 +7,7 @@ import mock
 import requests
 
 from ...management.rest import RestClient
-from ...exceptions import Auth0Error
+from ...exceptions import Auth0Error, RateLimitError
 
 
 class TestRest(unittest.TestCase):
@@ -149,6 +149,48 @@ class TestRest(unittest.TestCase):
         self.assertEqual(context.exception.error_code, 'code')
         self.assertEqual(context.exception.message, 'message')
 
+    @mock.patch('requests.get')
+    def test_get_rate_limit_error(self, mock_get):
+        rc = RestClient(jwt='a-token', telemetry=False)
+
+        mock_get.return_value.text = '{"statusCode": 429,' \
+                                     ' "errorCode": "code",' \
+                                     ' "message": "message"}'
+        mock_get.return_value.status_code = 429
+        mock_get.return_value.headers = {
+            'x-ratelimit-limit': '3',
+            'x-ratelimit-remaining': '6',
+            'x-ratelimit-reset': '9',
+        }
+
+        with self.assertRaises(Auth0Error) as context:
+            rc.get('the/url')
+
+        self.assertEqual(context.exception.status_code, 429)
+        self.assertEqual(context.exception.error_code, 'code')
+        self.assertEqual(context.exception.message, 'message')
+        self.assertIsInstance(context.exception, RateLimitError)
+        self.assertEqual(context.exception.reset_at, 9)
+
+    @mock.patch('requests.get')
+    def test_get_rate_limit_error_without_headers(self, mock_get):
+        rc = RestClient(jwt='a-token', telemetry=False)
+
+        mock_get.return_value.text = '{"statusCode": 429,' \
+                                     ' "errorCode": "code",' \
+                                     ' "message": "message"}'
+        mock_get.return_value.status_code = 429
+        
+        mock_get.return_value.headers = {}
+        with self.assertRaises(Auth0Error) as context:
+            rc.get('the/url')
+
+        self.assertEqual(context.exception.status_code, 429)
+        self.assertEqual(context.exception.error_code, 'code')
+        self.assertEqual(context.exception.message, 'message')
+        self.assertIsInstance(context.exception, RateLimitError)
+        self.assertEqual(context.exception.reset_at, -1)
+
     @mock.patch('requests.post')
     def test_post(self, mock_post):
         rc = RestClient(jwt='a-token', telemetry=False)
@@ -312,7 +354,6 @@ class TestRest(unittest.TestCase):
 
         mock_post.assert_called_once_with('the-url', data=data, files=files, headers=headers, timeout=5.0)
 
-
     @mock.patch('requests.put')
     def test_put(self, mock_put):
         rc = RestClient(jwt='a-token', telemetry=False)
@@ -326,7 +367,7 @@ class TestRest(unittest.TestCase):
 
         response = rc.put(url='the-url', data=data)
         mock_put.assert_called_with('the-url', json=data,
-                                      headers=headers, timeout=5.0)
+                                    headers=headers, timeout=5.0)
 
         self.assertEqual(response, ['a', 'b'])
 
@@ -335,8 +376,8 @@ class TestRest(unittest.TestCase):
         rc = RestClient(jwt='a-token', telemetry=False)
 
         mock_put.return_value.text = '{"statusCode": 999,' \
-                                       ' "errorCode": "code",' \
-                                       ' "message": "message"}'
+                                     ' "errorCode": "code",' \
+                                     ' "message": "message"}'
         mock_put.return_value.status_code = 999
 
         with self.assertRaises(Auth0Error) as context:
@@ -407,7 +448,7 @@ class TestRest(unittest.TestCase):
         mock_delete.return_value.status_code = 200
 
         data = {'some': 'data'}
-        params={'A': 'param', 'B': 'param'}
+        params = {'A': 'param', 'B': 'param'}
 
         response = rc.delete(url='the-url/ID', params=params, data=data)
         mock_delete.assert_called_with('the-url/ID', headers=headers, params=params, json=data, timeout=5.0)
@@ -436,7 +477,7 @@ class TestRest(unittest.TestCase):
             'Content-Type': 'application/json',
             'Authorization': 'Bearer a-token',
         }
-        
+
         self.assertEqual(rc.base_headers, expected_headers)
 
     def test_enabled_telemetry(self):
@@ -454,7 +495,7 @@ class TestRest(unittest.TestCase):
                                            sys.version_info.micro)
 
         client_info = {
-            'name': 'auth0-python', 
+            'name': 'auth0-python',
             'version': auth0_version,
             'env': {
                 'python': python_version