add support for variables in recurse

gql/parser.go

@@ -85,6 +85,7 @@ type GraphQuery struct {
 type RecurseArgs struct {
 	Depth     uint64
 	AllowLoop bool
+	varMap    map[string]string
 }
 
 // ShortestPathArgs stores the arguments needed to process the shortest path query.
@@ -404,6 +405,36 @@ func substituteVariables(gq *GraphQuery, vmap varMap) error {
 			return err
 		}
 	}
+	if gq.RecurseArgs.varMap != nil {
+		// update depth
+		varName, ok := gq.RecurseArgs.varMap["depth"]
+		if ok {
+			val, ok := vmap[varName]
+			if !ok {
+				return errors.Errorf("variable %s not defined", varName)
+			}
+			depth, err := strconv.ParseUint(val.Value, 0, 64)
+			if err != nil {
+				return err
+			}
+			gq.RecurseArgs.Depth = depth
+		}
+
+		// update loop
+		varName, ok = gq.RecurseArgs.varMap["loop"]
+		if ok {
+			val, ok := vmap[varName]
+			if !ok {
+				return errors.Errorf("variable %s not defined", varName)
+			}
+			allowLoop, err := strconv.ParseBool(val.Value)
+			if err != nil {
+				return err
+			}
+			gq.RecurseArgs.AllowLoop = allowLoop
+		}
+
+	}
 	return nil
 }
 
@@ -766,6 +797,18 @@ L2:
 	return gq, nil
 }
 
+// parseVarName returns the variable name.
+func parseVarName(it *lex.ItemIterator) (string, error) {
+	if !it.Next() {
+		return "", it.Errorf("Expected variable name")
+	}
+	item := it.Item()
+	if item.Typ != itemName {
+		return "", it.Errorf("Expected variable name")
+	}
+	return fmt.Sprintf("$%s", item.Val), nil
+}
+
 func parseRecurseArgs(it *lex.ItemIterator, gq *GraphQuery) error {
 	if ok := trySkipItemTyp(it, itemLeftRound); !ok {
 		// We don't have a (, we can return.
@@ -784,30 +827,59 @@ func parseRecurseArgs(it *lex.ItemIterator, gq *GraphQuery) error {
 		if ok := trySkipItemTyp(it, itemColon); !ok {
 			return it.Errorf("Expected colon(:) after %s", key)
 		}
-
-		if item, ok = tryParseItemType(it, itemName); !ok {
-			return item.Errorf("Expected value inside @recurse() for key: %s", key)
+		if !it.Next() {
+			return it.Errorf("Expected argument")
 		}
-		val = item.Val
 
+		// consume the next item
+		item = it.Item()
+		val = item.Val
 		switch key {
 		case "depth":
-			depth, err := strconv.ParseUint(val, 0, 64)
-			if err != nil {
-				return err
+			// check whether the argument is variable or value.
+			if item.Typ == itemDollar {
+				// consume the variable name.
+				varName, err := parseVarName(it)
+				if err != nil {
+					return err
+				}
+				if gq.RecurseArgs.varMap == nil {
+					gq.RecurseArgs.varMap = make(map[string]string)
+				}
+				gq.RecurseArgs.varMap["depth"] = varName
+			} else {
+				if item.Typ != itemName {
+					return item.Errorf("Expected value inside @recurse() for key: %s", key)
+				}
+				depth, err := strconv.ParseUint(val, 0, 64)
+				if err != nil {
+					return err
+				}
+				gq.RecurseArgs.Depth = depth
 			}
-			gq.RecurseArgs.Depth = depth
 		case "loop":
-			allowLoop, err := strconv.ParseBool(val)
-			if err != nil {
-				return err
+			if item.Typ == itemDollar {
+				// consume the variable name.
+				varName, err := parseVarName(it)
+				if err != nil {
+					return err
+				}
+				if gq.RecurseArgs.varMap == nil {
+					gq.RecurseArgs.varMap = make(map[string]string)
+				}
+				gq.RecurseArgs.varMap["loop"] = varName
+			} else {
+				allowLoop, err := strconv.ParseBool(val)
+				if err != nil {
+					return err
+				}
+				gq.RecurseArgs.AllowLoop = allowLoop
 			}
-			gq.RecurseArgs.AllowLoop = allowLoop
 		default:
 			return item.Errorf("Unexpected key: [%s] inside @recurse block", key)
 		}
 
-		if _, ok := tryParseItemType(it, itemRightRound); ok {
+		if _, ok = tryParseItemType(it, itemRightRound); ok {
 			return nil
 		}
 

gql/parser_test.go

@@ -4923,3 +4923,61 @@ func TestParseVarAfterCountQry(t *testing.T) {
 	_, err := Parse(Request{Str: query})
 	require.NoError(t, err)
 }
+
+func TestRecurseWithArgs(t *testing.T) {
+	query := `
+	{
+		me(func: gt(count(~genre), 30000), first: 1) @recurse(depth: $hello, loop: true) {
+			name@en
+			~genre (first:10) @filter(gt(count(starring), 2))
+			starring (first: 2)
+			performance.actor
+		}
+	}`
+	gq, err := Parse(Request{Str: query, Variables: map[string]string{"$hello": "1"}})
+	require.NoError(t, err)
+	require.Equal(t, gq.Query[0].RecurseArgs.Depth, uint64(1))
+
+	query = `
+	{
+		me(func: gt(count(~genre), 30000), first: 1) @recurse(depth: 1, loop: $hello) {
+			name@en
+			~genre (first:10) @filter(gt(count(starring), 2))
+			starring (first: 2)
+			performance.actor
+		}
+	}`
+	gq, err = Parse(Request{Str: query, Variables: map[string]string{"$hello": "true"}})
+	require.NoError(t, err)
+	require.Equal(t, gq.Query[0].RecurseArgs.AllowLoop, true)
+
+	query = `
+	{
+		me(func: gt(count(~genre), 30000), first: 1) @recurse(depth: $hello, loop: $hello1) {
+			name@en
+			~genre (first:10) @filter(gt(count(starring), 2))
+			starring (first: 2)
+			performance.actor
+		}
+	}`
+	gq, err = Parse(Request{Str: query, Variables: map[string]string{"$hello": "1", "$hello1": "true"}})
+	require.NoError(t, err)
+	require.Equal(t, gq.Query[0].RecurseArgs.AllowLoop, true)
+	require.Equal(t, gq.Query[0].RecurseArgs.Depth, uint64(1))
+}
+
+func TestRecurse(t *testing.T) {
+	query := `
+	{
+		me(func: gt(count(~genre), 30000), first: 1) @recurse(depth: 1, loop: true) {
+			name@en
+			~genre (first:10) @filter(gt(count(starring), 2))
+			starring (first: 2)
+			performance.actor
+		}
+	}`
+	gq, err := Parse(Request{Str: query})
+	require.NoError(t, err)
+	require.Equal(t, gq.Query[0].RecurseArgs.Depth, uint64(1))
+	require.Equal(t, gq.Query[0].RecurseArgs.AllowLoop, true)
+}