Add session token to the AWS credentials

Author: dex4er

prometrix/connect/aws_connect.py

@@ -12,10 +12,10 @@
 
 class AWSPrometheusConnect(CustomPrometheusConnect):
     def __init__(
-        self, access_key: str, secret_key: str, region: str, service_name: str, **kwargs
+        self, access_key: str, secret_key: str, region: str, service_name: str, token: Optional[str] = None, **kwargs
     ):
         super().__init__(**kwargs)
-        self._credentials = Credentials(access_key, secret_key)
+        self._credentials = Credentials(access_key, secret_key, token)
         self._sigv4auth = S3SigV4Auth(self._credentials, service_name, region)
 
     def signed_request(

prometrix/models/prometheus_config.py

@@ -32,6 +32,7 @@ class PrometheusConfig(BaseModel):
 class AWSPrometheusConfig(PrometheusConfig):
     access_key: str
     secret_access_key: str
+    token: Optional[str] = None
     service_name: str = "aps"
     aws_region: str
     supported_apis: List[PrometheusApis] = [

prometrix/utils.py

@@ -37,6 +37,7 @@ def get_custom_prometheus_connect(
             service_name=prom_config.service_name,
             region=prom_config.aws_region,
             config=prom_config,
+            token=prom_config.token,
         )
     else:
         prom = CustomPrometheusConnect(config=prom_config)