From b61941b64ca353cd71d6038d9c1855ab21a5218c Mon Sep 17 00:00:00 2001
From: Mark Jan van Kampen <markjanvankampen@gmail.com>
Date: Wed, 21 Feb 2024 20:52:37 +0100
Subject: [PATCH] chore(deps): upgrades pyrate-limiter (#1653)

---
 pyproject.toml                          | 2 +-
 src/kili/core/graphql/graphql_client.py | 7 ++++---
 tests/unit/test_graphql_client.py       | 4 ++--
 3 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index d54a1d247..2fc0e6e1b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -48,7 +48,7 @@ dependencies = [
   "gql[requests,websockets] >= 3.5.0b5, < 4.0.0",
   "filelock >= 3.0.0, < 4.0.0",
   "pip-system-certs >= 4.0.0, < 5.0.0; platform_system=='Windows'",
-  "pyrate-limiter >= 2, < 3",
+  "pyrate-limiter >= 3, < 4",
 ]
 urls = { homepage = "https://github.com/kili-technology/kili-python-sdk" }
 
diff --git a/src/kili/core/graphql/graphql_client.py b/src/kili/core/graphql/graphql_client.py
index ab502da9b..fdb690034 100644
--- a/src/kili/core/graphql/graphql_client.py
+++ b/src/kili/core/graphql/graphql_client.py
@@ -14,7 +14,7 @@
 from gql.transport.requests import RequestsHTTPTransport
 from gql.transport.requests import log as gql_requests_logger
 from graphql import DocumentNode, print_schema
-from pyrate_limiter import Duration, Limiter, RequestRate
+from pyrate_limiter import Duration, Limiter, Rate
 from tenacity import (
     retry,
     retry_all,
@@ -39,7 +39,7 @@
 # they need to be shared between all instances of Kili client within the same process
 
 # rate limiter to avoid sending too many queries to the backend
-_limiter = Limiter(RequestRate(MAX_CALLS_PER_MINUTE, Duration.MINUTE))
+_limiter = Limiter(Rate(MAX_CALLS_PER_MINUTE, Duration.MINUTE), max_delay=120 * 1000)
 
 # mutex to avoid multiple threads sending queries to the backend at the same time
 _execute_lock = threading.Lock()
@@ -303,7 +303,8 @@ def _execute_with_retries(
     def _raw_execute(
         self, document: DocumentNode, variables: Optional[Dict], **kwargs
     ) -> Dict[str, Any]:
-        with _limiter.ratelimit("GraphQLClient.execute", delay=True), _execute_lock:
+        _limiter.try_acquire("GraphQLClient.execute")
+        with _execute_lock:
             return self._gql_client.execute(
                 document=document,
                 variable_values=variables,
diff --git a/tests/unit/test_graphql_client.py b/tests/unit/test_graphql_client.py
index a9054bbf0..46f73df80 100644
--- a/tests/unit/test_graphql_client.py
+++ b/tests/unit/test_graphql_client.py
@@ -10,7 +10,7 @@
 import pytest_mock
 from gql import Client
 from gql.transport import exceptions
-from pyrate_limiter import Duration, Limiter, RequestRate
+from pyrate_limiter import Duration, Limiter, Rate
 
 from kili.adapters.http_client import HttpClient
 from kili.core.constants import MAX_CALLS_PER_MINUTE
@@ -155,7 +155,7 @@ def test_rate_limiting(mocker: pytest_mock.MockerFixture):
     mocker.patch("kili.core.graphql.graphql_client.gql", side_effect=lambda x: x)
     mocker.patch(
         "kili.core.graphql.graphql_client._limiter",
-        new=Limiter(RequestRate(MAX_CALLS_PER_MINUTE, Duration.SECOND * 5)),
+        new=Limiter(Rate(MAX_CALLS_PER_MINUTE, Duration.SECOND * 5), max_delay=120 * 1000),
     )
     client = GraphQLClient(
         endpoint="",