Add workload_identity class attribute

robusta-dev · Dec 12, 2023 · 3d93eaf · 3d93eaf
1 parent 3f63f33
commit 3d93eaf
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
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 != "" or  # Managed Identity Auth
+            config.azure_use_workload_id != ""  # Workload Identity Auth
         )
 
     @classmethod
@@ -60,6 +62,9 @@ def _post_azure_token_endpoint(cls, config: PrometheusConfig):
             }
         # Fallback to Azure Service Principal
         if not token:
+            if config.azure_use_workload_id:
+                logging.error(f"Could not open token file from {token_file}")
+                return {}
             data = {
                 "grant_type": "client_credentials",
                 "client_id": config.azure_client_id,
@@ -80,7 +85,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
@@ -64,6 +64,7 @@ class AzurePrometheusConfig(PrometheusConfig):
     azure_metadata_endpoint: str
     azure_token_endpoint: str
     azure_use_managed_id: Optional[str]
+    azure_use_workload_id: Optional[str]
     azure_client_id: Optional[str]
     azure_tenant_id: Optional[str]
     azure_client_secret: Optional[str]