From b870097d9df1e9f21a27b1a6b80f28d17a558f33 Mon Sep 17 00:00:00 2001
From: Robert Imschweiler <ro-i@posteo.de>
Date: Mon, 28 Jun 2021 22:59:40 +0200
Subject: [PATCH] zulip.Client.call_endpoint: Add retry_on_rate_limit_error.

If the call_endpoint method is called with the
"retry_on_rate_limit_error" parameter set to true, wait and retry
automatically on rate limit errors.
See https://chat.zulip.org/#narrow/stream/378-api-design/topic/
Rate.20limits/near/1217048 for the discussion.
---
 zulip/zulip/__init__.py | 33 +++++++++++++++++++++++++--------
 1 file changed, 25 insertions(+), 8 deletions(-)

diff --git a/zulip/zulip/__init__.py b/zulip/zulip/__init__.py
index 725896e34..536102346 100644
--- a/zulip/zulip/__init__.py
+++ b/zulip/zulip/__init__.py
@@ -697,9 +697,12 @@ def call_endpoint(
         method: str = "POST",
         request: Optional[Dict[str, Any]] = None,
         longpolling: bool = False,
+        retry_on_rate_limit_error: bool = False,
         files: Optional[List[IO[Any]]] = None,
         timeout: Optional[float] = None,
     ) -> Dict[str, Any]:
+        secs: float
+
         if request is None:
             request = dict()
         marshalled_request = {}
@@ -707,14 +710,28 @@ def call_endpoint(
             if v is not None:
                 marshalled_request[k] = v
         versioned_url = API_VERSTRING + (url if url is not None else "")
-        return self.do_api_query(
-            marshalled_request,
-            versioned_url,
-            method=method,
-            longpolling=longpolling,
-            files=files,
-            timeout=timeout,
-        )
+
+        while True:
+            result = self.do_api_query(
+                marshalled_request,
+                versioned_url,
+                method=method,
+                longpolling=longpolling,
+                files=files,
+                timeout=timeout,
+            )
+            if not retry_on_rate_limit_error or result["result"] == "success":
+                break
+            elif "code" in result and result["code"] == "RATE_LIMIT_HIT":
+                secs = result["retry-after"]
+            elif "X-RateLimit-Reset" in result:
+                secs = float(result["X-RateLimit-Reset"])
+            else:
+                break
+            logger.warning("hit API rate limit, waiting for %f seconds...", secs)
+            time.sleep(secs)
+
+        return result
 
     def call_on_each_event(
         self,