ROB-1126: headers support for centralized prometheus

.gitignore

@@ -21,3 +21,5 @@ skaffold.dev.yaml
 tests/last_used_tag.txt
 pytest.ini
 gen-config-test
+
+**/.claude/settings.local.json

docs/configuration/alertmanager-integration/_pull_integration.rst

@@ -24,6 +24,12 @@ If Robusta fails to auto-detect the Prometheus and Alertmanager urls - and you s
         # prometheus_auth: ...
         # alertmanager_auth: ...
 
+        # If using a multi-tenant prometheus or alertmanager, pass the org id to all queries
+        # prometheus_additional_headers:
+        #   X-Scope-OrgID: <org id>
+        # alertmanager_additional_headers:
+        #   X-Scope-OrgID: <org id>
+
 .. code-annotations::
     1. Example: http://alertmanager-Helm_release_name-kube-prometheus-alertmanager.default.svc.cluster.local:9093.
     2. Example: http://Helm_Release_Name-kube-prometheus-prometheus.default.svc.cluster.local:9090

docs/configuration/alertmanager-integration/coralogix_managed_prometheus.rst

@@ -82,6 +82,12 @@ Metrics querying lets Robusta pull metrics from Coralogix Managed Prometheus.
       # grafana_api_key: <YOUR GRAFANA EDITOR API KEY> # (1)
       # alertmanager_flavor: grafana
 
+        # If using a multi-tenant prometheus or alertmanager, pass the org id to all queries
+        # prometheus_additional_headers:
+        #   X-Scope-OrgID: <org id>
+        # alertmanager_additional_headers:
+        #   X-Scope-OrgID: <org id>
+
 .. code-annotations::
     1. This is necessary for Robusta to create silences when using Grafana Alerts, because of minor API differences in the AlertManager embedded in Grafana.
 

docs/configuration/alertmanager-integration/grafana-alert-manager.rst

@@ -83,6 +83,12 @@ Modify and add the following config to ``generated_values.yaml`` and :ref:`updat
         # prometheus_additional_labels:
         #   cluster: 'CLUSTER_NAME_HERE'
 
+        # If using a multi-tenant prometheus or alertmanager, pass the org id to all queries
+        # prometheus_additional_headers:
+        #   X-Scope-OrgID: <org id>
+        # alertmanager_additional_headers:
+        #   X-Scope-OrgID: <org id>
+
 .. code-annotations::
     1. This is necessary for Robusta to create silences when using Grafana Alerts, because of minor API differences in the AlertManager embedded in Grafana.
 

docs/configuration/alertmanager-integration/victoria-metrics.rst

@@ -39,6 +39,12 @@ Add the following to ``generated_values.yaml`` and :ref:`update Robusta <Simple
         # grafana_api_key: <YOUR GRAFANA EDITOR API KEY> # (1)
         # alertmanager_flavor: grafana
 
+        # If using a multi-tenant prometheus or alertmanager, pass the org id to all queries
+        # prometheus_additional_headers:
+        #   X-Scope-OrgID: <org id>
+        # alertmanager_additional_headers:
+        #   X-Scope-OrgID: <org id>
+
 .. code-annotations::
     1. This is necessary for Robusta to create silences when using Grafana Alerts, because of minor API differences in the AlertManager embedded in Grafana.
 

playbooks/robusta_playbooks/krr.py

@@ -178,11 +178,14 @@ class KRRSecret(BaseModel):
     secret_value: str
     command_flag: str
 
-    def __init__(self, env_var_name: str, secret_key: str, secret_value: str, command_flag: str):
-        secret_b64_str = base64.b64encode(bytes(secret_value, "utf-8")).decode("utf-8")
-        super().__init__(
-            env_var_name=env_var_name, secret_key=secret_key, secret_value=secret_b64_str, command_flag=command_flag
-        )
+    @validator('secret_value', pre=True, always=True, allow_reuse=True)
+    def encode_secret_value(cls, v: str) -> str:
+        if isinstance(v, str):
+            return base64.b64encode(bytes(v, "utf-8")).decode("utf-8")
+        raise ValueError("secret_value must be a string")
+
+class KRRSecretKeyValuePair(KRRSecret):
+    arg_key: str
 
 
 def _generate_krr_job_secret(scan_id: str, krr_secrets: Optional[List[KRRSecret]]) -> Optional[JobSecret]:
@@ -210,7 +213,14 @@ def _generate_krr_env_vars(
 def _generate_additional_env_args(krr_secrets: Optional[List[KRRSecret]]) -> str:
     if not krr_secrets:
         return ""
-    return " ".join(f"{secret.command_flag} '$({secret.env_var_name})'" for secret in krr_secrets)
+    additional_args = []
+    for secret in krr_secrets:
+        if isinstance(secret, KRRSecretKeyValuePair):
+            additional_args.append(f"{secret.command_flag} '{secret.arg_key}:$({secret.env_var_name})'")
+        else:
+            additional_args.append(f"{secret.command_flag} '$({secret.env_var_name})'")
+
+    return " ".join(additional_args)
 
 
 def _generate_cmd_line_args(prom_config: PrometheusConfig) -> str:
@@ -240,6 +250,22 @@ def _generate_prometheus_secrets(prom_config: PrometheusConfig) -> List[KRRSecre
                 command_flag="--prometheus-auth-header",
             )
         )
+
+    if prom_config.headers:
+        for header_name, header_value in prom_config.headers.items():
+
+            env_var_name = f"PROMETHEUS_HEADER_{header_name.upper().replace('-', '_')}"
+            secret_key = f"prometheus-header-{header_name.lower()}"
+
+            krr_secrets.append(
+                KRRSecretKeyValuePair(
+                    env_var_name=env_var_name,
+                    arg_key=header_name,
+                    secret_key=secret_key,
+                    secret_value=header_value,
+                    command_flag="--prometheus-headers",
+                )
+            )
 
     if isinstance(prom_config, AWSPrometheusConfig):
         krr_secrets.extend(

src/robusta/core/model/base_params.py

@@ -310,6 +310,7 @@ class PrometheusParams(ActionParams):
     :var prometheus_url: Prometheus url. If omitted, we will try to find a prometheus instance in the same cluster
     :var prometheus_auth: Prometheus auth header to be used in Authorization header. If omitted, we will not add any auth header
     :var prometheus_url_query_string: Additional query string parameters to be appended to the Prometheus connection URL
+    :var prometheus_additional_headers: additional HTTP headers (if defined) to add to every prometheus query
     :var prometheus_additional_labels: A dictionary of additional labels needed for multi-cluster prometheus
     :var add_additional_labels: adds the additional labels (if defined) to the query
 
@@ -325,6 +326,7 @@ class PrometheusParams(ActionParams):
     prometheus_url: Optional[str] = None
     prometheus_auth: Optional[SecretStr] = None
     prometheus_url_query_string: Optional[str] = None
+    prometheus_additional_headers: Optional[Dict[str, str]] = None
     prometheus_additional_labels: Optional[Dict[str, str]] = None
     add_additional_labels: bool = True
     prometheus_graphs_overrides: Optional[List[OverrideGraph]] = None

src/robusta/integrations/prometheus/utils.py

@@ -48,10 +48,13 @@ def generate_prometheus_config(prometheus_params: PrometheusParams) -> Prometheu
         "url": url,
         "disable_ssl": not PROMETHEUS_SSL_ENABLED,
         "additional_labels": prometheus_params.prometheus_additional_labels,
-        "prometheus_url_query_string": prometheus_params.prometheus_url_query_string,
+        "prometheus_url_query_string": prometheus_params.prometheus_url_query_string
     }
     if prometheus_params.prometheus_auth:
         baseconfig["prometheus_auth"] = prometheus_params.prometheus_auth.get_secret_value()
+
+    if prometheus_params.prometheus_additional_headers:
+        baseconfig["headers"] = prometheus_params.prometheus_additional_headers
 
     # aws config
     if AWS_ACCESS_KEY:

src/robusta/utils/silence_utils.py

@@ -56,11 +56,13 @@ def to_list(self) -> List[str]:
 class AlertManagerParams(ActionParams):
     """
     :var alertmanager_url: Alternative Alert Manager url to send requests.
+    :var alertmanager_additional_headers: additional HTTP headers (if defined) to add to every alertmanager request
     """
 
     alertmanager_flavor: str = None  # type: ignore
     alertmanager_url: Optional[str]
     alertmanager_auth: Optional[SecretStr] = None
+    alertmanager_additional_headers: Optional[Dict[str, str]] = None
     grafana_api_key: str = None  # type: ignore
 
     @validator("alertmanager_url", allow_reuse=True)
@@ -143,6 +145,9 @@ def gen_alertmanager_headers(params: AlertManagerParams) -> Dict:
 
     elif params.alertmanager_auth:
         headers.update({"Authorization": params.alertmanager_auth.get_secret_value()})
+
+    if params.alertmanager_additional_headers:
+        headers.update(params.alertmanager_additional_headers)
 
     return headers