feat: improved handling of retries

GMorris-professional · rhamzeh · commit fc600634c948 · 2025-04-30T11:35:39.000-04:00
diff --git a/openfga_sdk/api_client.py b/openfga_sdk/api_client.py
@@ -42,7 +42,7 @@
 DEFAULT_USER_AGENT = "openfga-sdk python/0.9.3"
 
 
-def random_time(loop_count, min_wait_in_ms):
+def random_time(loop_count, min_wait_in_ms) -> float:
     """
     Helper function to return the time (in s) to wait before retry
     """
@@ -253,11 +253,21 @@ async def __call_api(
             )
             else 0
         )
+        max_wait_in_sec = (
+            self.configuration.retry_params.max_wait_in_sec
+            if (
+                self.configuration.retry_params is not None
+                and self.configuration.retry_params.max_wait_in_sec is not None
+            )
+            else 120
+        )
         if _retry_params is not None:
             if _retry_params.max_retry is not None:
                 max_retry = _retry_params.max_retry
             if _retry_params.min_wait_in_ms is not None:
                 max_retry = _retry_params.min_wait_in_ms
+            if _retry_params.max_wait_in_sec is not None:
+                max_wait_in_sec = _retry_params.max_wait_in_sec
 
         _telemetry_attributes = TelemetryAttributes.fromRequest(
             user_agent=self.user_agent,
@@ -300,7 +310,14 @@ async def __call_api(
                         configuration=self.configuration.telemetry,
                     )
 
-                    await asyncio.sleep(random_time(retry, min_wait_in_ms))
+                    try:
+                        wait_time_in_sec = self._parse_retry_after_header(e.header)
+                    except ValueError:
+                        wait_time_in_sec = min(
+                            random_time(retry, min_wait_in_ms), max_wait_in_sec
+                        )
+
+                    await asyncio.sleep(wait_time_in_sec)
 
                     continue
                 e.body = e.body.decode("utf-8")
@@ -395,6 +412,25 @@ async def __call_api(
             else:
                 return (return_data, response_data.status, response_data.headers)
 
+    def _parse_retry_after_header(self, headers) -> int:
+        retry_after_header = headers.get("retry-after")
+        if not retry_after_header:
+            raise ValueError("Retry-After header is not present")
+
+        try:
+            parsed_http_date = self.__deserialize_datetime(retry_after_header).replace(
+                tzinfo=datetime.timezone.utc
+            )
+            now = datetime.datetime.now(datetime.timezone.utc)
+            wait_time_in_sec = (parsed_http_date - now).total_seconds()
+        except ApiException:
+            wait_time_in_sec = int(retry_after_header)
+
+        if wait_time_in_sec > 1800 or wait_time_in_sec < 1:
+            raise ValueError("Retry-After header is invalid")
+
+        return math.ceil(wait_time_in_sec)
+
     def sanitize_for_serialization(self, obj):
         """Builds a JSON POST object.
 
@@ -825,7 +861,7 @@ def __deserialize_datetime(self, string):
             return parse(string)
         except ImportError:
             return string
-        except ValueError:
+        except (TypeError, ValueError):
             raise rest.ApiException(
                 status=0,
                 reason=(f"Failed to parse `{string}` as datetime object"),
diff --git a/openfga_sdk/configuration.py b/openfga_sdk/configuration.py
@@ -41,11 +41,13 @@ class RetryParams:
 
     :param max_retry: Maximum number of retry
     :param min_wait_in_ms: Minimum wait (in ms) between retry
+    :param max_wait_in_sec: Maximum wait (in seconds) between retry
     """
 
-    def __init__(self, max_retry=3, min_wait_in_ms=100):
+    def __init__(self, max_retry=3, min_wait_in_ms=100, max_wait_in_sec=120):
         self._max_retry = max_retry
         self._min_wait_in_ms = min_wait_in_ms
+        self._max_wait_in_sec = max_wait_in_sec
 
     @property
     def max_retry(self):
@@ -95,6 +97,25 @@ def min_wait_in_ms(self, value):
 
         self._min_wait_in_ms = value
 
+    @property
+    def max_wait_in_sec(self):
+        """
+        Return the maximum allowed wait (in seconds) in between retry
+        """
+        return self._max_wait_in_sec
+
+    @max_wait_in_sec.setter
+    def max_wait_in_sec(self, value):
+        """
+        Update the maximum allowed wait (in seconds) in between retry
+        """
+        if not isinstance(value, int) or value < 0:
+            raise FgaValidationException(
+                "RetryParams.max_wait_in_sec must be an integer greater than or equal to 0"
+            )
+
+        self._max_wait_in_sec = value
+
 
 class Configuration:
     """NOTE: This class is auto generated by OpenAPI Generator
diff --git a/openfga_sdk/exceptions.py b/openfga_sdk/exceptions.py
@@ -10,20 +10,22 @@
 NOTE: This file was auto generated by OpenAPI Generator (https://openapi-generator.tech). DO NOT EDIT.
 """
 
-# Specifc FGA header to be parsed
+# Specific FGA header to be parsed
 X_RATELIMIT_LIMIT = "x-ratelimit-limit"
 X_RATELIMIT_REMAINING = "x_ratelimit_remaining"
 X_RATELIMIT_RESET = "x_ratelimit_reset"
 FGA_REQUEST_ID = "fga-request-id"
 FGA_QUERY_DURATION_MS = "fga-query-duration-ms"
 OPENFGA_AUTHORIZATION_MODEL_ID = "openfga_authorization_model_id"
+RETRY_AFTER = "retry-after"
 RESPONSE_HEADERS_TO_KEEP = [
     X_RATELIMIT_LIMIT,
     X_RATELIMIT_REMAINING,
     X_RATELIMIT_RESET,
     FGA_REQUEST_ID,
     FGA_QUERY_DURATION_MS,
     OPENFGA_AUTHORIZATION_MODEL_ID,
+    RETRY_AFTER,
 ]
 
 
diff --git a/openfga_sdk/sync/api_client.py b/openfga_sdk/sync/api_client.py
@@ -47,7 +47,6 @@ def random_time(loop_count, min_wait_in_ms) -> float:
     """
     minimum = math.ceil(2**loop_count * min_wait_in_ms)
     maximum = math.ceil(2 ** (loop_count + 1) * min_wait_in_ms)
-
     return random.randrange(minimum, maximum) / 1000
 
 
@@ -253,11 +252,21 @@ def __call_api(
             )
             else 0
         )
+        max_wait_in_sec = (
+            self.configuration.retry_params.max_wait_in_sec
+            if (
+                self.configuration.retry_params is not None
+                and self.configuration.retry_params.max_wait_in_sec is not None
+            )
+            else 120
+        )
         if _retry_params is not None:
             if _retry_params.max_retry is not None:
                 max_retry = _retry_params.max_retry
             if _retry_params.min_wait_in_ms is not None:
                 max_retry = _retry_params.min_wait_in_ms
+            if _retry_params.max_wait_in_sec is not None:
+                max_wait_in_sec = _retry_params.max_wait_in_sec
 
         _telemetry_attributes = TelemetryAttributes.fromRequest(
             user_agent=self.user_agent,
@@ -300,8 +309,14 @@ def __call_api(
                         configuration=self.configuration.telemetry,
                     )
 
-                    time.sleep(random_time(retry, min_wait_in_ms))
+                    try:
+                        wait_time_in_sec = self._parse_retry_after_header(e.header)
+                    except ValueError:
+                        wait_time_in_sec = min(
+                            random_time(retry, min_wait_in_ms), max_wait_in_sec
+                        )
 
+                    time.sleep(wait_time_in_sec)
                     continue
                 e.body = e.body.decode("utf-8")
                 response_type = response_types_map.get(e.status, None)
@@ -395,6 +410,25 @@ def __call_api(
             else:
                 return (return_data, response_data.status, response_data.headers)
 
+    def _parse_retry_after_header(self, headers) -> int:
+        retry_after_header = headers.get("retry-after")
+        if not retry_after_header:
+            raise ValueError("Retry-After header is not present")
+
+        try:
+            parsed_http_date = self.__deserialize_datetime(retry_after_header).replace(
+                tzinfo=datetime.timezone.utc
+            )
+            now = datetime.datetime.now(datetime.timezone.utc)
+            wait_time_in_sec = (parsed_http_date - now).total_seconds()
+        except ApiException:
+            wait_time_in_sec = int(retry_after_header)
+
+        if wait_time_in_sec > 1800 or wait_time_in_sec < 1:
+            raise ValueError("Retry-After header is invalid")
+
+        return math.ceil(wait_time_in_sec)
+
     def sanitize_for_serialization(self, obj):
         """Builds a JSON POST object.
 
@@ -825,7 +859,7 @@ def __deserialize_datetime(self, string):
             return parse(string)
         except ImportError:
             return string
-        except ValueError:
+        except (TypeError, ValueError):
             raise rest.ApiException(
                 status=0,
                 reason=(f"Failed to parse `{string}` as datetime object"),
diff --git a/test/api/open_fga_api_test.py b/test/api/open_fga_api_test.py
diff --git a/test/sync/open_fga_api_test.py b/test/sync/open_fga_api_test.py
