From 519c71ec56a0d3edfd070174492a6d4bbc97d01c Mon Sep 17 00:00:00 2001
From: LeaveMyYard <33721692+LeaveMyYard@users.noreply.github.com>
Date: Wed, 12 Jul 2023 10:54:37 +0300
Subject: Add ability to pass custom prometheus headers

---
 .../prometheus/metrics_service/prometheus_metrics_service.py      | 4 ++--
 robusta_krr/core/models/config.py                                 | 8 ++++++++
 robusta_krr/main.py                                               | 8 ++++++++
 3 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/robusta_krr/core/integrations/prometheus/metrics_service/prometheus_metrics_service.py b/robusta_krr/core/integrations/prometheus/metrics_service/prometheus_metrics_service.py
index 49daf87..ecc3e5d 100644
--- a/robusta_krr/core/integrations/prometheus/metrics_service/prometheus_metrics_service.py
+++ b/robusta_krr/core/integrations/prometheus/metrics_service/prometheus_metrics_service.py
@@ -83,10 +83,10 @@ class PrometheusMetricsService(MetricsService):
 
         self.info(f"Using {self.name()} at {self.url} for cluster {cluster or 'default'}")
 
-        headers = {}
+        headers = self.config.prometheus_other_headers
 
         if self.auth_header:
-            headers = {"Authorization": self.auth_header}
+            headers |= {"Authorization": self.auth_header}
         elif not self.config.inside_cluster and self.api_client is not None:
             self.api_client.update_params_for_auth(headers, {}, ["BearerToken"])
 
diff --git a/robusta_krr/core/models/config.py b/robusta_krr/core/models/config.py
index 38e3f4f..f64629d 100644
--- a/robusta_krr/core/models/config.py
+++ b/robusta_krr/core/models/config.py
@@ -24,6 +24,7 @@ class Config(pd.BaseSettings):
     # Prometheus Settings
     prometheus_url: Optional[str] = pd.Field(None)
     prometheus_auth_header: Optional[str] = pd.Field(None)
+    prometheus_other_headers: dict[str, str] = pd.Field(default_factory=dict)
     prometheus_ssl_enabled: bool = pd.Field(False)
     prometheus_cluster_label: Optional[str] = pd.Field(None)
     prometheus_label: str = pd.Field("cluster")
@@ -50,6 +51,13 @@ class Config(pd.BaseSettings):
     def Formatter(self) -> formatters.FormatterFunc:
         return formatters.find(self.format)
 
+    @pd.validator("prometheus_other_headers", pre=True)
+    def validate_prometheus_other_headers(cls, headers: Union[list[str], dict[str, str]]) -> dict[str, str]:
+        if isinstance(headers, dict):
+            return headers
+
+        return {k.strip().lower(): v.strip() for k, v in [header.split(":") for header in headers]}
+
     @pd.validator("namespaces")
     def validate_namespaces(cls, v: Union[list[str], Literal["*"]]) -> Union[list[str], Literal["*"]]:
         if v == []:
diff --git a/robusta_krr/main.py b/robusta_krr/main.py
index 185a819..d6fdf1f 100644
--- a/robusta_krr/main.py
+++ b/robusta_krr/main.py
@@ -85,6 +85,13 @@ def load_commands() -> None:
                     help="Prometheus authentication header.",
                     rich_help_panel="Prometheus Settings",
                 ),
+                prometheus_other_headers: Optional[List[str]] = typer.Option(
+                    None,
+                    "--prometheus-headers",
+                    "-H",
+                    help="Additional headers to add to Prometheus requests. Format as 'key: value', for example 'X-MyHeader: 123'. Trailing whitespaces will be stripped.",
+                    rich_help_panel="Prometheus Settings",
+                ),
                 prometheus_ssl_enabled: bool = typer.Option(
                     False,
                     "--prometheus-ssl-enabled",
@@ -125,6 +132,7 @@ def load_commands() -> None:
                     namespaces="*" if "*" in namespaces else namespaces,
                     prometheus_url=prometheus_url,
                     prometheus_auth_header=prometheus_auth_header,
+                    prometheus_other_headers=prometheus_other_headers,
                     prometheus_ssl_enabled=prometheus_ssl_enabled,
                     prometheus_cluster_label=prometheus_cluster_label,
                     prometheus_label=prometheus_label,
-- 
cgit v1.2.3