add support for variables in recurse

gql/parser.go

@@ -89,6 +89,8 @@ type GraphQuery struct {
 type RecurseArgs struct {
 	Depth     uint64
 	AllowLoop bool
+	varMap    map[string]string //varMap holds the variable args name. So, that we can substitute the
+	// argument in the substitution part.
 }
 
 // ShortestPathArgs stores the arguments needed to process the shortest path query.
@@ -408,6 +410,36 @@ func substituteVariables(gq *GraphQuery, vmap varMap) error {
 			return err
 		}
 	}
+	if gq.RecurseArgs.varMap != nil {
+		// Update the depth if the get the depth as a variable in the query.
+		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 errors.Wrapf(err, varName+" should be type of integer")
+			}
+			gq.RecurseArgs.Depth = depth
+		}
+
+		// Update the loop if the get the loop as a variable in the query.
+		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 errors.Wrapf(err, varName+"should be type of boolean")
+			}
+			gq.RecurseArgs.AllowLoop = allowLoop
+		}
+
+	}
 	return nil
 }
 
@@ -770,6 +802,31 @@ L2:
 	return gq, nil
 }
 
+// parseVarName returns the variable name.
+func parseVarName(it *lex.ItemIterator) (string, error) {
+	val := "$"
+	var consumeAtLeast bool
+	for {
+		items, err := it.Peek(1)
+		if err != nil {
+			return val, err
+		}
+		if items[0].Typ != itemName {
+			if !consumeAtLeast {
+				return "", it.Errorf("Expected variable name after $")
+			}
+			break
+		}
+		consumeAtLeast = true
+		val += items[0].Val
+		// Consume the current item.
+		if !it.Next() {
+			break
+		}
+	}
+	return val, nil
+}
+
 func parseRecurseArgs(it *lex.ItemIterator, gq *GraphQuery) error {
 	if ok := trySkipItemTyp(it, itemLeftRound); !ok {
 		// We don't have a (, we can return.
@@ -788,30 +845,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 errors.New("Value inside depth should be type of integer")
+				}
+				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 errors.New("Value inside loop should be type of boolean")
+				}
+				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

@@ -4924,6 +4924,115 @@ func TestParseVarAfterCountQry(t *testing.T) {
 	require.NoError(t, err)
 }
 
+func TestRecurseWithArgs(t *testing.T) {
+	query := `
+	{
+		me(func: eq(name, "sad")) @recurse(depth: $hello , loop: true) {
+		}
+	}`
+	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: eq(name, "sad"))@recurse(depth: 1 , loop: $hello) {
+		}
+	}`
+	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: eq(name, "sad"))@recurse(depth: $hello, loop: $hello1) {
+		}
+	}`
+	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))
+
+	query = `
+	{
+		me(func: eq(name, "sad"))@recurse(depth: $_hello_hello, loop: $hello1_heelo1) {
+		}
+	}`
+	gq, err = Parse(Request{Str: query, Variables: map[string]string{"$_hello_hello": "1",
+		"$hello1_heelo1": "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 TestRecurseWithArgsWithError(t *testing.T) {
+	query := `
+	{
+		me(func: eq(name, "sad"))@recurse(depth: $hello, loop: true) {
+		}
+	}`
+	_, err := Parse(Request{Str: query})
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "variable $hello not defined")
+
+	query = `
+	{
+		me(func: eq(name, "sad"))@recurse(depth: 1, loop: $hello) {
+		}
+	}`
+	_, err = Parse(Request{Str: query})
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "variable $hello not defined")
+
+	query = `
+	{
+		me(func: eq(name, "sad"))@recurse(depth: $hello, loop: $hello1) {
+		}
+	}`
+	_, err = Parse(Request{Str: query, Variables: map[string]string{"$hello": "sd", "$hello1": "true"}})
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "should be type of integer")
+
+	query = `
+	{
+		me(func: eq(name, "sad"))@recurse(depth: $hello, loop: $hello1) {
+		}
+	}`
+	_, err = Parse(Request{Str: query, Variables: map[string]string{"$hello": "1", "$hello1": "tre"}})
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "should be type of boolean")
+}
+
+func TestRecurse(t *testing.T) {
+	query := `
+	{
+		me(func: eq(name, "sad"))@recurse(depth: 1, loop: true) {
+		}
+	}`
+	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)
+}
+
+func TestRecurseWithError(t *testing.T) {
+	query := `
+	{
+		me(func: eq(name, "sad"))@recurse(depth: hello, loop: true) {
+		}
+	}`
+	_, err := Parse(Request{Str: query})
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "Value inside depth should be type of integer")
+	query = `
+	{
+		me(func: eq(name, "sad"))@recurse(depth: 1, loop: tre) {
+		}
+	}`
+	_, err = Parse(Request{Str: query})
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "Value inside loop should be type of boolean")
+}
 func TestParseExpandFilter(t *testing.T) {
 	query := `
 		{