ROB-1291 Prometheus tenant from the namespace label (option B: UI + alert graphs)

playbooks/robusta_playbooks/alerts_integration.py

@@ -45,6 +45,7 @@
     get_node_internal_ip,
 )
 from robusta.core.playbooks.oom_killer_utils import logs_enricher, start_log_enrichment
+from robusta.core.playbooks.prometheus_enrichment_utils import with_namespace_tenant
 from robusta.core.reporting import FindingSubject
 from robusta.core.reporting.base import Link, LinkType
 from robusta.core.reporting.blocks import TableBlockFormat
@@ -305,7 +306,8 @@ def graph_enricher(alert: PrometheusKubernetesAlert, params: GraphEnricherParams
     """
     promql_query = alert.get_prometheus_query()
     chart, prom_block = create_chart_from_prometheus_query(
-        params,
+        # The alert fired about its namespace, so its expression is graphed from that namespace's tenant
+        with_namespace_tenant(params, alert.alert.labels.get("namespace")),
         promql_query,
         alert.alert.startsAt,
         include_x_axis=False,

src/robusta/core/model/base_params.py

@@ -321,6 +321,7 @@ class PrometheusParams(ActionParams):
     :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
+    :var namespace: the namespace the queries are about. With PROMETHEUS_TENANT_LABEL set, its label picks the tenant
 
     :example prometheus_url: "http://prometheus-k8s.monitoring.svc.cluster.local:9090"
     :example prometheus_auth: Basic YWRtaW46cGFzc3dvcmQ=
@@ -337,6 +338,7 @@ class PrometheusParams(ActionParams):
     prometheus_additional_headers: Optional[Dict[str, str]] = None
     prometheus_additional_labels: Optional[Dict[str, str]] = None
     add_additional_labels: bool = True
+    namespace: Optional[str] = None
     prometheus_graphs_overrides: Optional[List[OverrideGraph]] = None
 
     @validator("prometheus_url", allow_reuse=True)

src/robusta/core/model/env_vars.py

@@ -49,6 +49,10 @@ def load_bool(env_var, default: bool):
 PROMETHEUS_ENABLED = os.environ.get("PROMETHEUS_ENABLED", "false").lower() == "true"
 MANAGED_CONFIGURATION_ENABLED = os.environ.get("MANAGED_CONFIGURATION_ENABLED", "false").lower() == "true"
 PROMETHEUS_SSL_ENABLED = os.environ.get("PROMETHEUS_SSL_ENABLED", "false").lower() == "true"
+# Multi-tenant Prometheus (e.g. Mimir): label on the action's namespace, sent as the tenant header.
+# Env vars, not global_config, so a request from the UI can't override them.
+PROMETHEUS_TENANT_LABEL = os.environ.get("PROMETHEUS_TENANT_LABEL")
+PROMETHEUS_TENANT_HEADER = os.environ.get("PROMETHEUS_TENANT_HEADER", "X-Scope-OrgID")
 
 INCOMING_REQUEST_TIME_WINDOW_SECONDS = int(os.environ.get("INCOMING_REQUEST_TIME_WINDOW_SECONDS", 3600))
 

src/robusta/core/playbooks/prometheus_enrichment_utils.py

@@ -16,7 +16,11 @@
     ResourceChartItemType,
     ResourceChartResourceType,
 )
-from robusta.core.model.env_vars import FLOAT_PRECISION_LIMIT, PROMETHEUS_REQUEST_TIMEOUT_SECONDS
+from robusta.core.model.env_vars import (
+    FLOAT_PRECISION_LIMIT,
+    PROMETHEUS_REQUEST_TIMEOUT_SECONDS,
+    PROMETHEUS_TENANT_LABEL,
+)
 from robusta.core.reporting.blocks import GraphBlock, PrometheusBlock, PrometheusBlockLineData
 from robusta.core.reporting.charts import XYChart
 from robusta.core.reporting.custom_rendering import charts_style
@@ -61,6 +65,16 @@ def __prepare_promql_query(provided_labels: Dict[Any, Any], promql_query_templat
     return promql_query
 
 
+def with_namespace_tenant(prometheus_params: PrometheusParams, namespace: Optional[str]) -> PrometheusParams:
+    """
+    Multi-tenant Prometheus: params whose queries go to the tenant of `namespace`, unless they already name one.
+    Returns a copy, since an action reuses its params for several graphs.
+    """
+    if not PROMETHEUS_TENANT_LABEL or prometheus_params.namespace or not namespace:
+        return prometheus_params
+    return prometheus_params.copy(update={"namespace": namespace})
+
+
 def get_node_internal_ip(node: Node) -> str:
     internal_ip = next(addr.address for addr in node.status.addresses if addr.type == "InternalIP")
     return internal_ip
@@ -495,6 +509,10 @@ def create_graph_enrichment(
     hide_legends: Optional[bool] = False,
     metrics_legends_labels: Optional[List[str]] = None,
 ) -> GraphBlock:
+    # Only a query that fills in $namespace is about that namespace. Node queries aren't, even when the labels have
+    # one (e.g. the exporter's namespace on a node alert), so they keep the default tenant
+    if "namespace" in Template(promql_query).get_identifiers():
+        prometheus_params = with_namespace_tenant(prometheus_params, labels.get("namespace"))
     promql_query = __prepare_promql_query(labels, promql_query)
     chart, prom_block = create_chart_from_prometheus_query(
         prometheus_params,

src/robusta/integrations/prometheus/utils.py

@@ -19,8 +19,11 @@
     DISABLE_PROMETHEUS_DISCOVERY,
     HOLMES_DISCOVERY_NAMESPACE,
     PROMETHEUS_SSL_ENABLED,
+    PROMETHEUS_TENANT_HEADER,
+    PROMETHEUS_TENANT_LABEL,
     SERVICE_CACHE_TTL_SEC,
 )
+from robusta.integrations.kubernetes.api_client_utils import get_namespace_labels
 from robusta.utils.service_discovery import find_service_url
 
 AZURE_RESOURCE = os.environ.get("AZURE_RESOURCE", "https://prometheus.monitor.azure.com")
@@ -39,6 +42,21 @@
 VICTORIA_METRICS_CONFIGURED = os.environ.get("VICTORIA_METRICS_CONFIGURED", "false").lower() == "true"
 
 
+def get_namespace_tenant(namespace: str) -> Optional[str]:
+    """
+    Multi-tenant Prometheus (e.g. Mimir): the tenant of a namespace, which is the value of its PROMETHEUS_TENANT_LABEL
+    label, read from the cached namespace data (NAMESPACE_DATA_TTL). None when the namespace has no such label.
+    Call it only with PROMETHEUS_TENANT_LABEL set.
+    """
+    try:
+        tenant = (get_namespace_labels(namespace) or {}).get(PROMETHEUS_TENANT_LABEL)
+    except KeyError:  # a namespace created after the last reload, or deleted
+        tenant = None
+    if not tenant:
+        logging.debug(f"No '{PROMETHEUS_TENANT_LABEL}' label on namespace {namespace}, using the default tenant")
+    return tenant
+
+
 def generate_prometheus_config(prometheus_params: PrometheusParams) -> PrometheusConfig:
     is_victoria_metrics = VICTORIA_METRICS_CONFIGURED
     url: Optional[str] = (
@@ -63,6 +81,12 @@ def generate_prometheus_config(prometheus_params: PrometheusParams) -> Prometheu
     if prometheus_params.prometheus_additional_headers:
         baseconfig["headers"] = prometheus_params.prometheus_additional_headers
 
+    if PROMETHEUS_TENANT_LABEL and prometheus_params.namespace:
+        tenant = get_namespace_tenant(prometheus_params.namespace)
+        if tenant:
+            # A new dict that replaces any tenant header from global_config or the request, leaving theirs unchanged
+            baseconfig["headers"] = {**baseconfig.get("headers", {}), PROMETHEUS_TENANT_HEADER: tenant}
+
     # aws config
     if AWS_REGION:
         return AWSPrometheusConfig(

tests/test_prometheus_tenant.py

@@ -0,0 +1,114 @@
+from datetime import datetime
+from unittest.mock import patch
+
+import pytest
+
+from robusta.core.model.base_params import PrometheusParams
+from robusta.core.playbooks import prometheus_enrichment_utils
+from robusta.integrations.prometheus import utils
+from robusta.integrations.prometheus.utils import generate_prometheus_config
+
+NAMESPACES = {"shop": {"team": "corporate"}, "billing": {"team": "retail"}, "unlabeled": {"app": "web"}}
+
+
+@pytest.fixture
+def tenant_label():
+    def labels(namespace):
+        return NAMESPACES[namespace]
+
+    with patch.object(utils, "PROMETHEUS_TENANT_LABEL", "team"), patch.object(
+        utils, "get_namespace_labels", side_effect=labels
+    ) as get_labels:
+        yield get_labels
+
+
+def _headers(**params):
+    return generate_prometheus_config(PrometheusParams(prometheus_url="http://prometheus:9090", **params)).headers
+
+
+def test_tenant_from_namespace_label(tenant_label):
+    assert _headers(namespace="shop") == {"X-Scope-OrgID": "corporate"}
+
+
+@pytest.mark.parametrize("label_key", ["product.productline", "example.com/productline"])
+def test_tenant_label_key_with_dots_or_prefix(label_key):
+    # The label key is used as written: dots and a DNS prefix are just part of the key
+    with patch.object(utils, "PROMETHEUS_TENANT_LABEL", label_key), patch.object(
+        utils, "get_namespace_labels", return_value={label_key: "corporate", "productline": "other"}
+    ):
+        assert _headers(namespace="shop") == {"X-Scope-OrgID": "corporate"}
+
+
+def test_tenant_replaces_configured_one_and_keeps_other_headers(tenant_label):
+    headers = _headers(namespace="billing", prometheus_additional_headers={"X-Scope-OrgID": "default", "X-Other": "1"})
+    assert headers == {"X-Scope-OrgID": "retail", "X-Other": "1"}
+
+
+def test_no_namespace_keeps_default_tenant(tenant_label):
+    assert _headers(prometheus_additional_headers={"X-Scope-OrgID": "default"}) == {"X-Scope-OrgID": "default"}
+    tenant_label.assert_not_called()
+
+
+@pytest.mark.parametrize("namespace", ["unlabeled", "missing"])
+def test_unlabeled_or_unknown_namespace_keeps_default_tenant(tenant_label, namespace):
+    headers = _headers(namespace=namespace, prometheus_additional_headers={"X-Scope-OrgID": "default"})
+    assert headers == {"X-Scope-OrgID": "default"}
+
+
+def test_feature_off_reads_nothing(tenant_label):
+    with patch.object(utils, "PROMETHEUS_TENANT_LABEL", None):
+        assert not _headers(namespace="shop")
+    tenant_label.assert_not_called()
+
+
+POD_QUERY = 'container_memory_working_set_bytes{namespace="$namespace", pod=~"$pod"}'
+NODE_QUERY = 'node_memory_MemAvailable_bytes{instance=~"$node_internal_ip:.*"}'
+
+
+@pytest.mark.parametrize(
+    "tenant_label_env, params_namespace, query, labels, expected",
+    [
+        ("team", None, POD_QUERY, {"namespace": "shop", "pod": "web"}, "shop"),  # pod graph: its namespace
+        ("team", None, 'up{namespace="${namespace}"}', {"namespace": "shop"}, "shop"),  # braced placeholder
+        ("team", None, 'up{label="$$namespace"}', {"namespace": "shop"}, None),  # escaped: a literal "$namespace"
+        ("team", "billing", POD_QUERY, {"namespace": "shop"}, "billing"),  # an explicit namespace wins
+        # node graph on a node alert that carries the exporter's namespace: not about it, default tenant
+        ("team", None, NODE_QUERY, {"namespace": "monitoring", "node_internal_ip": "10.0.0.1"}, None),
+        (None, None, POD_QUERY, {"namespace": "shop"}, None),  # feature off: params untouched
+    ],
+)
+def test_graph_enrichment_namespace(tenant_label_env, params_namespace, query, labels, expected):
+    params = PrometheusParams(namespace=params_namespace)
+    with patch.object(prometheus_enrichment_utils, "PROMETHEUS_TENANT_LABEL", tenant_label_env), patch.object(
+        prometheus_enrichment_utils, "create_chart_from_prometheus_query"
+    ) as create_chart:
+        create_chart.side_effect = Exception("stop")
+        with pytest.raises(Exception, match="stop"):
+            prometheus_enrichment_utils.create_graph_enrichment(
+                datetime.now(), labels, query, params, 60, None, None
+            )
+    assert create_chart.call_args.args[0].namespace == expected
+    assert params.namespace == params_namespace
+
+
+@pytest.mark.parametrize(
+    "tenant_label_env, params_namespace, namespace, expected",
+    [
+        ("team", None, "shop", "shop"),  # e.g. graph_enricher with the alert's namespace
+        ("team", "billing", "shop", "billing"),
+        ("team", None, None, None),
+        (None, None, "shop", None),
+    ],
+)
+def test_with_namespace_tenant(tenant_label_env, params_namespace, namespace, expected):
+    params = PrometheusParams(namespace=params_namespace)
+    with patch.object(prometheus_enrichment_utils, "PROMETHEUS_TENANT_LABEL", tenant_label_env):
+        result = prometheus_enrichment_utils.with_namespace_tenant(params, namespace)
+    assert result.namespace == expected
+    assert params.namespace == params_namespace
+
+
+def test_configured_headers_are_not_changed(tenant_label):
+    configured = {"X-Scope-OrgID": "default"}
+    _headers(namespace="shop", prometheus_additional_headers=configured)
+    assert configured == {"X-Scope-OrgID": "default"}