Merge pull request #12170 from fpetkovski/parser-inject-functions

parser: Allow parsing arbitrary functions

promql/parser/functions.go

@@ -387,7 +387,7 @@ var Functions = map[string]*Function{
 }
 
 // getFunction returns a predefined Function object for the given name.
-func getFunction(name string) (*Function, bool) {
+func getFunction(name string, functions map[string]*Function) (*Function, bool) {
-	function, ok := Functions[name]
+	function, ok := functions[name]
 	return function, ok
 }

promql/parser/generated_parser.y

@@ -339,7 +339,7 @@ grouping_label  : maybe_label
 
 function_call   : IDENTIFIER function_call_body
                         {
-                        fn, exist := getFunction($1.Val)
+                        fn, exist := getFunction($1.Val, yylex.(*parser).functions)
                         if !exist{
                                 yylex.(*parser).addParseErrf($1.PositionRange(),"unknown function with name %q", $1.Val)
                         }

promql/parser/generated_parser.y.go

@@ -1210,7 +1210,7 @@ yydefault:
 		yyDollar = yyS[yypt-2 : yypt+1]
 //line promql/parser/generated_parser.y:341
 		{
-			fn, exist := getFunction(yyDollar[1].item.Val)
+			fn, exist := getFunction(yyDollar[1].item.Val, yylex.(*parser).functions)
 			if !exist {
 				yylex.(*parser).addParseErrf(yyDollar[1].item.PositionRange(), "unknown function with name %q", yyDollar[1].item.Val)
 			}

promql/parser/parse.go

@@ -37,12 +37,20 @@ var parserPool = sync.Pool{
 	},
 }
 
+type Parser interface {
+	ParseExpr() (Expr, error)
+	Close()
+}
+
 type parser struct {
 	lex Lexer
 
 	inject    ItemType
 	injecting bool
 
+	// functions contains all functions supported by the parser instance.
+	functions map[string]*Function
+
 	// Everytime an Item is lexed that could be the end
 	// of certain expressions its end position is stored here.
 	lastClosing Pos
@@ -53,6 +61,63 @@ type parser struct {
 	parseErrors           ParseErrors
 }
 
+type Opt func(p *parser)
+
+func WithFunctions(functions map[string]*Function) Opt {
+	return func(p *parser) {
+		p.functions = functions
+	}
+}
+
+// NewParser returns a new parser.
+// nolint:revive
+func NewParser(input string, opts ...Opt) *parser {
+	p := parserPool.Get().(*parser)
+
+	p.functions = Functions
+	p.injecting = false
+	p.parseErrors = nil
+	p.generatedParserResult = nil
+
+	// Clear lexer struct before reusing.
+	p.lex = Lexer{
+		input: input,
+		state: lexStatements,
+	}
+
+	// Apply user define options.
+	for _, opt := range opts {
+		opt(p)
+	}
+
+	return p
+}
+
+func (p *parser) ParseExpr() (expr Expr, err error) {
+	defer p.recover(&err)
+
+	parseResult := p.parseGenerated(START_EXPRESSION)
+
+	if parseResult != nil {
+		expr = parseResult.(Expr)
+	}
+
+	// Only typecheck when there are no syntax errors.
+	if len(p.parseErrors) == 0 {
+		p.checkAST(expr)
+	}
+
+	if len(p.parseErrors) != 0 {
+		err = p.parseErrors
+	}
+
+	return expr, err
+}
+
+func (p *parser) Close() {
+	defer parserPool.Put(p)
+}
+
 // ParseErr wraps a parsing error with line and position context.
 type ParseErr struct {
 	PositionRange PositionRange
@@ -105,32 +170,15 @@ func (errs ParseErrors) Error() string {
 
 // ParseExpr returns the expression parsed from the input.
 func ParseExpr(input string) (expr Expr, err error) {
-	p := newParser(input)
+	p := NewParser(input)
-	defer parserPool.Put(p)
+	defer p.Close()
-	defer p.recover(&err)
+	return p.ParseExpr()
-
-	parseResult := p.parseGenerated(START_EXPRESSION)
-
-	if parseResult != nil {
-		expr = parseResult.(Expr)
-	}
-
-	// Only typecheck when there are no syntax errors.
-	if len(p.parseErrors) == 0 {
-		p.checkAST(expr)
-	}
-
-	if len(p.parseErrors) != 0 {
-		err = p.parseErrors
-	}
-
-	return expr, err
 }
 
 // ParseMetric parses the input into a metric
 func ParseMetric(input string) (m labels.Labels, err error) {
-	p := newParser(input)
+	p := NewParser(input)
-	defer parserPool.Put(p)
+	defer p.Close()
 	defer p.recover(&err)
 
 	parseResult := p.parseGenerated(START_METRIC)
@@ -148,8 +196,8 @@ func ParseMetric(input string) (m labels.Labels, err error) {
 // ParseMetricSelector parses the provided textual metric selector into a list of
 // label matchers.
 func ParseMetricSelector(input string) (m []*labels.Matcher, err error) {
-	p := newParser(input)
+	p := NewParser(input)
-	defer parserPool.Put(p)
+	defer p.Close()
 	defer p.recover(&err)
 
 	parseResult := p.parseGenerated(START_METRIC_SELECTOR)
@@ -164,22 +212,6 @@ func ParseMetricSelector(input string) (m []*labels.Matcher, err error) {
 	return m, err
 }
 
-// newParser returns a new parser.
-func newParser(input string) *parser {
-	p := parserPool.Get().(*parser)
-
-	p.injecting = false
-	p.parseErrors = nil
-	p.generatedParserResult = nil
-
-	// Clear lexer struct before reusing.
-	p.lex = Lexer{
-		input: input,
-		state: lexStatements,
-	}
-	return p
-}
-
 // SequenceValue is an omittable value in a sequence of time series values.
 type SequenceValue struct {
 	Value   float64
@@ -200,10 +232,10 @@ type seriesDescription struct {
 
 // ParseSeriesDesc parses the description of a time series.
 func ParseSeriesDesc(input string) (labels labels.Labels, values []SequenceValue, err error) {
-	p := newParser(input)
+	p := NewParser(input)
 	p.lex.seriesDesc = true
 
-	defer parserPool.Put(p)
+	defer p.Close()
 	defer p.recover(&err)
 
 	parseResult := p.parseGenerated(START_SERIES_DESCRIPTION)
@@ -799,7 +831,7 @@ func MustLabelMatcher(mt labels.MatchType, name, val string) *labels.Matcher {
 }
 
 func MustGetFunction(name string) *Function {
-	f, ok := getFunction(name)
+	f, ok := getFunction(name, Functions)
 	if !ok {
 		panic(fmt.Errorf("function %q does not exist", name))
 	}

promql/parser/parse_test.go

@@ -3714,7 +3714,7 @@ func TestParseSeries(t *testing.T) {
 }
 
 func TestRecoverParserRuntime(t *testing.T) {
-	p := newParser("foo bar")
+	p := NewParser("foo bar")
 	var err error
 
 	defer func() {
@@ -3728,7 +3728,7 @@ func TestRecoverParserRuntime(t *testing.T) {
 }
 
 func TestRecoverParserError(t *testing.T) {
-	p := newParser("foo bar")
+	p := NewParser("foo bar")
 	var err error
 
 	e := errors.New("custom error")
@@ -3776,3 +3776,20 @@ func TestExtractSelectors(t *testing.T) {
 		require.Equal(t, expected, ExtractSelectors(expr))
 	}
 }
+
+func TestParseCustomFunctions(t *testing.T) {
+	funcs := Functions
+	funcs["custom_func"] = &Function{
+		Name:       "custom_func",
+		ArgTypes:   []ValueType{ValueTypeMatrix},
+		ReturnType: ValueTypeVector,
+	}
+	input := "custom_func(metric[1m])"
+	p := NewParser(input, WithFunctions(funcs))
+	expr, err := p.ParseExpr()
+	require.NoError(t, err)
+
+	call, ok := expr.(*Call)
+	require.True(t, ok)
+	require.Equal(t, "custom_func", call.Func.Name)
+}