Add ParseWithVars that allows to preserve variables in metricsql expressions

Add `PrettifyExpr` that allows to return prettified string generated from existing expression

Author: AndrewChubatiuk

lexer.go

@@ -122,13 +122,12 @@ again:
 		}
 		goto tokenFoundLabel
 	}
-	if strings.HasPrefix(s, "$__interval") {
-		lex.sTail = s[len("$__interval"):]
-		return "$__interval", nil
-	}
-	if strings.HasPrefix(s, "$__rate_interval") {
-		lex.sTail = s[len("$__rate_interval"):]
-		return "$__interval", nil
+	if isVariable(s) {
+		token, err = scanVariable(s)
+		if err != nil {
+			return "", err
+		}
+		goto tokenFoundLabel
 	}
 	return "", fmt.Errorf("cannot recognize %q", s)
 
@@ -301,6 +300,43 @@ func scanPositiveNumber(s string) (string, error) {
 	return s[:j], nil
 }
 
+func isVariable(s string) bool {
+	return len(s) > 1 && s[0] == '$'
+}
+
+func isVariableChar(c byte) bool {
+	return (c >= 'A' && c <= 'Z') || (c >= 'a' && c <= 'z') || (c >= '0' && c <= '9') || c == '_'
+}
+
+func scanVariable(s string) (string, error) {
+	if len(s) < 2 {
+		return "", fmt.Errorf("too small string for a variable %q", s)
+	}
+	i := 1
+	hasBraces := s[i] == '{'
+	if hasBraces {
+		i++
+	}
+	for i < len(s) {
+		if hasBraces {
+			if s[i] == '}' {
+				i++
+				break
+			}
+			if !isVariableChar(s[i]) {
+				return "", fmt.Errorf("not allowed symbol in variable %q", s)
+			}
+		} else if !isVariableChar(s[i]) {
+			break
+		}
+		i++
+	}
+	if hasBraces && i <= 3 || i <= 2 {
+		return "", fmt.Errorf("impossible variable name %q", s)
+	}
+	return s[:i], nil
+}
+
 func scanNumMultiplier(s string) int {
 	if len(s) > 3 {
 		s = s[:3]
@@ -534,7 +570,7 @@ func scanSpecialIntegerPrefix(s string) (skipChars int, isHex bool) {
 }
 
 func isPositiveDuration(s string) bool {
-	if s == "$__interval" {
+	if s == "$__interval" || s == "$__rate_interval" {
 		return true
 	}
 	n := scanDuration(s)
@@ -610,7 +646,7 @@ func DurationValue(s string, step int64) (int64, error) {
 }
 
 func parseSingleDuration(s string, step int64) (float64, error) {
-	if s == "$__interval" {
+	if s == "$__interval" || s == "$__rate_interval" {
 		return float64(step), nil
 	}
 
@@ -676,8 +712,8 @@ func scanSingleDuration(s string, canBeNegative bool) int {
 	if s[0] == '-' && canBeNegative {
 		i++
 	}
-	if s[i:] == "$__interval" {
-		return i + len("$__interval")
+	if s[i:] == "$__interval" || s[i:] == "$__rate_interval" {
+		return i + len(s[i:])
 	}
 	for i < len(s) && isDecimalChar(s[i]) {
 		i++

lexer_test.go

@@ -98,6 +98,33 @@ func TestScanPositiveNumberFailure(t *testing.T) {
 	f("12.34e-")
 }
 
+func TestScanVariableSuccess(t *testing.T) {
+	f := func(s, vsExpected string) {
+		t.Helper()
+		vs, err := scanVariable(s)
+		if err != nil {
+			t.Fatalf("unexpected error in scanVariable(%q): %s", s, err)
+		}
+		if vs != vsExpected {
+			t.Fatalf("unexpected variable scanned from %q; got %q; want %q", s, vs, vsExpected)
+		}
+	}
+	f("$__rate_interval", "$__rate_interval")
+	f("${foobar},test", "${foobar}")
+}
+
+func TestScanVariableFailure(t *testing.T) {
+	f := func(s string) {
+		t.Helper()
+		vs, err := scanVariable(s)
+		if err == nil {
+			t.Fatalf("expecting non-nil error in scanVariable(%q); got result %q", s, vs)
+		}
+	}
+	f("")
+	f("${foobar,")
+}
+
 func TestParsePositiveNumberSuccess(t *testing.T) {
 	f := func(s string, vExpected float64) {
 		t.Helper()

parser.go

@@ -13,8 +13,13 @@ import (
 //
 // MetricsQL is backwards-compatible with PromQL.
 func Parse(s string) (Expr, error) {
+	return ParseWithVars(s, false)
+}
+
+// ParseWithVars provides an option to keep variables in expressions.
+func ParseWithVars(s string, keepVars bool) (Expr, error) {
 	// Parse s
-	e, err := parseInternal(s)
+	e, err := parseInternal(s, keepVars)
 	if err != nil {
 		return nil, err
 	}
@@ -32,8 +37,10 @@ func Parse(s string) (Expr, error) {
 	return e, nil
 }
 
-func parseInternal(s string) (Expr, error) {
-	var p parser
+func parseInternal(s string, keepVars bool) (Expr, error) {
+	p := parser{
+		keepVars: keepVars,
+	}
 	p.lex.Init(s)
 	if err := p.lex.Next(); err != nil {
 		return nil, fmt.Errorf(`cannot find the first token: %s`, err)
@@ -258,7 +265,8 @@ func simplifyConstantsInplace(args []Expr) {
 // postconditions for all parser.parse* funcs:
 // - p.lex.Token should point to the next token after the parsed token.
 type parser struct {
-	lex lexer
+	lex      lexer
+	keepVars bool
 }
 
 func isWith(s string) bool {
@@ -473,6 +481,15 @@ func (p *parser) parseSingleExprWithoutRollupSuffix() (Expr, error) {
 	if isIdentPrefix(p.lex.Token) {
 		return p.parseIdentExpr()
 	}
+	if isVariable(p.lex.Token) {
+		e := &NumberExpr{
+			s: p.lex.Token,
+		}
+		if err := p.lex.Next(); err != nil {
+			return nil, err
+		}
+		return e, nil
+	}
 	switch p.lex.Token {
 	case "(":
 		return p.parseParensExpr()
@@ -1537,7 +1554,12 @@ func (p *parser) parseWindowAndStep() (*DurationExpr, *DurationExpr, bool, error
 	}
 	var window *DurationExpr
 	if !strings.HasPrefix(p.lex.Token, ":") {
-		if p.lex.Token == "$__interval" {
+		if p.lex.Token == "$__interval" || p.lex.Token == "$__rate_interval" {
+			if p.keepVars {
+				window = &DurationExpr{
+					s: p.lex.Token,
+				}
+			}
 			// Skip $__interval, since it must be treated as missing lookbehind window,
 			// e.g. rate(m[$__interval]) must be equivalent to rate(m).
 			// In this case VictoriaMetrics automatically adjusts the lookbehind window
@@ -1656,8 +1678,14 @@ func (p *parser) parsePositiveDuration() (*DurationExpr, error) {
 		}
 	}
 	// Verify duration value.
-	if s == "$__interval" {
-		s = "1i"
+	if s == "$__interval" || s == "$__rate_interval" {
+		if p.keepVars {
+			return &DurationExpr{
+				s: s,
+			}, nil
+		} else {
+			s = "1i"
+		}
 	}
 	return newDurationExpr(s)
 }

parser_test.go

@@ -23,6 +23,7 @@ func TestParseSuccess(t *testing.T) {
 	}
 
 	// metricExpr
+	same(`topk($topk, max(sum(vmalert_recording_rules_last_evaluation_samples{job=~"$job",instance=~"$instance",group=~"$group",file=~"$file"}) by(job,instance,group,file,recording) > 0) by(job,group,file,recording))`)
 	same(`{}`)
 	same(`{}[5m]`)
 	same(`{}[5m:]`)
@@ -643,6 +644,31 @@ func TestParseSuccess(t *testing.T) {
 	another(`rate(m[$__interval:5m])`, `rate(m[:5m])`)
 }
 
+func TestParseWithVarsSuccess(t *testing.T) {
+	another := func(s string, sExpected string) {
+		t.Helper()
+
+		e, err := ParseWithVars(s, true)
+		if err != nil {
+			t.Fatalf("unexpected error when parsing %s: %s", s, err)
+		}
+		res := e.AppendString(nil)
+		if string(res) != sExpected {
+			t.Fatalf("unexpected string constructed;\ngot\n%s\nwant\n%s", res, sExpected)
+		}
+	}
+	same := func(s string) {
+		t.Helper()
+		another(s, s)
+	}
+
+	// $__interval and $__rate_interval must be replaced with 1i
+	same(`rate(m[$__interval] offset $__interval) * $__interval`)
+	another(`increase(m[$__rate_interval] offset -$__rate_interval) + -$__rate_interval`, `increase(m[$__rate_interval] offset -$__rate_interval) + (0 - $__rate_interval)`)
+	same(`rate(m[$__rate_interval:5m])`)
+	same(`rate(m[$__interval:5m])`)
+}
+
 func TestParseError(t *testing.T) {
 	f := func(s string) {
 		t.Helper()

prettifier.go

@@ -2,13 +2,18 @@ package metricsql
 
 // Prettify returns prettified representation of MetricsQL query q.
 func Prettify(q string) (string, error) {
-	e, err := parseInternal(q)
+	return PrettifyWithVars(q, false)
+}
+
+// PrettifyWithVars returns prettified representation of MetricsQL query q and keeps expression variables depending on keepVars value.
+func PrettifyWithVars(q string, keepVars bool) (string, error) {
+	e, err := parseInternal(q, keepVars)
 	if err != nil {
 		return "", err
 	}
 	e = removeParensExpr(e)
 	b := appendPrettifiedExpr(nil, e, 0, false)
-	return string(b), nil
+	return string(b)
 }
 
 // maxPrettifiedLineLen is the maximum length of a single line returned by Prettify().