functions passed into macros that have mutliple args will parse correctly (#98)

Author: bossinc

macros.go

@@ -126,22 +126,16 @@ var DefaultMacros Macros = Macros{
 	"column":     macroColumn,
 }
 
-func trimAll(s []string) []string {
-	r := make([]string, len(s))
-	for i, v := range s {
-		r[i] = strings.TrimSpace(v)
-	}
-
-	return r
+type Macro struct {
+	Name string
+	Args []string
 }
 
-var pair = map[rune]rune{')': '('}
-
 // 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) ([][]string, error) {
+func getMacroMatches(input string, name string) ([]Macro, error) {
 	macroName := fmt.Sprintf("\\$__%s\\b", name)
-	matchedMacros := [][]string{}
+	matchedMacros := []Macro{}
 	rgx, err := regexp.Compile(macroName)
 
 	if err != nil {
@@ -201,19 +195,43 @@ func getMacroMatches(input string, name string) ([][]string, error) {
 			macroEnd = matched[matchedIndex][1] - macroStart
 		}
 		macroString := inputCopy[0:macroEnd]
-		macroMatch := []string{macroString}
+		macroMatch := Macro{Name: macroString}
 
 		args := ""
 		// if opening parenthesis was found, extract contents as arguments
 		if argStart > 0 {
 			args = inputCopy[argStart : macroEnd-1]
 		}
-		macroMatch = append(macroMatch, args)
+		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{}
+		}
+	}
+	argsArray = append(argsArray, strings.TrimSpace(string(phrase)))
+	return argsArray
+}
+
 // Interpolate returns an interpolated query string given a backend.DataQuery
 func Interpolate(driver Driver, query *Query) (string, error) {
 	macros := Macros{}
@@ -223,41 +241,23 @@ func Interpolate(driver Driver, query *Query) (string, error) {
 	rawSQL := query.RawSQL
 
 	for key, macro := range macros {
-		matches, err := getMatches(key, rawSQL)
-
+		matches, err := getMacroMatches(rawSQL, key)
 		if err != nil {
 			return rawSQL, err
 		}
-		for _, match := range matches {
-			if len(match) == 0 {
-				// There were no matches for this macro
-				continue
-			}
-
-			args := []string{}
-			if len(match) > 1 {
-				// This macro has arguments
-				args = trimAll(strings.Split(match[1], ","))
-			}
+		if len(matches) == 0 {
+			continue
+		}
 
-			res, err := macro(query.WithSQL(rawSQL), args)
+		for _, match := range matches {
+			res, err := macro(query.WithSQL(rawSQL), match.Args)
 			if err != nil {
 				return rawSQL, err
 			}
 
-			rawSQL = strings.Replace(rawSQL, match[0], res, -1)
+			rawSQL = strings.Replace(rawSQL, match.Name, res, -1)
 		}
 	}
 
 	return rawSQL, nil
 }
-
-func getMatches(macroName, rawSQL string) ([][]string, error) {
-	parsedInput, err := getMacroMatches(rawSQL, macroName)
-
-	if err != nil {
-		return nil, err
-	}
-
-	return parsedInput, err
-}

macros_test.go

@@ -32,6 +32,13 @@ func (h *MockDB) Macros() (macros Macros) {
 		"timeGroup": func(query *Query, args []string) (out string, err error) {
 			return "grouped!", nil
 		},
+		"multiParams": func(query *Query, args []string) (out string, err error) {
+			r := "bar"
+			for _, v := range args {
+				r += "_" + v
+			}
+			return r, nil
+		},
 	}
 }
 
@@ -71,6 +78,8 @@ func TestInterpolate(t *testing.T) {
 		{input: "select $__column from $__table", output: "select my_col from my_table", name: "table and column macros"},
 		{input: "select * from table where ( datetime >= $__foo() ) AND ( datetime <= $__foo() ) limit 100", output: "select * from table where ( datetime >= bar ) AND ( datetime <= bar ) limit 100", name: "macro functions inside more complex clauses"},
 		{input: "select * from table where ( datetime >= $__foo ) AND ( datetime <= $__foo ) limit 100", output: "select * from table where ( datetime >= bar ) AND ( datetime <= bar ) limit 100", name: "macros inside more complex clauses"},
+		{input: "select * from foo where $__multiParams(foo, bar)", output: "select * from foo where bar_foo_bar", name: "macro with multiple parameters"},
+		{input: "select * from foo where $__params(FUNC(foo, bar))", output: "select * from foo where bar_FUNC(foo, bar)", name: "function in macro with multiple parameters"},
 	}
 	for i, tc := range tests {
 		driver := MockDB{}
@@ -90,14 +99,14 @@ func TestInterpolate(t *testing.T) {
 func TestGetMatches(t *testing.T) {
 	t.Run("FindAllStringSubmatch returns DefaultMacros", func(t *testing.T) {
 		for macroName := range DefaultMacros {
-			matches, err := getMatches(macroName, fmt.Sprintf("$__%s", macroName))
+			matches, err := getMacroMatches(fmt.Sprintf("$__%s", macroName), macroName)
 
 			assert.NoError(t, err)
-			assert.Equal(t, [][]string{{fmt.Sprintf("$__%s", macroName), ""}}, matches)
+			assert.Equal(t, []Macro{{Name: fmt.Sprintf("$__%s", macroName), Args: []string{""}}}, matches)
 		}
 	})
 	t.Run("does not return matches for macro name which is substring", func(t *testing.T) {
-		matches, err := getMatches("timeFilter", "$__timeFilterEpoch(time_column)")
+		matches, err := getMacroMatches("$__timeFilterEpoch(time_column)", "timeFilter")
 
 		assert.NoError(t, err)
 		assert.Nil(t, matches)