parse variables in expressions

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

@@ -473,6 +473,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 +1546,7 @@ 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" {
 			// 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,7 +1665,7 @@ func (p *parser) parsePositiveDuration() (*DurationExpr, error) {
 		}
 	}
 	// Verify duration value.
-	if s == "$__interval" {
+	if s == "$__interval" || s == "$__rate_interval" {
 		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:]`)