Merge pull request #488 from siemens/feat/rate-limit

Gauvain Pocentek · web-flow · commit 86a825143fda · 2018-04-24T11:25:08.000+02:00
feat: obey the rate limit
diff --git a/docs/api-usage.rst b/docs/api-usage.rst
@@ -326,3 +326,26 @@ The following sample illustrates how to use a client-side certificate:
 
 Reference:
 http://docs.python-requests.org/en/master/user/advanced/#client-side-certificates
+
+Rate limits
+-----------
+
+python-gitlab will obey the rate limit of the GitLab server by default.
+On receiving a 429 response (Too Many Requests), python-gitlab will sleep for the amount of time
+in the Retry-After header, that GitLab sends back.
+
+If you don't want to wait, you can disable the rate-limiting feature, by supplying the
+``obey_rate_limit`` argument.
+
+.. code-block:: python
+
+   import gitlab
+   import requests
+
+   gl = gitlab.gitlab(url, token, api_version=4)
+   gl.projects.list(all=True, obey_rate_limit=False)
+
+
+.. warning::
+
+   You will get an Exception, if you then go over the rate limit of your GitLab instance.
diff --git a/gitlab/__init__.py b/gitlab/__init__.py
@@ -23,6 +23,7 @@
 import itertools
 import json
 import re
+import time
 import warnings
 
 import requests
@@ -698,24 +699,35 @@ def copy_dict(dest, src):
         prepped.url = sanitized_url(prepped.url)
         settings = self.session.merge_environment_settings(
             prepped.url, {}, streamed, verify, None)
-        result = self.session.send(prepped, timeout=timeout, **settings)
 
-        if 200 <= result.status_code < 300:
-            return result
+        # obey the rate limit by default
+        obey_rate_limit = kwargs.get("obey_rate_limit", True)
 
-        try:
-            error_message = result.json()['message']
-        except (KeyError, ValueError, TypeError):
-            error_message = result.content
-
-        if result.status_code == 401:
-            raise GitlabAuthenticationError(response_code=result.status_code,
-                                            error_message=error_message,
-                                            response_body=result.content)
-
-        raise GitlabHttpError(response_code=result.status_code,
-                              error_message=error_message,
-                              response_body=result.content)
+        while True:
+            result = self.session.send(prepped, timeout=timeout, **settings)
+
+            if 200 <= result.status_code < 300:
+                return result
+
+            if 429 == result.status_code and obey_rate_limit:
+                wait_time = int(result.headers["Retry-After"])
+                time.sleep(wait_time)
+                continue
+
+            try:
+                error_message = result.json()['message']
+            except (KeyError, ValueError, TypeError):
+                error_message = result.content
+
+            if result.status_code == 401:
+                raise GitlabAuthenticationError(
+                    response_code=result.status_code,
+                    error_message=error_message,
+                    response_body=result.content)
+
+            raise GitlabHttpError(response_code=result.status_code,
+                                  error_message=error_message,
+                                  response_body=result.content)
 
     def http_get(self, path, query_data={}, streamed=False, **kwargs):
         """Make a GET request to the Gitlab server.
diff --git a/tools/python_test_v4.py b/tools/python_test_v4.py
@@ -646,3 +646,28 @@
 
 # events
 gl.events.list()
+
+# rate limit
+settings = gl.settings.get()
+settings.throttle_authenticated_api_enabled = True
+settings.throttle_authenticated_api_requests_per_period = 1
+settings.throttle_authenticated_api_period_in_seconds = 3
+settings.save()
+projects = list()
+for i in range(0, 20):
+    projects.append(gl.projects.create(
+        {'name': str(i) + "ok"}))
+
+error_message = None
+for i in range(20, 40):
+    try:
+        projects.append(
+            gl.projects.create(
+                {'name': str(i) + 'shouldfail'}, obey_rate_limit=False))
+    except gitlab.GitlabCreateError as e:
+        error_message = e.error_message
+        break
+assert 'Retry later' in error_message
+[current_project.delete() for current_project in projects]
+settings.throttle_authenticated_api_enabled = False
+settings.save()
