Add support for reading from and to arguments for shortest path query from variables.

gql/parser.go

@@ -53,21 +53,22 @@ type GraphQuery struct {
 
 	Args map[string]string
 	// Query can have multiple sort parameters.
-	Order        []*pb.Order
-	Children     []*GraphQuery
-	Filter       *FilterTree
-	MathExp      *MathTree
-	Normalize    bool
-	Recurse      bool
-	RecurseArgs  RecurseArgs
-	Cascade      bool
-	IgnoreReflex bool
-	Facets       *pb.FacetParams
-	FacetsFilter *FilterTree
-	GroupbyAttrs []GroupByAttr
-	FacetVar     map[string]string
-	FacetOrder   string
-	FacetDesc    bool
+	Order            []*pb.Order
+	Children         []*GraphQuery
+	Filter           *FilterTree
+	MathExp          *MathTree
+	Normalize        bool
+	Recurse          bool
+	RecurseArgs      RecurseArgs
+	ShortestPathArgs ShortestPathArgs
+	Cascade          bool
+	IgnoreReflex     bool
+	Facets           *pb.FacetParams
+	FacetsFilter     *FilterTree
+	GroupbyAttrs     []GroupByAttr
+	FacetVar         map[string]string
+	FacetOrder       string
+	FacetDesc        bool
 
 	// Internal fields below.
 	// If gq.fragment is nonempty, then it is a fragment reference / spread.
@@ -90,6 +91,16 @@ type RecurseArgs struct {
 	AllowLoop bool
 }
 
+// SHortestPathArgs stores the arguments needed to process the shortest path query.
+type ShortestPathArgs struct {
+	// From, To can have a uid or a uid function as the argument.
+	// 1. from: 0x01
+	// 2. from: uid(0x01)
+	// 3. from: uid(p) // a variable
+	From *Function
+	To   *Function
+}
+
 // GroupByAttr stores the arguments needed to process the @groupby directive.
 type GroupByAttr struct {
 	Attr  string
@@ -653,6 +664,15 @@ func (gq *GraphQuery) collectVars(v *Vars) {
 	if gq.MathExp != nil {
 		gq.MathExp.collectVars(v)
 	}
+
+	shortestPathFrom := gq.ShortestPathArgs.From
+	if shortestPathFrom != nil && len(shortestPathFrom.NeedsVar) > 0 {
+		v.Needs = append(v.Needs, shortestPathFrom.NeedsVar[0].Name)
+	}
+	shortestPathTo := gq.ShortestPathArgs.To
+	if shortestPathTo != nil && len(shortestPathTo.NeedsVar) > 0 {
+		v.Needs = append(v.Needs, shortestPathTo.NeedsVar[0].Name)
+	}
 }
 
 func (f *MathTree) collectVars(v *Vars) {
@@ -2318,6 +2338,11 @@ func attrAndLang(attrData string) (attr string, langs []string) {
 	return
 }
 
+func isEmpty(gq *GraphQuery) bool {
+	return gq.Func == nil && len(gq.NeedsVar) == 0 && len(gq.Args) == 0 &&
+		gq.ShortestPathArgs.From == nil && gq.ShortestPathArgs.To == nil
+}
+
 // getRoot gets the root graph query object after parsing the args.
 func getRoot(it *lex.ItemIterator) (gq *GraphQuery, rerr error) {
 	gq = &GraphQuery{
@@ -2365,7 +2390,7 @@ func getRoot(it *lex.ItemIterator) (gq *GraphQuery, rerr error) {
 			key = item.Val
 			expectArg = false
 		} else if item.Typ == itemRightRound {
-			if gq.Func == nil && len(gq.NeedsVar) == 0 && len(gq.Args) == 0 {
+			if isEmpty(gq) {
 				// Used to do aggregation at root which would be fetched in another block.
 				gq.IsEmpty = true
 			}
@@ -2392,7 +2417,8 @@ func getRoot(it *lex.ItemIterator) (gq *GraphQuery, rerr error) {
 			return nil, item.Errorf("Expecting a colon. Got: %v", item)
 		}
 
-		if key == "func" {
+		switch key {
+		case "func":
 			// Store the generator function.
 			if gq.Func != nil {
 				return gq, item.Errorf("Only one function allowed at root")
@@ -2406,7 +2432,55 @@ func getRoot(it *lex.ItemIterator) (gq *GraphQuery, rerr error) {
 			}
 			gq.Func = gen
 			gq.NeedsVar = append(gq.NeedsVar, gen.NeedsVar...)
-		} else {
+		case "from", "to":
+			if gq.Alias != "shortest" {
+				return gq, item.Errorf("from/to only allowed for shortest path queries")
+			}
+
+			fn := &Function{}
+			peekIt, err := it.Peek(1)
+			if err != nil {
+				return nil, item.Errorf("Invalid query")
+			}
+
+			assignShortestPathFn := func(fn *Function, key string) {
+				if key == "from" {
+					gq.ShortestPathArgs.From = fn
+				} else {
+					gq.ShortestPathArgs.To = fn
+				}
+			}
+
+			if peekIt[0].Val == uid {
+				gen, err := parseFunction(it, gq)
+				if err != nil {
+					return gq, err
+				}
+				fn.NeedsVar = gen.NeedsVar
+				fn.Name = gen.Name
+				assignShortestPathFn(fn, key)
+				continue
+			}
+
+			// This means it's not a uid function, so it has to be an actual uid.
+			it.Next()
+			item := it.Item()
+			val := collectName(it, item.Val)
+			uid, err := strconv.ParseUint(val, 0, 64)
+			switch e := err.(type) {
+			case nil:
+				fn.UID = append(fn.UID, uid)
+			case *strconv.NumError:
+				if e.Err == strconv.ErrRange {
+					return nil, item.Errorf("The uid value %q is too large.", val)
+				}
+				return nil,
+					item.Errorf("from/to in shortest path can only accept uid function or an uid. Got: %s",
+						val)
+			}
+			assignShortestPathFn(fn, key)
+
+		default:
 			var val string
 			if !it.Next() {
 				return nil, it.Errorf("Invalid query")

gql/parser_test.go

@@ -1071,13 +1071,47 @@ func TestParseShortestPath(t *testing.T) {
 	require.NoError(t, err)
 	require.NotNil(t, res.Query)
 	require.Equal(t, 1, len(res.Query))
-	require.Equal(t, "0x0a", res.Query[0].Args["from"])
-	require.Equal(t, "0x0b", res.Query[0].Args["to"])
+	require.Equal(t, uint64(0xa), res.Query[0].ShortestPathArgs.From.UID[0])
+	require.Equal(t, uint64(0xb), res.Query[0].ShortestPathArgs.To.UID[0])
 	require.Equal(t, "3", res.Query[0].Args["numpaths"])
 	require.Equal(t, "3", res.Query[0].Args["minweight"])
 	require.Equal(t, "6", res.Query[0].Args["maxweight"])
 }
 
+func TestParseShortestPathWithUidVars(t *testing.T) {
+	query := `{
+		a as var(func: uid(0x01))
+		b as var(func: uid(0x02))
+
+		shortest(from: uid(a), to: uid(b)) {
+			password
+			friend
+		}
+
+	}`
+	res, err := Parse(Request{Str: query})
+	require.NoError(t, err)
+	q := res.Query[2]
+	require.NotNil(t, q.ShortestPathArgs.From)
+	require.Equal(t, 1, len(q.ShortestPathArgs.From.NeedsVar))
+	require.Equal(t, "a", q.ShortestPathArgs.From.NeedsVar[0].Name)
+	require.Equal(t, "uid", q.ShortestPathArgs.From.Name)
+	require.NotNil(t, q.ShortestPathArgs.To)
+	require.Equal(t, 1, len(q.ShortestPathArgs.To.NeedsVar))
+}
+
+func TestParseShortestPathInvalidFnError(t *testing.T) {
+	query := `{
+		shortest(from: eq(a), to: uid(b)) {
+			password
+			friend
+		}
+
+	}`
+	_, err := Parse(Request{Str: query})
+	require.Error(t, err)
+}
+
 func TestParseMultipleQueries(t *testing.T) {
 	query := `
 	{

query/query.go

@@ -128,12 +128,14 @@ type params struct {
 	Langs      []string
 
 	// directives.
-	Normalize    bool
-	Recurse      bool
-	RecurseArgs  gql.RecurseArgs
-	Cascade      bool
-	IgnoreReflex bool
-
+	Normalize        bool
+	Recurse          bool
+	RecurseArgs      gql.RecurseArgs
+	ShortestPathArgs gql.ShortestPathArgs
+	Cascade          bool
+	IgnoreReflex     bool
+
+	// From and To are uids of the nodes between which we run the shortest path query.
 	From           uint64
 	To             uint64
 	Facet          *pb.FacetParams
@@ -858,22 +860,6 @@ func (args *params) fill(gq *gql.GraphQuery) error {
 			args.numPaths = int(numPaths)
 		}
 
-		if v, ok := gq.Args["from"]; ok {
-			from, err := strconv.ParseUint(v, 0, 64)
-			if err != nil {
-				return err
-			}
-			args.From = uint64(from)
-		}
-
-		if v, ok := gq.Args["to"]; ok {
-			to, err := strconv.ParseUint(v, 0, 64)
-			if err != nil {
-				return err
-			}
-			args.To = uint64(to)
-		}
-
 		if v, ok := gq.Args["maxweight"]; ok {
 			maxWeight, err := strconv.ParseFloat(v, 64)
 			if err != nil {
@@ -893,6 +879,16 @@ func (args *params) fill(gq *gql.GraphQuery) error {
 		} else if !ok {
 			args.MinWeight = -math.MaxFloat64
 		}
+
+		if gq.ShortestPathArgs.From == nil || gq.ShortestPathArgs.To == nil {
+			return errors.Errorf("from/to can't be nil for shortest path")
+		}
+		if len(gq.ShortestPathArgs.From.UID) > 0 {
+			args.From = gq.ShortestPathArgs.From.UID[0]
+		}
+		if len(gq.ShortestPathArgs.To.UID) > 0 {
+			args.To = gq.ShortestPathArgs.To.UID[0]
+		}
 	}
 
 	if v, ok := gq.Args["first"]; ok {
@@ -961,23 +957,24 @@ func newGraph(ctx context.Context, gq *gql.GraphQuery) (*SubGraph, error) {
 	// For the root, the name to be used in result is stored in Alias, not Attr.
 	// The attr at root (if present) would stand for the source functions attr.
 	args := params{
-		Alias:         gq.Alias,
-		Cascade:       gq.Cascade,
-		GetUid:        isDebug(ctx),
-		IgnoreReflex:  gq.IgnoreReflex,
-		IsEmpty:       gq.IsEmpty,
-		Langs:         gq.Langs,
-		NeedsVar:      append(gq.NeedsVar[:0:0], gq.NeedsVar...),
-		Normalize:     gq.Normalize,
-		Order:         gq.Order,
-		ParentVars:    make(map[string]varValue),
-		Recurse:       gq.Recurse,
-		RecurseArgs:   gq.RecurseArgs,
-		Var:           gq.Var,
-		groupbyAttrs:  gq.GroupbyAttrs,
-		isGroupBy:     gq.IsGroupby,
-		uidCount:      gq.UidCount,
-		uidCountAlias: gq.UidCountAlias,
+		Alias:            gq.Alias,
+		Cascade:          gq.Cascade,
+		GetUid:           isDebug(ctx),
+		IgnoreReflex:     gq.IgnoreReflex,
+		IsEmpty:          gq.IsEmpty,
+		Langs:            gq.Langs,
+		NeedsVar:         append(gq.NeedsVar[:0:0], gq.NeedsVar...),
+		Normalize:        gq.Normalize,
+		Order:            gq.Order,
+		ParentVars:       make(map[string]varValue),
+		Recurse:          gq.Recurse,
+		RecurseArgs:      gq.RecurseArgs,
+		ShortestPathArgs: gq.ShortestPathArgs,
+		Var:              gq.Var,
+		groupbyAttrs:     gq.GroupbyAttrs,
+		isGroupBy:        gq.IsGroupby,
+		uidCount:         gq.UidCount,
+		uidCountAlias:    gq.UidCountAlias,
 	}
 
 	for argk := range gq.Args {
@@ -1694,7 +1691,48 @@ func (sg *SubGraph) recursiveFillVars(doneVars map[string]varValue) error {
 	return nil
 }
 
+func (sg *SubGraph) fillShortestPathVars(mp map[string]varValue) error {
+	// The uidVar.Uids can be nil if the variable didn't return any uids. This would mean
+	// sg.Params.From or sg.Params.To is 0 and the query would return an empty result.
+	if sg.Params.ShortestPathArgs.From != nil && len(sg.Params.ShortestPathArgs.From.NeedsVar) > 0 {
+		fromVar := sg.Params.ShortestPathArgs.From.NeedsVar[0].Name
+		uidVar, ok := mp[fromVar]
+		if !ok {
+			return errors.Errorf("value of from var(%s) should have already been populated",
+				fromVar)
+		}
+		if uidVar.Uids != nil {
+			if len(uidVar.Uids.Uids) > 1 {
+				return errors.Errorf("from variable(%s) should only expand to 1 uid", fromVar)
+			}
+			sg.Params.From = uidVar.Uids.Uids[0]
+		}
+	}
+
+	if sg.Params.ShortestPathArgs.To != nil && len(sg.Params.ShortestPathArgs.To.NeedsVar) > 0 {
+		toVar := sg.Params.ShortestPathArgs.To.NeedsVar[0].Name
+		uidVar, ok := mp[toVar]
+		if !ok {
+			return errors.Errorf("value of to var(%s) should have already been populated",
+				toVar)
+		}
+		if uidVar.Uids != nil {
+			if len(uidVar.Uids.Uids) > 1 {
+				return errors.Errorf("to variable(%s) should only expand to 1 uid", toVar)
+			}
+			sg.Params.To = uidVar.Uids.Uids[0]
+		}
+	}
+	return nil
+}
+
 func (sg *SubGraph) fillVars(mp map[string]varValue) error {
+	if sg.Params.Alias == "shortest" {
+		if err := sg.fillShortestPathVars(mp); err != nil {
+			return err
+		}
+	}
+
 	var lists []*pb.List
 	for _, v := range sg.Params.NeedsVar {
 		if l, ok := mp[v.Name]; ok {