added irsa support to prometrix

Author: Avi-Robusta

prometrix/connect/aws_connect.py

@@ -2,6 +2,7 @@
 from typing import Any, Dict, Optional, List
 
 import requests
+import boto3
 from botocore.auth import S3SigV4Auth
 from botocore.awsrequest import AWSRequest
 from botocore.credentials import Credentials
@@ -12,19 +13,44 @@
 
 class AWSPrometheusConnect(CustomPrometheusConnect):
     def __init__(
-        self, access_key: str, secret_key: str, region: str, service_name: str, token: Optional[str] = None, **kwargs
+        self,
+        access_key: str,
+        secret_key: str,
+        region: str,
+        service_name: str,
+        token: Optional[str] = None,
+        **kwargs,
     ):
+        """
+        AMP connector using S3SigV4Auth.
+        Supports both static AWS keys (for backwards compatibility) and IRSA.
+        """
         super().__init__(**kwargs)
-        self._credentials = Credentials(access_key, secret_key, token)
-        self._sigv4auth = S3SigV4Auth(self._credentials, service_name, region)
+        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)
+        else:
+            # Preferred: IRSA or any AWS provider chain
+            session = boto3.Session()
+            creds = session.get_credentials()
+            if not creds:
+                raise RuntimeError("No AWS credentials found (neither static keys nor IRSA)")
+            self._credentials = creds
+
+    def _build_auth(self) -> S3SigV4Auth:
+        """Builds fresh SigV4 auth with current credentials (handles rotation)."""
+        frozen = self._credentials.get_frozen_credentials()
+        return S3SigV4Auth(frozen, self.service_name, self.region)
 
     def signed_request(
         self, method, url, data=None, params=None, verify=False, headers=None
     ):
-        request = AWSRequest(
-            method=method, url=url, data=data, params=params, headers=headers
-        )
-        self._sigv4auth.add_auth(request)
+        request = AWSRequest(method=method, url=url, data=data, params=params, headers=headers)
+        auth = self._build_auth()
+        auth.add_auth(request)
         return requests.request(
             method=method,
             url=url,

prometrix/models/prometheus_config.py

@@ -30,8 +30,8 @@ class PrometheusConfig(BaseModel):
 
 
 class AWSPrometheusConfig(PrometheusConfig):
-    access_key: str
-    secret_access_key: str
+    access_key: Optional[str] = None
+    secret_access_key: Optional[str] = None
     token: Optional[str] = None
     service_name: str = "aps"
     aws_region: str

pyproject.toml

@@ -1,12 +1,12 @@
 [tool.poetry]
 name = "prometrix"
-version = "0.2.1"
+version = "0.2.2"
 authors = ["Avi Kotlicky <avi@robusta.dev>"]
 readme = "README.md"
 packages = [{include = "prometrix"}]
 description = "A Python Prometheus client for all Prometheus instances."
 
-[project.urls]
+[tool.poetry.urls]
 "Homepage" = "https://github.com/robusta-dev/prometrix"
 "Bug Tracker" = "https://github.com/robusta-dev/prometrix/issues"