From 0d2650158d945f58e590dbfafbd9e1a84487f0d9 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 | 30 ++++++++++++++++++++++--------
 1 file changed, 22 insertions(+), 8 deletions(-)

diff --git a/zulip/zulip/__init__.py b/zulip/zulip/__init__.py
index 725896e34..be118ca4e 100644
--- a/zulip/zulip/__init__.py
+++ b/zulip/zulip/__init__.py
@@ -697,6 +697,7 @@ 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]:
@@ -707,14 +708,27 @@ 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"
+                or "retry-after" not in result
+            ):
+                break
+            wait_before_retry_secs: float = result["retry-after"]
+            logger.warning("Hit API rate limit, waiting for %f seconds...", wait_before_retry_secs)
+            time.sleep(wait_before_retry_secs)
+
+        return result
 
     def call_on_each_event(
         self,