omnata-plugin-runtime 0.4.5a100__tar.gz → 0.4.6__tar.gz

{omnata_plugin_runtime-0.4.5a100 → omnata_plugin_runtime-0.4.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: omnata-plugin-runtime
-Version: 0.4.5a100
+Version: 0.4.6
 Summary: Classes and common runtime components for building and running Omnata Plugins
 Author: James Weakley
 Author-email: james.weakley@omnata.com

{omnata_plugin_runtime-0.4.5a100 → omnata_plugin_runtime-0.4.6}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "omnata-plugin-runtime"
-version = "0.4.5-a100"
+version = "0.4.6"
 description = "Classes and common runtime components for building and running Omnata Plugins"
 authors = ["James Weakley <james.weakley@omnata.com>"]
 readme = "README.md"

{omnata_plugin_runtime-0.4.5a100 → omnata_plugin_runtime-0.4.6}/src/omnata_plugin_runtime/omnata_plugin.py

@@ -21,6 +21,7 @@ import queue
 import threading
 import time
 import hashlib
+import requests
 from abc import ABC, abstractmethod
 from decimal import Decimal
 from functools import partial, wraps, reduce
@@ -263,6 +264,16 @@ class SyncRequest(ABC):
         Takes into account the run deadline and cancellation status.
         This is an alternative which can be used when the target API does not publish specific rate limits, and instead just asks you to respond to 429s as they are sent.
         """
+        if self.test_replay_mode:
+            # when in test replay mode, we want to make the same requests but without any waiting
+            return RateLimitedSession(
+                run_deadline=self._run_deadline,
+                thread_cancellation_token=self._thread_cancellation_token,
+                max_retries=max_retries,
+                backoff_factor=0,
+                statuses_to_include=statuses_to_include,
+                respect_retry_after_header=False
+            )
         return RateLimitedSession(
             run_deadline=self._run_deadline,
             thread_cancellation_token=self._thread_cancellation_token,

{omnata_plugin_runtime-0.4.5a100 → omnata_plugin_runtime-0.4.6}/src/omnata_plugin_runtime/rate_limiting.py

@@ -8,8 +8,9 @@ import re
 import threading
 from email.utils import parsedate_to_datetime
 from logging import getLogger
-from typing import List, Literal, Optional, Dict, Tuple
+from typing import Any, List, Literal, Optional, Dict, Tuple
 import requests
+import time
 from pydantic import Field, root_validator
 from pydantic.json import pydantic_encoder
 from .configuration import SubscriptableBaseModel
@@ -349,6 +350,43 @@ class RetryLaterException(Exception):
         self.message = message
         super().__init__(self.message)
 
+class RetryWithLogging(Retry):
+    """
+        Adding extra logs before making a retry request     
+    """
+    def __init__(self, *args: Any, **kwargs: Any) -> Any:
+        self.thread_cancellation_token:Optional[threading.Event] = None
+        return super().__init__(*args, **kwargs)
+    
+    def new(self, **kw):
+        new_retry = super().new(**kw)
+        new_retry.thread_cancellation_token = self.thread_cancellation_token
+        return new_retry
+    
+    def sleep_for_retry(self, response=None):
+        retry_after = self.get_retry_after(response)
+        if retry_after:
+            logger.info(f"Retrying after {retry_after} seconds due to Retry-After header")
+            if self.thread_cancellation_token is None:
+                time.sleep(retry_after)
+            else:
+                if self.thread_cancellation_token.wait(retry_after):
+                    raise InterruptedWhileWaitingException(message="The sync was interrupted while waiting for rate limiting to expire")
+            return True
+        return False
+
+    def _sleep_backoff(self):
+        backoff = self.get_backoff_time()
+        if backoff <= 0:
+            return
+        logger.info(f"Retrying after {backoff} seconds due to backoff time")
+        if self.thread_cancellation_token is None:
+            time.sleep(backoff)
+        else:
+            if self.thread_cancellation_token.wait(backoff):
+                raise InterruptedWhileWaitingException(message="The sync was interrupted while waiting for rate limiting to expire")
+
+
 class RateLimitedSession(requests.Session):
     """
     Creates a requests session that will automatically handle rate limiting.
@@ -356,7 +394,13 @@ class RateLimitedSession(requests.Session):
     The thread_cancellation_token is observed when waiting, as well as the overall run deadline.
     In case this is used across threads, the retry count will be tracked per request URL (minus query parameters). It will be cleared when the request is successful.
     """
-    def __init__(self, run_deadline:datetime.datetime, thread_cancellation_token:threading.Event, max_retries=5, backoff_factor=1, statuses_to_include:List[int] = [429]):
+    def __init__(self,
+                 run_deadline:datetime.datetime,
+                 thread_cancellation_token:threading.Event,
+                 max_retries=5,
+                 backoff_factor=1,
+                 statuses_to_include:List[int] = [429],
+                 respect_retry_after_header:bool = True):
         super().__init__()
         self.max_retries = max_retries
         self.backoff_factor = backoff_factor
@@ -366,12 +410,14 @@ class RateLimitedSession(requests.Session):
         self.thread_cancellation_token = thread_cancellation_token
         self.statuses_to_include = statuses_to_include
 
-        retry_strategy = Retry(
+        retry_strategy = RetryWithLogging(
             total=max_retries,
             backoff_factor=backoff_factor,
             status_forcelist=statuses_to_include,
-            allowed_methods=["HEAD", "GET", "OPTIONS", "POST", "PUT", "DELETE"]
+            allowed_methods=["HEAD", "GET", "OPTIONS", "POST", "PUT", "DELETE"],
+            respect_retry_after_header=respect_retry_after_header
         )
+        retry_strategy.thread_cancellation_token = thread_cancellation_token
         adapter = HTTPAdapter(max_retries=retry_strategy)
         self.mount("https://", adapter)
         self.mount("http://", adapter)
@@ -398,7 +444,7 @@ class RateLimitedSession(requests.Session):
     def request(self, method, url, **kwargs):
         while True:
             response = super().request(method, url, **kwargs)
-            
+            # TODO: this is probably all redundant as the Retry object should handle this at a lower level (urllib3)
             if response.status_code in self.statuses_to_include:
                 if 'Retry-After' in response.headers:
                     retry_after = response.headers['Retry-After']