Adds prettify support for quoted UTF8 label and metric names

lexer.go

@@ -399,6 +399,29 @@ func unescapeIdent(s string) string {
 	}
 }
 
+func hasEscapedChars(s string) bool {
+	i := 0
+	for i < len(s) {
+		r, size := utf8.DecodeRuneInString(s[i:])
+		if i == 0 && !isFirstIdentChar(r) || i > 0 && !isIdentChar(r) {
+			return true
+		}
+		i += size
+	}
+	return false
+}
+
+func appendQuotedIdent(dst []byte, s string) []byte {
+	dst = utf8.AppendRune(dst, '"')
+	for i := 0; i < len(s); {
+		r, size := utf8.DecodeRuneInString(s[i:])
+		dst = utf8.AppendRune(dst, r)
+		i += size
+	}
+	dst = utf8.AppendRune(dst, '"')
+	return dst
+}
+
 func appendEscapedIdent(dst []byte, s string) []byte {
 	i := 0
 	for i < len(s) {
@@ -413,6 +436,13 @@ func appendEscapedIdent(dst []byte, s string) []byte {
 	return dst
 }
 
+func ifEscapedCharsAppendQuotedIdent(dst []byte, s string) []byte {
+	if hasEscapedChars(s) {
+		return appendQuotedIdent(dst, s)
+	}
+	return appendEscapedIdent(dst, s)
+}
+
 func (lex *lexer) Prev() {
 	lex.nextTokens = append(lex.nextTokens, lex.Token)
 	lex.Token = lex.prevTokens[len(lex.prevTokens)-1]

parser.go

@@ -1462,7 +1462,7 @@ type labelFilterExpr struct {
 }
 
 func (lfe *labelFilterExpr) AppendString(dst []byte) []byte {
-	dst = appendEscapedIdent(dst, lfe.Label)
+	dst = ifEscapedCharsAppendQuotedIdent(dst, lfe.Label)
 	if lfe.Value == nil {
 		return dst
 	}
@@ -2267,15 +2267,28 @@ type MetricExpr struct {
 func appendLabelFilterss(dst []byte, lfss [][]*labelFilterExpr) []byte {
 	offset := 0
 	metricName := getMetricNameFromLabelFilterss(lfss)
+	metricNameHasEscapedChars := hasEscapedChars(metricName)
+
 	if metricName != "" {
 		offset = 1
-		dst = appendEscapedIdent(dst, metricName)
+		if !metricNameHasEscapedChars {
+			dst = appendEscapedIdent(dst, metricName)
+		} else {
+			dst = append(dst, '{')
+			dst = appendQuotedIdent(dst, metricName)
+		}
 	}
 	if isOnlyMetricNameInLabelFilterss(lfss) {
+		if metricNameHasEscapedChars {
+			dst = append(dst, '}')
+		}
 		return dst
 	}
-
-	dst = append(dst, '{')
+	if !metricNameHasEscapedChars {
+		dst = append(dst, '{')
+	} else {
+		dst = append(dst, ',', ' ')
+	}
 	for i, lfs := range lfss {
 		lfs = lfs[offset:]
 		if len(lfs) == 0 {
@@ -2335,6 +2348,9 @@ func mustGetMetricName(lfss []*labelFilterExpr) string {
 	}
 	lfs := lfss[0]
 	if lfs.Label != "__name__" || lfs.Value == nil || len(lfs.Value.tokens) != 1 {
+		if lfs.IsPossibleMetricName {
+			return lfs.Label
+		}
 		return ""
 	}
 	metricName, err := extractStringValue(lfs.Value.tokens[0])

prettifier.go

@@ -160,13 +160,21 @@ func appendPrettifiedExpr(dst []byte, e Expr, indent int, needParens bool) []byt
 		lfss := t.labelFilterss
 		offset := 0
 		metricName := getMetricNameFromLabelFilterss(lfss)
+		metricNamehasEscapedChars := hasEscapedChars(metricName)
+
 		if metricName != "" {
 			offset = 1
 		}
 		dst = appendIndent(dst, indent)
-		dst = appendEscapedIdent(dst, metricName)
+		if !metricNamehasEscapedChars {
+			dst = appendEscapedIdent(dst, metricName)
+		}
 		if !isOnlyMetricNameInLabelFilterss(lfss) {
 			dst = append(dst, "{\n"...)
+			if metricNamehasEscapedChars {
+				dst = ifEscapedCharsAppendQuotedIdent(dst, metricName)
+				dst = append(dst, ',', ' ', '\n')
+			}
 			for i, lfs := range lfss {
 				lfs = lfs[offset:]
 				if len(lfs) == 0 {

prettifier_test.go

@@ -101,7 +101,16 @@ func TestPrettifySuccess(t *testing.T) {
 
 	// Verify that short queries remain single-line
 	same(`foo`)
+	another(`{"foo"}`, "foo")
+	another(`{"3foo"}`, `{"3foo"}`)
+	another(`{"foo", bar="baz"}`, `foo{bar="baz"}`)
+	another(`{"foo", "bar"="baz"}`, `foo{bar="baz"}`)
 	same(`foo{bar="baz"}`)
+	same(`foo{"3bar"="baz"}`)
+	another(`foo{"bar3"="baz"}`, `foo{bar3="baz"}`)
+	same(`"metr\"ic"`)
+	same(`'metr"ic'`)
+	same(`{"3foo", bar="baz"}`)
 	same(`foo{bar="baz",x="y" or q="w",r="t"}`)
 	same(`foo{bar="baz"} + rate(x{y="x"}[5m] offset 1h)`)
 
@@ -114,6 +123,19 @@ func TestPrettifySuccess(t *testing.T) {
 	another(`process_cpu_seconds_total{foo="bar",xjljljlkjopiwererrewre="asdfdsfdsfsdfdsfjkljlk"}`,
 		`process_cpu_seconds_total{
   foo="bar",xjljljlkjopiwererrewre="asdfdsfdsfsdfdsfjkljlk"
+}`)
+	another(`process_cpu_seconds_total{"3foo"="bar",xjljljlkjopiwererrewre="asdfdsfdsfsdfdsfjkljlk"}`,
+		`process_cpu_seconds_total{
+  "3foo"="bar",xjljljlkjopiwererrewre="asdfdsfdsfsdfdsfjkljlk"
+}`)
+	another(`{"process_cpu_seconds_total", "foo"="bar",xjljljlkjopiwererrewre="asdfdsfdsfsdfdsfjkljlk"}`,
+		`process_cpu_seconds_total{
+  foo="bar",xjljljlkjopiwererrewre="asdfdsfdsfsdfdsfjkljlk"
+}`)
+	another(`{"3process_cpu_seconds_total", "foo"="bar",xjljljlkjopiwererrewre="asdfdsfdsfsdfdsfjkljlk"}`,
+		`{
+"3process_cpu_seconds_total", 
+  foo="bar",xjljljlkjopiwererrewre="asdfdsfdsfsdfdsfjkljlk"
 }`)
 	another(`process_cpu_seconds_total{foo="bar",xjljljlkjopiwererrewre="asdfdsfdsfsdfdsfjkljlk",very_long_label_aaaaaaaaaaaaaaa="fdsfdsffdsfs"}`,
 		`process_cpu_seconds_total{
@@ -149,6 +171,16 @@ func TestPrettifySuccess(t *testing.T) {
   sum(rate(node_cpu_seconds_total{mode!="idle"}[5m]) keep_metric_names)
 ) keep_metric_names`)
 
+	another(`(sum(rate({"3process_cpu_seconds_total", instance="foo",job="bar"}[5m] offset 1h @ start())) by (x) / on(x) group_right(y) prefix "x" sum(rate(node_cpu_seconds_total{mode!="idle"}[5m]) keep_metric_names)) keep_metric_names`,
+		`(
+  sum(
+    rate(
+      {"3process_cpu_seconds_total", instance="foo",job="bar"}[5m] offset 1h @ start()
+    )
+  ) by(x)
+    / on(x) group_right(y) prefix "x"
+  sum(rate(node_cpu_seconds_total{mode!="idle"}[5m]) keep_metric_names)
+) keep_metric_names`)
 	another(`process_cpu_seconds_total{aaaaaaaaaaaaaaaaaa="bbbbbb"} offset 5m + (rate(xxxxxxxxxxxxxxxx{yyyyyyyy="aaaaaaa"}) keep_metric_names)`,
 		`(process_cpu_seconds_total{aaaaaaaaaaaaaaaaaa="bbbbbb"} offset 5m)
   +