From c4e91fd4d6294cdaa03f02677e98e19d86d9b471 Mon Sep 17 00:00:00 2001
From: arik <alon.arik@gmail.com>
Date: Tue, 26 Oct 2021 18:49:51 +0300
Subject: [PATCH] cluster_zone and grafana annotations custom tags

---
 helm/robusta/templates/_helpers.tpl |  3 +++
 helm/robusta/values.yaml            |  1 +
 playbooks/grafana_enrichment.py     | 14 +++++++++++---
 3 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/helm/robusta/templates/_helpers.tpl b/helm/robusta/templates/_helpers.tpl
index 4b613883c..45dd7ea1e 100644
--- a/helm/robusta/templates/_helpers.tpl
+++ b/helm/robusta/templates/_helpers.tpl
@@ -26,6 +26,9 @@ global_config:
   {{- end }}
 {{- end }}
   cluster_name: {{ required "A valid .Values.clusterName entry is required!" .Values.clusterName }}
+  {{- if .Values.clusterZone }}
+  cluster_zone: {{ .Values.clusterZone }}
+  {{- end }}
 active_playbooks:
 {{ toYaml .Values.playbooks | indent 2 }}
 {{ end }}
diff --git a/helm/robusta/values.yaml b/helm/robusta/values.yaml
index 54ae6fc74..582f27a35 100644
--- a/helm/robusta/values.yaml
+++ b/helm/robusta/values.yaml
@@ -1,5 +1,6 @@
 # global parameters
 clusterName: ""
+clusterZone: ""
 
 # slack integration params
 slackApiKey: ""
diff --git a/playbooks/grafana_enrichment.py b/playbooks/grafana_enrichment.py
index 1944d301c..eb5b437eb 100644
--- a/playbooks/grafana_enrichment.py
+++ b/playbooks/grafana_enrichment.py
@@ -3,15 +3,17 @@
 from robusta.api import *
 
 
-class Params(BaseModel):
+class GrafanaAnnotationsParams(BaseModel):
     grafana_url: str = None
     grafana_api_key: SecretStr
     grafana_dashboard_uid: str
     cluster_name: str
+    cluster_zone: str = None
+    custom_tags: List[str] = None
 
 
 @on_deployment_update
-def add_deployment_lines_to_grafana(event: DeploymentEvent, action_params: Params):
+def add_deployment_lines_to_grafana(event: DeploymentEvent, action_params: GrafanaAnnotationsParams):
     """
     Add annotations to grafana whenever a new application version is deployed so that you can easily see changes in performance.
     """
@@ -31,9 +33,15 @@ def add_deployment_lines_to_grafana(event: DeploymentEvent, action_params: Param
     grafana = Grafana(
         action_params.grafana_api_key.get_secret_value(), action_params.grafana_url
     )
+    tags = [event.obj.metadata.name, event.obj.metadata.namespace, action_params.cluster_name]
+    if action_params.cluster_zone:
+        tags.append(action_params.cluster_zone)
+    if action_params.custom_tags:
+        tags.extend(action_params.custom_tags)
+
     grafana.add_line_to_dashboard(
         action_params.grafana_dashboard_uid, msg,
-        tags=[event.obj.metadata.name, event.obj.metadata.namespace, action_params.cluster_name]
+        tags=tags
     )