Refactor getMacroMatches & fix a couple of bugs (#863)

This addresses #859 and #861 and removes some duplicated logic.

Author: njvrzm

data/sqlutil/macros.go

@@ -124,110 +124,72 @@ var DefaultMacros = Macros{
 	"column":     macroColumn,
 }
 
-type Macro struct {
-	Name string
-	Args []string
+type macroMatch struct {
+	full string
+	args []string
 }
 
 // getMacroMatches extracts macro strings with their respective arguments from the sql input given
 // It manually parses the string to find the closing parenthesis of the macro (because regex has no memory)
-func getMacroMatches(input string, name string) ([]Macro, error) {
-	macroName := fmt.Sprintf("\\$__%s\\b", name)
-	matchedMacros := []Macro{}
-	rgx, err := regexp.Compile(macroName)
+func getMacroMatches(input string, name string) ([]macroMatch, error) {
+	rgx, err := regexp.Compile(fmt.Sprintf(`\$__%s\b`, name))
 
 	if err != nil {
 		return nil, err
 	}
 
-	// get all matching macro instances
-	matched := rgx.FindAllStringIndex(input, -1)
+	var matches []macroMatch
+	for _, window := range rgx.FindAllStringIndex(input, -1) {
+		start, end := window[0], window[1]
+		args, length := parseArgs(input[end:])
+		if length < 0 {
+			return nil, fmt.Errorf("failed to parse macro arguments (missing close bracket?)")
+		}
+		matches = append(matches, macroMatch{full: input[start : end+length], args: args})
+	}
+	return matches, nil
+}
 
-	if matched == nil {
-		return nil, nil
+// parseArgs looks for a bracketed argument list at the beginning of argString.
+// If one is present, returns a list of whitespace-trimmed arguments and the
+// length of the string comprising the bracketed argument list.
+func parseArgs(argString string) ([]string, int) {
+	if !strings.HasPrefix(argString, "(") {
+		return nil, 0 // single empty arg for backwards compatibility
 	}
 
-	for matchedIndex := 0; matchedIndex < len(matched); matchedIndex++ {
-		var macroEnd = 0
-		var argStart = 0
-		// quick exit from the loop, when we encounter a closing bracket before an opening one (ie "($__macro)", where we can skip the closing one from the result)
-		var forceBreak = false
-		macroStart := matched[matchedIndex][0]
-		inputCopy := input[macroStart:]
-		cache := make([]rune, 0)
+	var args []string
+	depth := 0
+	arg := []rune{}
 
-		// find the opening and closing arguments brackets
-		for idx, r := range inputCopy {
-			if len(cache) == 0 && macroEnd > 0 || forceBreak {
-				break
+	for i, r := range argString {
+		switch r {
+		case '(':
+			depth++
+			if depth == 1 {
+				// don't include the outer bracket in the arg
+				continue
 			}
-			switch r {
-			case '(':
-				cache = append(cache, r)
-				if argStart == 0 {
-					argStart = idx + 1
-				}
-			case ' ':
-				// when we are inside an argument, we do not want to exit on space
-				if argStart != 0 {
-					continue
-				}
-				fallthrough
-			case ')':
-				l := len(cache)
-				if l == 0 {
-					macroEnd = 0
-					forceBreak = true
-					break
-				}
-				cache = cache[:l-1]
-				macroEnd = idx + 1
-			default:
+		case ')':
+			depth--
+			if depth == 0 {
+				// closing bracket
+				args = append(args, strings.TrimSpace(string(arg)))
+				return args, i + 1
+			}
+		case ',':
+			if depth == 1 {
+				// a comma at this level is separating args
+				args = append(args, strings.TrimSpace(string(arg)))
+				arg = []rune{}
 				continue
 			}
 		}
-
-		// macroEnd equals to 0 means there are no parentheses, so just set it
-		// to the end of the regex match
-		if macroEnd == 0 {
-			macroEnd = matched[matchedIndex][1] - macroStart
-		}
-		macroString := inputCopy[0:macroEnd]
-		macroMatch := Macro{Name: macroString}
-
-		args := ""
-		// if opening parenthesis was found, extract contents as arguments
-		if argStart > 0 {
-			args = inputCopy[argStart : macroEnd-1]
-		}
-		macroMatch.Args = parseArgs(args)
-		matchedMacros = append(matchedMacros, macroMatch)
-	}
-	return matchedMacros, nil
-}
-
-func parseArgs(args string) []string {
-	argsArray := []string{}
-	phrase := []rune{}
-	bracketCount := 0
-	for _, v := range args {
-		phrase = append(phrase, v)
-		if v == '(' {
-			bracketCount++
-			continue
-		}
-		if v == ')' {
-			bracketCount--
-			continue
-		}
-		if v == ',' && bracketCount == 0 {
-			removeComma := phrase[:len(phrase)-1]
-			argsArray = append(argsArray, string(removeComma))
-			phrase = []rune{}
-		}
+		arg = append(arg, r)
 	}
-	argsArray = append(argsArray, strings.TrimSpace(string(phrase)))
-	return argsArray
+	// If we get here, we have seen an open bracket but not a close bracket. This
+	// would formerly cause a panic; now it is treated as an error.
+	return nil, -1
 }
 
 // Interpolate returns an interpolated query string given a backend.DataQuery
@@ -243,17 +205,14 @@ func Interpolate(query *Query, macros Macros) (string, error) {
 		if err != nil {
 			return rawSQL, err
 		}
-		if len(matches) == 0 {
-			continue
-		}
 
 		for _, match := range matches {
-			res, err := macro(query.WithSQL(rawSQL), match.Args)
+			res, err := macro(query.WithSQL(rawSQL), match.args)
 			if err != nil {
 				return rawSQL, err
 			}
 
-			rawSQL = strings.ReplaceAll(rawSQL, match.Name, res)
+			rawSQL = strings.ReplaceAll(rawSQL, match.full, res)
 		}
 	}
 

data/sqlutil/macros_test.go

@@ -6,7 +6,6 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
 )
 
 func staticMacro(output string) MacroFunc {
@@ -26,8 +25,9 @@ var macros = Macros{
 		}
 		return "bar"
 	}),
-	"f": staticMacro("f(1)"),
-	"g": staticMacro("g(1)"),
+	"f":   staticMacro("f(1)"),
+	"g":   staticMacro("g(1)"),
+	"num": staticMacro("10000"),
 	"multiParams": argMacro(func(args []string) string {
 		return strings.Join(append([]string{"bar"}, args...), "_")
 	}),
@@ -39,9 +39,10 @@ func TestInterpolate(t *testing.T) {
 	tableName := "my_table"
 	tableColumn := "my_col"
 	type test struct {
-		name   string
-		input  string
-		output string
+		name    string
+		input   string
+		output  string
+		wantErr bool
 	}
 	tests := []test{
 		{
@@ -174,6 +175,26 @@ func TestInterpolate(t *testing.T) {
 			output: "select * from foo where f(1) > g(1)",
 			name:   "don't consume args after a space (see https://github.com/grafana/sqlds/issues/82)",
 		},
+		{
+			input:  "select * from foo where $__num*(table.a + table.b) > 1000000",
+			output: "select * from foo where 10000*(table.a + table.b) > 1000000",
+			name:   "don't consume args after other non-paren characters",
+		},
+		{
+			input:   "select * from foo where $__params(whoops",
+			name:    "error (not panic) on missing close bracket",
+			wantErr: true,
+		},
+		{
+			input:   "select * from foo where $__params(FUNC(foo, bar)",
+			name:    "error on missing close bracket, nested",
+			wantErr: true,
+		},
+		{
+			input:   "select * from foo where $__params(FUNC(foo, bar)) > $__timeTo(uhoh",
+			name:    "error on missing close bracket after good macros",
+			wantErr: true,
+		},
 	}
 	for i, tc := range tests {
 		t.Run(fmt.Sprintf("[%d/%d] %s", i+1, len(tests), tc.name), func(t *testing.T) {
@@ -183,8 +204,75 @@ func TestInterpolate(t *testing.T) {
 				Column: tableColumn,
 			}
 			interpolatedQuery, err := Interpolate(query, macros)
-			require.Nil(t, err)
-			assert.Equal(t, tc.output, interpolatedQuery)
+			assert.Equal(t, err != nil, tc.wantErr, "wantErr != gotErr")
+			if !tc.wantErr {
+				assert.Equal(t, tc.output, interpolatedQuery)
+			}
+		})
+	}
+}
+
+func Test_parseArgs(t *testing.T) {
+	var tests = []struct {
+		name       string
+		input      string
+		wantArgs   []string
+		wantLength int
+	}{
+		{
+			name:       "no parens, no args",
+			input:      "foo bar",
+			wantArgs:   nil,
+			wantLength: 0,
+		},
+		{
+			name:       "parens not at beginning, still no args",
+			input:      "foo(bar)",
+			wantArgs:   nil,
+			wantLength: 0,
+		},
+		{
+			name:       "even just a space is enough",
+			input:      " (bar)",
+			wantArgs:   nil,
+			wantLength: 0,
+		},
+		{
+			name:       "simple one-arg case",
+			input:      "(bar)",
+			wantArgs:   []string{"bar"},
+			wantLength: 5,
+		},
+		{
+			name:       "multiple args, spaces",
+			input:      "(bar, baz, quux)",
+			wantArgs:   []string{"bar", "baz", "quux"},
+			wantLength: 16,
+		},
+		{
+			name:       "nested parens are not parsed further",
+			input:      "(bar(some,thing))",
+			wantArgs:   []string{"bar(some,thing)"},
+			wantLength: 17,
+		},
+		{
+			name:       "stuff after the closing bracket is ignored",
+			input:      "(arg1, arg2), not_an_arg, nope",
+			wantArgs:   []string{"arg1", "arg2"},
+			wantLength: 12,
+		},
+		{
+			name:       "missing close bracket is an error",
+			input:      "(arg1, arg2",
+			wantArgs:   nil,
+			wantLength: -1,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			args, length := parseArgs(tt.input)
+			assert.Equalf(t, tt.wantArgs, args, "parseArgs(%v) wrong args:\nwant: %v\ngot:  %v\n", tt.input, tt.wantArgs, args)
+			assert.Equalf(t, tt.wantLength, length, "parseArgs(%v) wrong length:\nwant: %d\ngot:  %d\n", tt.input, tt.wantLength, length)
 		})
 	}
 }
@@ -195,7 +283,7 @@ func TestGetMacroMatches(t *testing.T) {
 			matches, err := getMacroMatches(fmt.Sprintf("$__%s", macroName), macroName)
 
 			assert.NoError(t, err)
-			assert.Equal(t, []Macro{{fmt.Sprintf("$__%s", macroName), []string{""}}}, matches)
+			assert.Equal(t, []macroMatch{{fmt.Sprintf("$__%s", macroName), nil}}, matches)
 		}
 	})
 	t.Run("does not return matches for macro name which is substring", func(t *testing.T) {