diff --git a/prometrix/connect/aws_connect.py b/prometrix/connect/aws_connect.py index 18839d0..951eeef 100644 --- a/prometrix/connect/aws_connect.py +++ b/prometrix/connect/aws_connect.py @@ -15,7 +15,6 @@ SA_TOKEN_PATH = os.environ.get("SA_TOKEN_PATH", "/var/run/secrets/eks.amazonaws.com/serviceaccount/token") AWS_ASSUME_ROLE = os.environ.get("AWS_ASSUME_ROLE") -AWS_REFRESH_CREDS_SEC = int(os.environ.get("AWS_REFRESH_CREDS_SEC", "900")) # 15 minutes class AWSPrometheusConnect(CustomPrometheusConnect): def __init__( @@ -32,43 +31,20 @@ def __init__( self.region = region self.service_name = service_name - 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 + if access_key and secret_key: + # Backwards compatibility: use static keys + self._credentials = Credentials(access_key, secret_key, token) + 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 role_to_assume = assume_role_arn or AWS_ASSUME_ROLE - self._role_to_assume = 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 + if role_to_assume: + self._assume_role(role_to_assume) def _assume_role(self, role_arn: str) -> None: try: @@ -99,12 +75,6 @@ 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() >= AWS_REFRESH_CREDS_SEC: - logging.debug("%d seconds passed; re-initializing AWS credentials", AWS_REFRESH_CREDS_SEC) - 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) @@ -123,28 +93,6 @@ def signed_request( params=params, ) - def _request_with_refresh(self, *, method, url, data=None, params=None, headers=None, verify=False): - resp = self.signed_request( - method=method, - url=url, - data=data, - params=params, - verify=verify, - headers=headers, - ) - if resp is not None and resp.status_code in (400, 401, 403): - logging.warning("Auth failure %s, re-initializing credentials", resp.status_code) - self.init_credentials() - resp = self.signed_request( - method=method, - url=url, - data=data, - params=params, - verify=verify, - headers=headers, - ) - return resp - def _custom_query(self, query: str, params: dict = None): """ Send a custom query to a Prometheus Host. @@ -165,7 +113,7 @@ def _custom_query(self, query: str, params: dict = None): data = None query = str(query) # using the query API to get raw data - response = self._request_with_refresh( + response = self.signed_request( method="POST", url="{0}/api/v1/query".format(self.url), data={**{"query": query}, **params}, @@ -204,7 +152,7 @@ def safe_custom_query_range( params = params or {} query = str(query) - response = self._request_with_refresh( + response = self.signed_request( method="POST", url="{0}/api/v1/query_range".format(self.url), data={ diff --git a/pyproject.toml b/pyproject.toml index 7bc9882..a4183f4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "prometrix" -version = "0.2.7" +version = "0.2.8" authors = ["Avi Kotlicky "] readme = "README.md" packages = [{include = "prometrix"}]