refreshes credentials for irsa

Author: Avi-Robusta

prometrix/connect/aws_connect.py

@@ -34,15 +34,20 @@ def __init__(
         if access_key and secret_key:
             # Backwards compatibility: use static keys
             self._credentials = Credentials(access_key, secret_key, token)
+            self._has_static_keys = True
+            self._session = None
         else:
             # IRSA
             session = boto3.Session()
             creds = session.get_credentials()
             if not creds:
                 raise RuntimeError("No AWS credentials found (neither static keys nor IRSA)")
             self._credentials = creds
+            self._has_static_keys = False
+            self._session = session
 
         role_to_assume = assume_role_arn or AWS_ASSUME_ROLE
+        self._role_to_assume = role_to_assume
         if role_to_assume:
             self._assume_role(role_to_assume)
 
@@ -93,6 +98,24 @@ def signed_request(
             params=params,
         )
 
+    def _refresh_credentials(self) -> None:
+        """
+            Boto should automatically refresh expired credentials but when assuming role it cant be done automatically
+        """
+        try:
+            if not self._has_static_keys and self._session is not None:
+                # this is also needed for assume role if base credentials fails
+                refreshed = self._session.get_credentials()
+                if refreshed:
+                    self._credentials = refreshed
+        except Exception:
+            logging.exception("Failed to refresh session credentials")
+        if self._role_to_assume:
+            try:
+                self._assume_role(self._role_to_assume)
+            except Exception:
+                logging.exception("Failed to refresh assume role")
+
     def _custom_query(self, query: str, params: dict = None):
         """
         Send a custom query to a Prometheus Host.
@@ -121,6 +144,16 @@ def _custom_query(self, query: str, params: dict = None):
             verify=self.ssl_verification,
             headers=self.headers,
         )
+        if response is not None and response.status_code == 403:
+            self._refresh_credentials()
+            response = self.signed_request(
+                method="POST",
+                url="{0}/api/v1/query".format(self.url),
+                data={**{"query": query}, **params},
+                params={},
+                verify=self.ssl_verification,
+                headers=self.headers,
+            )
         return response
 
     def safe_custom_query_range(
@@ -162,6 +195,18 @@ def safe_custom_query_range(
             params={},
             headers=self.headers,
         )
+        if response is not None and response.status_code == 403:
+            self._refresh_credentials()
+            response = self.signed_request(
+                method="POST",
+                url="{0}/api/v1/query_range".format(self.url),
+                data={
+                    **{"query": query, "start": start, "end": end, "step": step},
+                    **params,
+                },
+                params={},
+                headers=self.headers,
+            )
         if response.status_code == 200:
             return response.json()["data"]
         else: