fix refresh creds

Author: Avi-Robusta

prometrix/connect/aws_connect.py

@@ -31,25 +31,43 @@ def __init__(
         self.region = region
         self.service_name = service_name
 
-        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
+        self._initial_access_key = access_key
+        self._initial_secret_key = secret_key
+        self._initial_token = token
+        self._has_static_keys = bool(access_key and secret_key)
+        self._session = None
+        self._credentials = None
 
         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)
+
+        self._last_init_at = None
+
+        self.init_credentials()
+
+    def init_credentials(self) -> None:
+
+        try:
+            if self._has_static_keys:
+                self._credentials = Credentials(self._initial_access_key, self._initial_secret_key, self._initial_token)
+                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._session = session
+
+            role_to_assume = self._role_to_assume
+            if role_to_assume:
+                self._assume_role(role_to_assume)
+
+            self._last_init_at = datetime.utcnow()
+        except Exception:
+            logging.exception("Failed to initialize credentials")
+            raise
 
     def _assume_role(self, role_arn: str) -> None:
         try:
@@ -80,6 +98,12 @@ def _assume_role(self, role_arn: str) -> None:
 
     def _build_auth(self) -> SigV4Auth:
         """Builds fresh SigV4 auth with current credentials (handles rotation)."""
+        try:
+            if self._last_init_at is None or (datetime.utcnow() - self._last_init_at).total_seconds() >= 900:
+                logging.debug("Fifteen minutes passed; re-initializing AWS credentials")
+                self.init_credentials()
+        except Exception:
+            logging.exception("Time-based credential refresh failed")
         frozen = self._credentials.get_frozen_credentials()
         return SigV4Auth(frozen, self.service_name, self.region)
 
@@ -126,7 +150,8 @@ def _request_with_refresh(self, *, method, url, data=None, params=None, headers=
             headers=headers,
         )
         if resp is not None and resp.status_code in (400, 401, 403):
-            self._refresh_credentials()
+            logging.warning("Auth failure %s, re-initializing credentials", resp.status_code)
+            self.init_credentials()
             resp = self.signed_request(
                 method=method,
                 url=url,

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "prometrix"
-version = "0.2.6"
+version = "0.2.7"
 authors = ["Avi Kotlicky <[email protected]>"]
 readme = "README.md"
 packages = [{include = "prometrix"}]