Add support for Azure Workload Identity

EppO · fmonbillard · commit 9279ea13b3cf · 2025-09-25T12:14:56.000-04:00
diff --git a/README.md b/README.md
@@ -26,7 +26,7 @@ pip install prometrix
 ```
 
 
-> **⚠️ Note:** For Python **3.8 support**, you must use version **0.2.3 or below** of `prometrix`.  
+> **⚠️ Note:** For Python **3.8 support**, you must use version **0.2.3 or below** of `prometrix`.
 > From `0.2.4` onward, `prometrix` requires **Python ≥ 3.9**.
 
 Usage
@@ -84,6 +84,8 @@ azure_config = AzurePrometheusConfig(
 azure_client = get_custom_prometheus_connect(azure_config)
 ```
 
+prometrix supports Azure Managed Identity and Azure Workload Identity. See [Azure Authentication Setup](/azure-auth-setup.rst) for more details.
+
 Similar configuration and creation can be done for EKS, Thanos, and Victoria Metrics Prometheus.
 
 > **_NOTE:_** You need to replace the placeholder values (e.g., YOUR_CORALOGIX_PROMETHEUS_TOKEN) with your actual credentials and endpoints.
diff --git a/azure-auth-setup.rst b/azure-auth-setup.rst
@@ -0,0 +1,156 @@
+Azure managed Prometheus
+*************************
+
+In order to authenticate against the Azure Monitor Workspace Query endpoint, you have multiple options:
+
+- Create an Azure Active Directory authentication app (Option #1)
+
+  - Pros:
+
+    - Quick setup. Just need to create an app, get the credentials and add them to the manifests
+    - Other pods can't use the Service Principal without having the secret
+
+  - Cons:
+
+    - Requires a service principal (Azure AD permission)
+    - Need the client secret in the kubernetes manifests
+    - Client secret expires, you need to manage its rotation
+
+- Use Kubelet's Managed Identity (Option #2)
+
+  - Pros:
+
+    - Quick setup. Get the Managed Identity Client ID and add them to the manifests
+    - No need to manage secrets. Removing the password element decreases the risk of the credentials being compromised
+
+  - Cons:
+
+    - Managed Identity is bound to the AKS nodepool, so any pods can use it if they know/get the client ID
+
+- Use Azure AD Workload Identity (Option #3)
+
+  - Pros:
+
+    - Most secure option as Managed Identity is only bound to the pod. No other pods can use it
+    - No need to manage secrets. Removing the password element decreases the risk of the credentials being compromised
+
+  - Cons:
+
+    - Extra setup needed: need AKS cluster with Workload Identity add-on enabled, get the OIDC issuer URL and add it to the manifests
+
+Get the Azure prometheus query endpoint
+=========================================
+
+1. Go to `Azure Monitor workspaces <https://portal.azure.com/#view/HubsExtension/BrowseResource/resourceType/microsoft.monitor%2Faccounts>`_ and choose your monitored workspace.
+2. In your monitored workspace, `overview`, find the ``Query endpoint`` and copy it.
+
+Option #1: Create an Azure authentication app
+==============================================
+
+We will now create an Azure authentication app and get the necesssary credentials so Robusta can access Prometheus data.
+
+1. Follow this Azure guide to `Register an app with Azure Active Directory <https://learn.microsoft.com/en-us/azure/azure-monitor/essentials/prometheus-self-managed-grafana-azure-active-directory#register-an-app-with-azure-active-directory>`_
+
+.. code-block:: python
+
+    # Create a custom Prometheus client for Azure Prometheus using Service Principal Authentication
+    azure_config = AzurePrometheusConfig(
+        url="https://azure-prometheus.example.com", # Replace with your Azure Monitor workspace query endpoint
+        disable_ssl=False,
+        azure_resource="https://prometheus.monitor.azure.com", # Default resource for Azure Monitor
+        azure_token_endpoint="https://azure-token.example.com",
+        azure_client_id="YOUR_AZURE_CLIENT_ID",
+        azure_tenant_id="YOUR_AZURE_TENANT_ID",
+        azure_client_secret="YOUR_AZURE_CLIENT_SECRET",
+        additional_labels={"job": "azure-prometheus"},
+    )
+    azure_client = get_custom_prometheus_connect(azure_config)
+
+3. Complete the `Allow your app access to your workspace <https://learn.microsoft.com/en-us/azure/azure-monitor/essentials/prometheus-self-managed-grafana-azure-active-directory#allow-your-app-access-to-your-workspace>`_ step, so your app can query data from your Azure Monitor workspace.
+
+Option #2: Use Kubelet's Managed Identity
+==============================================
+
+1. Get the AKS kubelet's Managed Identity Client ID:
+
+.. code-block:: bash
+
+  az aks show -g <resource-group> -n <cluster-name> --query identityProfile.kubeletidentity.clientId -o tsv
+
+2. Set the following settings based from the previous step.
+
+.. code-block:: python
+
+  # Create a custom Prometheus client for Azure Prometheus using Azure Managed Identity
+  azure_config = AzurePrometheusConfig(
+      url="https://azure-prometheus.example.com", # Replace with your Azure Monitor workspace query endpoint
+      disable_ssl=False,
+      azure_use_managed_id=True,
+      azure_resource="https://prometheus.monitor.azure.com", # Default resource for Azure Monitor
+      azure_metadata_endpoint="http://169.254.169.254/metadata/identity/oauth2/token", # Default endpoint for Managed Identity
+      azure_client_id="YOUR_AZURE_CLIENT_ID", # Client ID from step 1
+      azure_tenant_id="YOUR_AZURE_TENANT_ID",
+      additional_labels={"job": "azure-prometheus"},
+  )
+  azure_client = get_custom_prometheus_connect(azure_config)
+
+3. Give access to your Managed Identity on your Azure Monitor Workspace:
+
+   #. Open the Access Control (IAM) page for your Azure Monitor workspace in the Azure portal.
+   #. Select Add role assignment.
+   #. Select Monitoring Data Reader and select Next.
+   #. For Assign access to, select Managed identity.
+   #. Select + Select members.
+   #. Select the Managed Identity you got from step 1
+   #. Select Review + assign to save the configuration.
+
+Option #3: Use Azure Workload Identity (Recommended)
+==============================================
+
+1. Requirements
+
+AKS cluster needs to have Workload Identity add-on and OIDC issuer enabled. You can use `--enable-oidc-issuer --enable-workload-identity` with `az aks create` or `az aks update` to enable them.
+
+2. Create a new Managed Identity. Change the Identity name, resource group and location to match your environment.
+
+.. code-block:: bash
+
+  export SUBSCRIPTION="$(az account show --query id --output tsv)"
+  az identity create --name <identity-name> --resource-group <resource-group> --location "eastus" --subscription "${SUBSCRIPTION}"
+  az identity show --name <identity-name> --resource-group <resource-group> -query clientId -o tsv # keep this value for the step #3
+
+3. Set the following settings based from the previous step.
+
+.. code-block:: python
+
+  # Create a custom Prometheus client for Azure Prometheus using Azure Managed Identity
+  azure_config = AzurePrometheusConfig(
+      url="https://azure-prometheus.example.com", # Replace with your Azure Monitor workspace query endpoint
+      disable_ssl=False,
+      azure_use_workload_id=True,
+      azure_resource="https://prometheus.monitor.azure.com", # Default resource for Azure Monitor
+      azure_token_endpoint="https://azure-token.example.com",
+      azure_client_id="YOUR_AZURE_CLIENT_ID", # Client ID from step 1
+      azure_tenant_id="YOUR_AZURE_TENANT_ID",
+      additional_labels={"job": "azure-prometheus"},
+  )
+  azure_client = get_custom_prometheus_connect(azure_config)
+
+4. Federate the Service Account with the Managed Identity. Replace the values with the ones from the step #1.
+
+.. code-block:: bash
+
+  export AKS_OIDC_ISSUER="$(az aks show -g <resource-group> -n <cluster-name> --query "oidcIssuerProfile.issuerUrl" -otsv)" # Replace with the corresponding values of your AKS clusters.
+  MY_NAMESPACE="mynamespace" # Replace with the namespace where your application is deployed
+  MY_SERVICE_ACCOUNT="my-service-account" # Replace with the service account name used by your application
+  az identity federated-credential create --name <federated-identity-name> --identity-name "robusta-id" --resource-group <resource-group> --issuer ${AKS_OIDC_ISSUER} --subject system:serviceaccount:$MY_NAMESPACE:$MY_SERVICE_ACCOUNT
+
+5. Give access to your Managed Identity on your workspace:
+
+   #. Open the Access Control (IAM) page for your Azure Monitor workspace in the Azure portal.
+   #. Select Add role assignment.
+   #. Select Monitoring Data Reader and select Next.
+   #. For Assign access to, select Managed identity.
+   #. Select + Select members.
+   #. Select the Managed Identity you got from step 2
+   #. Select Review + assign to save the configuration.
diff --git a/prometrix/auth.py b/prometrix/auth.py
@@ -16,7 +16,9 @@ def azure_authorization(cls, config: PrometheusConfig) -> bool:
         if not isinstance(config, AzurePrometheusConfig):
             return False
         return (config.azure_client_id != "" and config.azure_tenant_id != "") and (
-            config.azure_client_secret != "" or config.azure_use_managed_id != ""
+            config.azure_client_secret != "" or     # Service Principal Auth
+            config.azure_use_managed_id != False or # Managed Identity Auth
+            config.azure_use_workload_id != False   # Workload Identity Auth
         )
 
     @classmethod
@@ -48,15 +50,33 @@ def _get_azure_metadata_endpoint(cls, config: PrometheusConfig):
     @no_type_check
     @classmethod
     def _post_azure_token_endpoint(cls, config: PrometheusConfig):
-        return requests.post(
-            url=config.azure_token_endpoint,
-            headers={"Content-Type": "application/x-www-form-urlencoded"},
-            data={
+        # Try Azure Workload Identity
+        with open("/var/run/secrets/azure/tokens/azure-identity-token", "r") as token_file:
+            token = token_file.read()
+            data = {
+                "grant_type": "client_credentials",
+                "client_assertion_type": "urn:ietf:params:oauth:client-assertion-type:jwt-bearer",
+                "client_assertion": token,
+                "client_id": config.azure_client_id,
+                "scope": f"{config.azure_resource}/.default",
+            }
+        # Fallback to Azure Service Principal
+        if not token:
+            if config.azure_use_workload_id:
+                return {
+                    "ok": False,
+                    "reason": f"Could not open token file from {token_file}",
+                }
+            data = {
                 "grant_type": "client_credentials",
                 "client_id": config.azure_client_id,
                 "client_secret": config.azure_client_secret,
                 "resource": config.azure_resource,
-            },
+            }
+        return requests.post(
+            url=config.azure_token_endpoint,
+            headers={"Content-Type": "application/x-www-form-urlencoded"},
+            data=data,
         )
 
     @classmethod
@@ -67,7 +87,7 @@ def request_new_token(cls, config: PrometheusConfig) -> bool:
             try:
                 if config.azure_use_managed_id:
                     res = cls._get_azure_metadata_endpoint(config)
-                else:
+                else:  # Service Principal and Workload Identity
                     res = cls._post_azure_token_endpoint(config)
             except Exception:
                 logging.exception(
diff --git a/prometrix/models/prometheus_config.py b/prometrix/models/prometheus_config.py
@@ -1,4 +1,5 @@
 from enum import Enum
+import os
 from typing import Dict, List, Optional
 
 try:
@@ -69,12 +70,13 @@ class VictoriaMetricsPrometheusConfig(PrometheusConfig):
 # Does not support labels according to the docs, See below for apis
 # https://learn.microsoft.com/en-us/azure/azure-monitor/essentials/prometheus-api-promql#supported-apis
 class AzurePrometheusConfig(PrometheusConfig):
-    azure_resource: str
-    azure_metadata_endpoint: str
-    azure_token_endpoint: str
-    azure_use_managed_id: Optional[str] = None
+    azure_resource: str = "https://prometheus.monitor.azure.com"
+    azure_metadata_endpoint: str = "http://169.254.169.254/metadata/identity/oauth2/token"
+    azure_token_endpoint: str = f"https://login.microsoftonline.com/{os.environ.get('AZURE_TENANT_ID')}/oauth2/token"
+    azure_use_managed_id: Optional[bool] = False
+    azure_use_workload_id: Optional[bool] = False
     azure_client_id: Optional[str] = None
-    azure_tenant_id: Optional[str] = None
+    azure_tenant_id: Optional[str] = os.environ.get('AZURE_TENANT_ID', '')
     azure_client_secret: Optional[str] = None
     supported_apis: List[PrometheusApis] = [
         PrometheusApis.QUERY,
