diff --git a/robusta_krr/core/integrations/prometheus/prometheus_utils.py b/robusta_krr/core/integrations/prometheus/prometheus_utils.py index e561af24..ec6f20e2 100644 --- a/robusta_krr/core/integrations/prometheus/prometheus_utils.py +++ b/robusta_krr/core/integrations/prometheus/prometheus_utils.py @@ -3,7 +3,12 @@ from typing import TYPE_CHECKING import boto3 -from prometrix import AWSPrometheusConfig, CoralogixPrometheusConfig, PrometheusConfig, VictoriaMetricsPrometheusConfig +from prometrix import ( + AWSPrometheusConfig, + CoralogixPrometheusConfig, + PrometheusConfig, + VictoriaMetricsPrometheusConfig, +) from robusta_krr.core.models.config import settings @@ -40,6 +45,7 @@ def generate_prometheus_config( region = settings.eks_managed_prom_region if settings.eks_managed_prom_region else session.region_name access_key = settings.eks_access_key if settings.eks_access_key else credentials.access_key secret_key = settings.eks_secret_key.get_secret_value() if settings.eks_secret_key else credentials.secret_key + token = settings.eks_token.get_secret_value() if settings.eks_token else credentials.token service_name = settings.eks_service_name if settings.eks_secret_key else "aps" if not region: raise Exception("No eks region specified") @@ -47,6 +53,7 @@ def generate_prometheus_config( return AWSPrometheusConfig( access_key=access_key, secret_access_key=secret_key, + token=token, aws_region=region, service_name=service_name, **baseconfig, diff --git a/robusta_krr/core/models/config.py b/robusta_krr/core/models/config.py index 32241ed1..5735387f 100644 --- a/robusta_krr/core/models/config.py +++ b/robusta_krr/core/models/config.py @@ -44,6 +44,7 @@ class Config(pd.BaseSettings): eks_managed_prom_profile_name: Optional[str] = pd.Field(None) eks_access_key: Optional[str] = pd.Field(None) eks_secret_key: Optional[pd.SecretStr] = pd.Field(None) + eks_token: Optional[pd.SecretStr] = pd.Field(None) eks_service_name: Optional[str] = pd.Field(None) eks_managed_prom_region: Optional[str] = pd.Field(None) coralogix_token: Optional[pd.SecretStr] = pd.Field(None) diff --git a/robusta_krr/main.py b/robusta_krr/main.py index 7159bdde..d88e09b1 100644 --- a/robusta_krr/main.py +++ b/robusta_krr/main.py @@ -171,6 +171,12 @@ def run_strategy( help="Sets the secret key for eks prometheus connection.", rich_help_panel="Prometheus EKS Settings", ), + eks_token: Optional[str] = typer.Option( + None, + "--eks-token", + help="Sets the session token for eks prometheus connection.", + rich_help_panel="Prometheus EKS Settings", + ), eks_service_name: Optional[str] = typer.Option( "aps", "--eks-service-name", @@ -234,13 +240,24 @@ def run_strategy( rich_help_panel="Output Settings", ), verbose: bool = typer.Option( - False, "--verbose", "-v", help="Enable verbose mode", rich_help_panel="Logging Settings" + False, + "--verbose", + "-v", + help="Enable verbose mode", + rich_help_panel="Logging Settings", ), quiet: bool = typer.Option( - False, "--quiet", "-q", help="Enable quiet mode", rich_help_panel="Logging Settings" + False, + "--quiet", + "-q", + help="Enable quiet mode", + rich_help_panel="Logging Settings", ), log_to_stderr: bool = typer.Option( - False, "--logtostderr", help="Pass logs to stderr", rich_help_panel="Logging Settings" + False, + "--logtostderr", + help="Pass logs to stderr", + rich_help_panel="Logging Settings", ), width: Optional[int] = typer.Option( None, @@ -270,7 +287,10 @@ def run_strategy( ) -> None: f"""Run KRR using the `{_strategy_name}` strategy""" if not show_severity and format != "csv": - raise click.BadOptionUsage("--exclude-severity", "--exclude-severity works only with format=csv") + raise click.BadOptionUsage( + "--exclude-severity", + "--exclude-severity works only with format=csv", + ) try: config = Config( @@ -292,6 +312,7 @@ def run_strategy( eks_managed_prom_profile_name=eks_managed_prom_profile_name, eks_access_key=eks_access_key, eks_secret_key=eks_secret_key, + eks_token=eks_token, eks_service_name=eks_service_name, coralogix_token=coralogix_token, openshift=openshift, @@ -329,7 +350,14 @@ def run_strategy( kind=inspect.Parameter.KEYWORD_ONLY, default=OptionInfo( default=field_meta.default, - param_decls=list(set([f"--{field_name}", f"--{field_name.replace('_', '-')}"])), + param_decls=list( + set( + [ + f"--{field_name}", + f"--{field_name.replace('_', '-')}", + ] + ) + ), help=f"{field_meta.field_info.description}", rich_help_panel="Strategy Settings", ), diff --git a/tests/formatters/test_csv_formatter.py b/tests/formatters/test_csv_formatter.py index 150b2aa5..98cffbcf 100644 --- a/tests/formatters/test_csv_formatter.py +++ b/tests/formatters/test_csv_formatter.py @@ -119,6 +119,7 @@ "eks_managed_prom_profile_name": null, "eks_access_key": null, "eks_secret_key": null, + "eks_token": null, "eks_service_name": "aps", "eks_managed_prom_region": null, "coralogix_token": null, @@ -235,7 +236,7 @@ def test_csv_headers(override_config: dict[str, Any], expected_headers: list[str "Old Pods": "1", "Type": "Deployment", "Container": "mock-container-1", - 'Severity': 'CRITICAL', + "Severity": "CRITICAL", "CPU Diff": "-87m", "CPU Requests": "(-43m) 50m -> 6m", "CPU Limits": "2.0 -> ?", @@ -271,7 +272,7 @@ def test_csv_headers(override_config: dict[str, Any], expected_headers: list[str "Old Pods": "1", "Type": "Deployment", "Container": "mock-container-1", - 'Severity': 'CRITICAL', + "Severity": "CRITICAL", "CPU Diff": "-87m", "CPU Requests": "(-43m) 50m -> 6m", "CPU Limits": "2.0 -> ?",