omnata-plugin-runtime 0.3.27a80__tar.gz → 0.3.28__tar.gz

{omnata_plugin_runtime-0.3.27a80 → omnata_plugin_runtime-0.3.28}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: omnata-plugin-runtime
-Version: 0.3.27a80
+Version: 0.3.28
 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.3.27a80 → omnata_plugin_runtime-0.3.28}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "omnata-plugin-runtime"
-version = "0.3.27-a80"
+version = "0.3.28"
 description = "Classes and common runtime components for building and running Omnata Plugins"
 authors = ["James Weakley <james.weakley@omnata.com>"]
 readme = "README.md"
@@ -32,6 +32,7 @@ wheel = "<=0.41.2" # latest version available on Snowflake Anaconda
 [tool.poetry.dev-dependencies]
 pytest = "^6.2.4"
 deepdiff = "^6"
+requests-mock = ">=1.9.3"
 
 [tool.pytest.ini_options]
 addopts = ["--import-mode=importlib"]

{omnata_plugin_runtime-0.3.27a80 → omnata_plugin_runtime-0.3.28}/src/omnata_plugin_runtime/omnata_plugin.py

@@ -77,6 +77,7 @@ from .rate_limiting import (
     HttpMethodType,
     InterruptedWhileWaitingException,
     RateLimitState,
+    RateLimitedSession
 )
 
 logger = getLogger(__name__)
@@ -251,6 +252,25 @@ class SyncRequest(ABC):
             f"thread_cancellation_token: {self._thread_cancellation_token.is_set()}"
         )
 
+    def get_ratelimit_retrying_http_session(self,
+                                            max_retries: int = 5,
+                                            backoff_factor: int = 1,
+                                            statuses_to_include: List[int] = [429]
+):
+        """
+        Returns a requests.Session object which can respond to 429 responses by waiting and retrying.
+        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.
+        """
+        return RateLimitedSession(
+            run_deadline=self._run_deadline,
+            thread_cancellation_token=self._thread_cancellation_token,
+            max_retries=max_retries,
+            backoff_factor=backoff_factor,
+            statuses_to_include=statuses_to_include
+        )
+    
+
     def __apply_results_worker(self, cancellation_token:threading.Event):
         """
         Designed to be run in a thread, this method polls the results every 20 seconds and sends them back to Snowflake.

{omnata_plugin_runtime-0.3.27a80 → omnata_plugin_runtime-0.3.28}/src/omnata_plugin_runtime/rate_limiting.py

@@ -13,6 +13,9 @@ import requests
 from pydantic import Field, root_validator
 from pydantic.json import pydantic_encoder
 from .configuration import SubscriptableBaseModel
+import pytz
+from requests.adapters import HTTPAdapter
+from urllib3.util.retry import Retry
 
 logger = getLogger(__name__)
 
@@ -350,13 +353,96 @@ class RetryLaterException(Exception):
         self.message = message
         super().__init__(self.message)
 
+class RateLimitedSession(requests.Session):
+    """
+    Creates a requests session that will automatically handle rate limiting.
+    It will retry requests that return a 429 status code, and will wait until the rate limit is reset.
+    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]):
+        super().__init__()
+        self.max_retries = max_retries
+        self.backoff_factor = backoff_factor
+        self.retries:Dict[str,int] = {}
+        self._retries_lock = threading.Lock()
+        self.run_deadline = run_deadline
+        self.thread_cancellation_token = thread_cancellation_token
+        self.statuses_to_include = statuses_to_include
+
+        retry_strategy = Retry(
+            total=max_retries,
+            backoff_factor=backoff_factor,
+            status_forcelist=statuses_to_include,
+            allowed_methods=["HEAD", "GET", "OPTIONS", "POST", "PUT", "DELETE"]
+        )
+        adapter = HTTPAdapter(max_retries=retry_strategy)
+        self.mount("https://", adapter)
+        self.mount("http://", adapter)
+    
+    def set_retries(self, url:str, retries:int):
+        # ensure that the query parameters are not included in the key
+        url = re.sub(r'\?.*$', '', url)
+        with self._retries_lock:
+            self.retries[url] = retries
+    
+    def get_retries(self, url:str):
+        # ensure that the query parameters are not included in the key
+        url = re.sub(r'\?.*$', '', url)
+        with self._retries_lock:
+            return self.retries.get(url,0)
+    
+    def increment_retries(self, url:str) -> int:
+        # ensure that the query parameters are not included in the key
+        url = re.sub(r'\?.*$', '', url)
+        with self._retries_lock:
+            self.retries[url] = self.retries.get(url,0) + 1
+            return self.retries[url]
+
+    def request(self, method, url, **kwargs):
+        while True:
+            response = super().request(method, url, **kwargs)
+            
+            if response.status_code in self.statuses_to_include:
+                if 'Retry-After' in response.headers:
+                    retry_after = response.headers['Retry-After']
+                    if retry_after.isdigit():
+                        wait_time = int(retry_after)
+                    else:
+                        # Retry-After can be a date in the format specified in https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Retry-After
+                        # e.g. Fri, 31 Dec 1999 23:59:59 GMT
+                        # we'll parse it using the standard datetime parser
+                        retry_datetime = datetime.datetime.strptime(retry_after, "%a, %d %b %Y %H:%M:%S %Z")
+                        retry_datetime = retry_datetime.replace(tzinfo=pytz.UTC)
+                        current_datetime = datetime.datetime.now(pytz.UTC)
+                        wait_time = (retry_datetime - current_datetime).total_seconds()
+                    # first check that the wait time is not longer than the run deadline
+                    if datetime.datetime.now(pytz.UTC) + datetime.timedelta(seconds=wait_time) > self.run_deadline:
+                        raise InterruptedWhileWaitingException(message=f"The rate limiting wait time ({wait_time} seconds) would exceed the run deadline")
+                    logger.info(f"Waiting for {wait_time} seconds before retrying {method} request to {url}")
+                    # if wait() returns true, it means that the thread was cancelled
+                    if self.thread_cancellation_token.wait(wait_time):
+                        raise InterruptedWhileWaitingException(message="The sync was interrupted while waiting for rate limiting to expire")
+                else:
+                    current_url_retries = self.increment_retries(url)
+                    if current_url_retries >= self.max_retries:
+                        raise requests.exceptions.RetryError(f"Maximum retries reached: {self.max_retries}")
+                    backoff_time = self.backoff_factor * (2 ** (current_url_retries - 1))
+                    if datetime.datetime.now(pytz.UTC) + datetime.timedelta(seconds=backoff_time) > self.run_deadline:
+                        raise InterruptedWhileWaitingException(message=f"The rate limiting backoff time ({backoff_time} seconds) would exceed the run deadline")
+                    logger.info(f"Waiting for {backoff_time} seconds before retrying {method} request to {url}")
+                    if self.thread_cancellation_token.wait(backoff_time):
+                        raise InterruptedWhileWaitingException(message="The sync was interrupted while waiting for rate limiting backoff")
+            else:
+                self.set_retries(url,0)  # Reset retries if the request is successful
+                return response
+
 
 class InterruptedWhileWaitingException(Exception):
     """
     Indicates that while waiting for rate limiting to expire, the sync was interrupted
     """
-    def __init__(self):
-        message = "The sync was interrupted while waiting"
+    def __init__(self, message:str = "The sync was interrupted while waiting"):
         self.message = message
         super().__init__(self.message)