From 98402e9f8a23ebaa56bb42dae8f95bdec6ee8b22 Mon Sep 17 00:00:00 2001
From: Xu Bin <140785332+Reditiny@users.noreply.github.com>
Date: Fri, 6 Dec 2024 16:07:47 +0800
Subject: [PATCH] fix: quotation mark for alert rule var (#2339)

---
 alert/eval/eval.go      | 11 +++++++++--
 alert/eval/eval_test.go |  8 ++++----
 2 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/alert/eval/eval.go b/alert/eval/eval.go
index e612e459..dbbec71e 100644
--- a/alert/eval/eval.go
+++ b/alert/eval/eval.go
@@ -377,7 +377,7 @@ func (arw *AlertRuleWorker) VarFillingAfterQuery(query models.PromQuery, readerC
 				for _, paramKey := range ParamKeys {
 					val := string(seqVals[i].Metric[model.LabelName(varToLabel[paramKey])])
 					cur = append(cur, val)
-					curRealQuery = strings.Replace(curRealQuery, fmt.Sprintf("\"$%s\"", paramKey), fmt.Sprintf("\"%s\"", val), -1)
+					curRealQuery = fillVar(curRealQuery, paramKey, val)
 				}
 
 				if _, ok := paramPermutation[strings.Join(cur, "-")]; ok {
@@ -1271,7 +1271,7 @@ func (arw *AlertRuleWorker) VarFillingBeforeQuery(query models.PromQuery, reader
 				realPromql := curPromql
 				split := strings.Split(paramPermutationKeys, "-")
 				for j := range ParamKeys {
-					realPromql = strings.Replace(realPromql, fmt.Sprintf("$%s", ParamKeys[j]), split[j], -1)
+					realPromql = fillVar(realPromql, ParamKeys[j], split[j])
 				}
 				keyToPromql[paramPermutationKeys] = realPromql
 			}
@@ -1373,6 +1373,7 @@ func ExtractVarMapping(promql string) map[string]string {
 
 			key := strings.TrimSpace(kv[0])
 			value := strings.Trim(strings.TrimSpace(kv[1]), "\"")
+			value = strings.Trim(value, "'")
 
 			// 检查值是否为变量(以$开头)
 			if strings.HasPrefix(value, "$") {
@@ -1387,3 +1388,9 @@ func ExtractVarMapping(promql string) map[string]string {
 
 	return varMapping
 }
+
+func fillVar(curRealQuery string, paramKey string, val string) string {
+	curRealQuery = strings.Replace(curRealQuery, fmt.Sprintf("'$%s'", paramKey), fmt.Sprintf("'%s'", val), -1)
+	curRealQuery = strings.Replace(curRealQuery, fmt.Sprintf("\"$%s\"", paramKey), fmt.Sprintf("\"%s\"", val), -1)
+	return curRealQuery
+}
diff --git a/alert/eval/eval_test.go b/alert/eval/eval_test.go
index cd2c6b5c..ba5a11a9 100644
--- a/alert/eval/eval_test.go
+++ b/alert/eval/eval_test.go
@@ -340,7 +340,7 @@ func Test_removeVal(t *testing.T) {
 		{
 			name: "removeVal7",
 			args: args{
-				promql: "mem{test1=\"test1\",test2=\"test2\",test3=\"$test3\"} > $val",
+				promql: "mem{test1=\"test1\",test2=\"test2\",test3='$test3'} > $val",
 			},
 			want: "mem{test1=\"test1\",test2=\"test2\"} > $val",
 		},
@@ -361,16 +361,16 @@ func Test_removeVal(t *testing.T) {
 		{
 			name: "removeVal10",
 			args: args{
-				promql: "mem{test1=\"test1\",test2=\"$test2\"} > $val1 and mem{test3=\"test3\",test4=\"test4\"} > $val2",
+				promql: "mem{test1=\"test1\",test2='$test2'} > $val1 and mem{test3=\"test3\",test4=\"test4\"} > $val2",
 			},
 			want: "mem{test1=\"test1\"} > $val1 and mem{test3=\"test3\",test4=\"test4\"} > $val2",
 		},
 		{
 			name: "removeVal11",
 			args: args{
-				promql: "mem{test1=\"test1\",test2=\"test2\"} > $val1 and mem{test3=\"$test3\",test4=\"test4\"} > $val2",
+				promql: "mem{test1='test1',test2=\"test2\"} > $val1 and mem{test3=\"$test3\",test4=\"test4\"} > $val2",
 			},
-			want: "mem{test1=\"test1\",test2=\"test2\"} > $val1 and mem{test4=\"test4\"} > $val2",
+			want: "mem{test1='test1',test2=\"test2\"} > $val1 and mem{test4=\"test4\"} > $val2",
 		},
 		{
 			name: "removeVal12",