From 68d25a89894e141f208b62ca3c2c220a37cbc013 Mon Sep 17 00:00:00 2001
From: Srikanth Chekuri <srikanth.chekuri92@gmail.com>
Date: Wed, 30 Oct 2024 14:12:45 +0530
Subject: [PATCH] fix: add support for {{.Labels.<key>}} with dots in `key` for
 template (#6282)

---
 pkg/query-service/rules/templates.go      | 22 +++++++++++++++++-----
 pkg/query-service/rules/templates_test.go | 11 +++++++++++
 2 files changed, 28 insertions(+), 5 deletions(-)

diff --git a/pkg/query-service/rules/templates.go b/pkg/query-service/rules/templates.go
index 519ff3859b..b5487011ce 100644
--- a/pkg/query-service/rules/templates.go
+++ b/pkg/query-service/rules/templates.go
@@ -234,15 +234,27 @@ func AlertTemplateData(labels map[string]string, value string, threshold string)
 // If there is a go template block, it won't be replaced.
 // The example for existing go template block is: {{$threshold}} or {{$value}} or any other valid go template syntax.
 func (te *TemplateExpander) preprocessTemplate() {
-	re := regexp.MustCompile(`({{.*?}})|(\$(\w+(?:\.\w+)*))`)
-	te.text = re.ReplaceAllStringFunc(te.text, func(match string) string {
+	// Handle the $variable syntax
+	reDollar := regexp.MustCompile(`({{.*?}})|(\$(\w+(?:\.\w+)*))`)
+	te.text = reDollar.ReplaceAllStringFunc(te.text, func(match string) string {
 		if strings.HasPrefix(match, "{{") {
 			// If it's a Go template block, leave it unchanged
 			return match
 		}
-		// Otherwise, it's our custom $variable syntax
-		path := strings.Split(match[1:], ".")
-		return "{{index $labels \"" + strings.Join(path, ".") + "\"}}"
+		path := match[1:] // Remove the '$'
+		return fmt.Sprintf(`{{index $labels "%s"}}`, path)
+	})
+
+	// Handle the {{.Labels.service.name}} syntax
+	reLabels := regexp.MustCompile(`{{\s*\.Labels\.([a-zA-Z0-9_.]+)(.*?)}}`)
+	te.text = reLabels.ReplaceAllStringFunc(te.text, func(match string) string {
+		submatches := reLabels.FindStringSubmatch(match)
+		if len(submatches) < 3 {
+			return match // Should not happen
+		}
+		path := submatches[1]
+		rest := submatches[2]
+		return fmt.Sprintf(`{{index .Labels "%s"%s}}`, path, rest)
 	})
 }
 
diff --git a/pkg/query-service/rules/templates_test.go b/pkg/query-service/rules/templates_test.go
index 737a3d20ca..66d958e8f3 100644
--- a/pkg/query-service/rules/templates_test.go
+++ b/pkg/query-service/rules/templates_test.go
@@ -63,3 +63,14 @@ func TestTemplateExpander_WithMissingKey(t *testing.T) {
 	}
 	require.Equal(t, "test  exceeds 100 and observed at 200", result)
 }
+
+func TestTemplateExpander_WithLablesDotSyntax(t *testing.T) {
+	defs := "{{$labels := .Labels}}{{$value := .Value}}{{$threshold := .Threshold}}"
+	data := AlertTemplateData(map[string]string{"service.name": "my-service"}, "200", "100")
+	expander := NewTemplateExpander(context.Background(), defs+"test {{.Labels.service.name}} exceeds {{$threshold}} and observed at {{$value}}", "test", data, times.Time(time.Now().Unix()), nil)
+	result, err := expander.Expand()
+	if err != nil {
+		t.Fatal(err)
+	}
+	require.Equal(t, "test my-service exceeds 100 and observed at 200", result)
+}