feat(client): add retry handling to GraphQL client

nejch · max-wittig · commit 8898c38b97ed · 2024-09-12T17:06:54.000+02:00
diff --git a/gitlab/client.py b/gitlab/client.py
@@ -26,6 +26,7 @@
 
 try:
     import gql
+    import gql.transport.exceptions
     import graphql
     import httpx
 
@@ -746,7 +747,9 @@ def http_request(
             if 200 <= result.status_code < 300:
                 return result.response
 
-            if retry.handle_retry_on_status(result):
+            if retry.handle_retry_on_status(
+                result.status_code, result.headers, result.reason
+            ):
                 continue
 
             error_message = result.content
@@ -1329,4 +1332,28 @@ def execute(
         self, request: Union[str, graphql.Source], *args: Any, **kwargs: Any
     ) -> Any:
         parsed_document = self._gql(request)
-        return self._client.execute(parsed_document, *args, **kwargs)
+        retry = utils.Retry(
+            max_retries=3, obey_rate_limit=True, retry_transient_errors=False
+        )
+
+        while True:
+            try:
+                result = self._client.execute(parsed_document, *args, **kwargs)
+            except gql.transport.exceptions.TransportServerError as e:
+                if retry.handle_retry_on_status(
+                    status_code=e.code, headers=self._transport.response_headers
+                ):
+                    continue
+
+                if e.code == 401:
+                    raise gitlab.exceptions.GitlabAuthenticationError(
+                        response_code=e.code,
+                        error_message=str(e),
+                    )
+
+                raise gitlab.exceptions.GitlabHttpError(
+                    response_code=e.code,
+                    error_message=str(e),
+                )
+
+            return result
diff --git a/gitlab/utils.py b/gitlab/utils.py
@@ -6,12 +6,22 @@
 import traceback
 import urllib.parse
 import warnings
-from typing import Any, Callable, Dict, Iterator, Literal, Optional, Tuple, Type, Union
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Iterator,
+    Literal,
+    MutableMapping,
+    Optional,
+    Tuple,
+    Type,
+    Union,
+)
 
 import requests
 
 from gitlab import const, types
-from gitlab._backends import requests_backend
 
 
 class _StdoutStream:
@@ -100,33 +110,40 @@ def __init__(
         self.retry_transient_errors = retry_transient_errors
 
     def _retryable_status_code(
-        self,
-        result: requests_backend.RequestsResponse,
+        self, status_code: Optional[int], reason: str = ""
     ) -> bool:
-        if result.status_code == 429 and self.obey_rate_limit:
+        if status_code == 429 and self.obey_rate_limit:
             return True
 
         if not self.retry_transient_errors:
             return False
-        if result.status_code in const.RETRYABLE_TRANSIENT_ERROR_CODES:
+        if status_code in const.RETRYABLE_TRANSIENT_ERROR_CODES:
             return True
-        if result.status_code == 409 and "Resource lock" in result.reason:
+        if status_code == 409 and "Resource lock" in reason:
             return True
 
         return False
 
-    def handle_retry_on_status(self, result: requests_backend.RequestsResponse) -> bool:
-        if not self._retryable_status_code(result):
+    def handle_retry_on_status(
+        self,
+        status_code: Optional[int],
+        headers: Optional[MutableMapping[str, str]] = None,
+        reason: str = "",
+    ) -> bool:
+        if not self._retryable_status_code(status_code, reason):
             return False
 
+        if headers is None:
+            headers = {}
+
         # Response headers documentation:
         # https://docs.gitlab.com/ee/user/admin_area/settings/user_and_ip_rate_limits.html#response-headers
         if self.max_retries == -1 or self.cur_retries < self.max_retries:
             wait_time = 2**self.cur_retries * 0.1
-            if "Retry-After" in result.headers:
-                wait_time = int(result.headers["Retry-After"])
-            elif "RateLimit-Reset" in result.headers:
-                wait_time = int(result.headers["RateLimit-Reset"]) - time.time()
+            if "Retry-After" in headers:
+                wait_time = int(headers["Retry-After"])
+            elif "RateLimit-Reset" in headers:
+                wait_time = int(headers["RateLimit-Reset"]) - time.time()
             self.cur_retries += 1
             time.sleep(wait_time)
             return True
diff --git a/requirements-lint.txt b/requirements-lint.txt
@@ -8,6 +8,7 @@ mypy==1.11.2
 pylint==3.2.7
 pytest==8.3.2
 responses==0.25.3
+respx==0.21.1
 types-PyYAML==6.0.12.20240808
 types-requests==2.32.0.20240907
 types-setuptools==74.1.0.20240907
diff --git a/requirements-test.txt b/requirements-test.txt
@@ -7,4 +7,5 @@ pytest-github-actions-annotate-failures==0.2.0
 pytest==8.3.2
 PyYaml==6.0.2
 responses==0.25.3
+respx==0.21.1
 wheel==0.44.0
diff --git a/tests/unit/test_graphql.py b/tests/unit/test_graphql.py
@@ -1,8 +1,15 @@
+import httpx
 import pytest
+import respx
 
 import gitlab
 
 
+@pytest.fixture
+def gl_gql() -> gitlab.GraphQL:
+    return gitlab.GraphQL("https://gitlab.example.com")
+
+
 def test_import_error_includes_message(monkeypatch: pytest.MonkeyPatch):
     monkeypatch.setattr(gitlab.client, "_GQL_INSTALLED", False)
     with pytest.raises(ImportError, match="GraphQL client could not be initialized"):
@@ -12,3 +19,42 @@ def test_import_error_includes_message(monkeypatch: pytest.MonkeyPatch):
 def test_graphql_as_context_manager_exits():
     with gitlab.GraphQL() as gl:
         assert isinstance(gl, gitlab.GraphQL)
+
+
+def test_graphql_retries_on_429_response(
+    gl_gql: gitlab.GraphQL, respx_mock: respx.MockRouter
+):
+    url = "https://gitlab.example.com/api/graphql"
+    responses = [
+        httpx.Response(429, headers={"retry-after": "1"}),
+        httpx.Response(
+            200, json={"data": {"currentUser": {"id": "gid://gitlab/User/1"}}}
+        ),
+    ]
+    respx_mock.post(url).mock(side_effect=responses)
+    gl_gql.execute("query {currentUser {id}}")
+
+
+def test_graphql_raises_when_max_retries_exceeded(
+    gl_gql: gitlab.GraphQL, respx_mock: respx.MockRouter
+):
+    url = "https://gitlab.example.com/api/graphql"
+    responses = [
+        httpx.Response(502),
+        httpx.Response(502),
+        httpx.Response(502),
+        httpx.Response(502),
+        httpx.Response(502),
+    ]
+    respx_mock.post(url).mock(side_effect=responses)
+    with pytest.raises(gitlab.GitlabHttpError):
+        gl_gql.execute("query {currentUser {id}}")
+
+
+def test_graphql_raises_on_401_response(
+    gl_gql: gitlab.GraphQL, respx_mock: respx.MockRouter
+):
+    url = "https://gitlab.example.com/api/graphql"
+    respx_mock.post(url).mock(return_value=httpx.Response(401))
+    with pytest.raises(gitlab.GitlabAuthenticationError):
+        gl_gql.execute("query {currentUser {id}}")
diff --git a/tests/unit/test_retry.py b/tests/unit/test_retry.py
@@ -2,34 +2,24 @@
 from unittest import mock
 
 import pytest
-import requests
 
 from gitlab import utils
-from gitlab._backends import requests_backend
 
 
 def test_handle_retry_on_status_ignores_unknown_status_code():
     retry = utils.Retry(max_retries=1, retry_transient_errors=True)
-    response = requests.Response()
-    response.status_code = 418
-    backend_response = requests_backend.RequestsResponse(response)
-
-    assert retry.handle_retry_on_status(backend_response) is False
+    assert retry.handle_retry_on_status(418) is False
 
 
 def test_handle_retry_on_status_accepts_retry_after_header(
     monkeypatch: pytest.MonkeyPatch,
 ):
     mock_sleep = mock.Mock()
     monkeypatch.setattr(time, "sleep", mock_sleep)
-
     retry = utils.Retry(max_retries=1)
-    response = requests.Response()
-    response.status_code = 429
-    response.headers["Retry-After"] = "1"
-    backend_response = requests_backend.RequestsResponse(response)
+    headers = {"Retry-After": "1"}
 
-    assert retry.handle_retry_on_status(backend_response) is True
+    assert retry.handle_retry_on_status(429, headers=headers) is True
     assert isinstance(mock_sleep.call_args[0][0], int)
 
 
@@ -40,19 +30,12 @@ def test_handle_retry_on_status_accepts_ratelimit_reset_header(
     monkeypatch.setattr(time, "sleep", mock_sleep)
 
     retry = utils.Retry(max_retries=1)
-    response = requests.Response()
-    response.status_code = 429
-    response.headers["RateLimit-Reset"] = str(int(time.time() + 1))
-    backend_response = requests_backend.RequestsResponse(response)
+    headers = {"RateLimit-Reset": str(int(time.time() + 1))}
 
-    assert retry.handle_retry_on_status(backend_response) is True
+    assert retry.handle_retry_on_status(429, headers=headers) is True
     assert isinstance(mock_sleep.call_args[0][0], float)
 
 
 def test_handle_retry_on_status_returns_false_when_max_retries_reached():
     retry = utils.Retry(max_retries=0)
-    response = requests.Response()
-    response.status_code = 429
-    backend_response = requests_backend.RequestsResponse(response)
-
-    assert retry.handle_retry_on_status(backend_response) is False
+    assert retry.handle_retry_on_status(429) is False
