Add support for multiple uids in uid_in function

gql/parser.go

@@ -1556,11 +1556,11 @@ loop:
 	return nil
 }
 
-// parseIneqArgs will try to parse the arguments inside an array ([]). If the values
+// parseFuncArgs will try to parse the arguments inside an array ([]). If the values
 // are prefixed with $ they are treated as Gql variables, otherwise they are used as scalar values.
 // Returns nil on success while appending arguments to the function Args slice. Otherwise
 // returns an error, which can be a parsing or value error.
-func parseIneqArgs(it *lex.ItemIterator, g *Function) error {
+func parseFuncArgs(it *lex.ItemIterator, g *Function) error {
 	var expectArg, isDollar bool
 
 	expectArg = true
@@ -1764,7 +1764,10 @@ L:
 					err = parseGeoArgs(it, function)
 
 				case IsInequalityFn(function.Name):
-					err = parseIneqArgs(it, function)
+					err = parseFuncArgs(it, function)
+
+				case function.Name == "uid_in":
+					err = parseFuncArgs(it, function)
 
 				default:
 					err = itemInFunc.Errorf("Unexpected character [ while parsing request.")

gql/parser_test.go

@@ -4869,6 +4869,41 @@ func TestParseGraphQLVarArray(t *testing.T) {
 	}
 }
 
+func TestParseGraphQLVarArrayUID_IN(t *testing.T) {
+	tests := []struct {
+		q    string
+		vars map[string]string
+		args int
+	}{
+		// uid_in test cases (uids and predicate inside uid_in are dummy)
+		{q: `query test($a: string){q(func: uid_in(director.film, [$a])) {name}}`,
+			vars: map[string]string{"$a": "0x4e472a"}, args: 1},
+		{q: `query test($a: string, $b: string){q(func: uid_in(director.film, [$a, $b])) {name}}`,
+			vars: map[string]string{"$a": "0x4e472a", "$b": "0x4e9545"}, args: 2},
+		{q: `query test($a: string){q(func: uid_in(name, [$a, "0x4e9545"])) {name}}`,
+			vars: map[string]string{"$a": "0x4e472a"}, args: 2},
+		{q: `query test($a: string){q(func: uid_in(name, ["0x4e9545", $a])) {name}}`,
+			vars: map[string]string{"$a": "0x4e472a"}, args: 2},
+	}
+	for _, tc := range tests {
+		gq, err := Parse(Request{Str: tc.q, Variables: tc.vars})
+		require.NoError(t, err)
+		require.Equal(t, 1, len(gq.Query))
+		require.Equal(t, "uid_in", gq.Query[0].Func.Name)
+		require.Equal(t, tc.args, len(gq.Query[0].Func.Args))
+		found := false
+		for _, val := range tc.vars {
+			for _, arg := range gq.Query[0].Func.Args {
+				if val == arg.Value {
+					found = true
+					break
+				}
+			}
+		}
+		require.True(t, found, "vars not matched: %v", tc.vars)
+	}
+}
+
 func TestParseGraphQLValueArray(t *testing.T) {
 	q := `
 	{

query/query1_test.go

@@ -1037,6 +1037,58 @@ func TestUidInFunction2(t *testing.T) {
 		js)
 }
 
+func TestUidInFunction3a(t *testing.T) {
+
+	query := `
+	{
+		me(func: UID(1, 23, 24)) @filter(uid_in(school, 5000, 5001)) {
+			name
+		}
+	}`
+	js := processQueryNoErr(t, query)
+	require.JSONEq(t, `{"data": {"me":[{"name":"Michonne"},{"name":"Rick Grimes"},{"name":"Glenn Rhee"}]}}`, js)
+}
+
+func TestUidInFunction3b(t *testing.T) {
+
+	query := `
+	{
+		me(func: UID(1, 23, 24)) @filter(uid_in(school, [5000, 5001])) {
+			name
+		}
+	}`
+	js := processQueryNoErr(t, query)
+	require.JSONEq(t, `{"data": {"me":[{"name":"Michonne"},{"name":"Rick Grimes"},{"name":"Glenn Rhee"}]}}`, js)
+}
+func TestUidInFunction4a(t *testing.T) {
+
+	query := `
+	{
+		me(func: uid(1, 23, 24 )) {
+			friend @filter(uid_in(school, 5000, 5001)) {
+				name
+			}
+		}
+	}`
+	js := processQueryNoErr(t, query)
+	require.JSONEq(t, `{"data": {"me":[{"friend":[{"name":"Rick Grimes"}, {"name":"Glenn Rhee"},{"name":"Daryl Dixon"},{"name":"Andrea"}]},{"friend":[{"name":"Michonne"}]}]}}`,
+		js)
+}
+
+func TestUidInFunction4b(t *testing.T) {
+
+	query := `
+	{
+		me(func: uid(1, 23, 24 )) {
+			friend @filter(uid_in(school, [5000, 5001])) {
+				name
+			}
+		}
+	}`
+	js := processQueryNoErr(t, query)
+	require.JSONEq(t, `{"data": {"me":[{"friend":[{"name":"Rick Grimes"}, {"name":"Glenn Rhee"},{"name":"Daryl Dixon"},{"name":"Andrea"}]},{"friend":[{"name":"Michonne"}]}]}}`,
+		js)
+}
 func TestUidInFunctionAtRoot(t *testing.T) {
 
 	query := `

worker/task.go

@@ -701,6 +701,7 @@ func (qs *queryState) handleUidPostings(
 			}
 			var key []byte
 			switch srcFn.fnType {
+			//Check with pawan/ashish if this needs any changes
 			case notAFunction, compareScalarFn, hasFn, uidInFn:
 				if q.Reverse {
 					key = x.ReverseKey(q.Attr, q.UidList.Uids[i])
@@ -757,11 +758,12 @@ func (qs *queryState) handleUidPostings(
 					tlist := &pb.List{Uids: []uint64{q.UidList.Uids[i]}}
 					out.UidMatrix = append(out.UidMatrix, tlist)
 				}
+			//Check with Pawan/Ashish what this function do? I suppose itreturns the relevant posting/uid list to append to uid matrix
 			case srcFn.fnType == uidInFn:
 				if i == 0 {
-					span.Annotate(nil, "UidInFn")
+					span.Annotate(nil, "UidInFn") //Check with Pawan/Ashish what does this do?
 				}
-				reqList := &pb.List{Uids: []uint64{srcFn.uidPresent}}
+				reqList := &pb.List{Uids: srcFn.uidsPresent}
 				topts := posting.ListOptions{
 					ReadTs:    args.q.ReadTs,
 					AfterUid:  0,
@@ -1583,7 +1585,7 @@ type functionContext struct {
 	ineqValueToken string
 	n              int
 	threshold      int64
-	uidPresent     uint64
+	uidsPresent    []uint64
 	fname          string
 	fnType         FuncType
 	regex          *cregexp.Regexp
@@ -1816,17 +1818,19 @@ func parseSrcFn(ctx context.Context, q *pb.Query) (*functionContext, error) {
 		}
 		checkRoot(q, fc)
 	case uidInFn:
-		if err = ensureArgsCount(q.SrcFunc, 1); err != nil {
-			return nil, err
-		}
-		fc.uidPresent, err = strconv.ParseUint(q.SrcFunc.Args[0], 0, 64)
-		if err != nil {
-			if e, ok := err.(*strconv.NumError); ok && e.Err == strconv.ErrSyntax {
-				return nil, errors.Errorf("Value %q in %s is not a number",
-					q.SrcFunc.Args[0], q.SrcFunc.Name)
+		//TODO (Anurag) Implement an argsCount for min 1 argument
+		for _, arg := range q.SrcFunc.Args {
+			uidParsed, err := strconv.ParseUint(arg, 0, 64)
+			if err != nil {
+				if e, ok := err.(*strconv.NumError); ok && e.Err == strconv.ErrSyntax {
+					return nil, errors.Errorf("Value %q in %s is not a number",
+						arg, q.SrcFunc.Name)
+				}
+				return nil, err
 			}
-			return nil, err
+			fc.uidsPresent = append(fc.uidsPresent, uidParsed)
 		}
+		sort.Slice(fc.uidsPresent, func(i, j int) bool { return fc.uidsPresent[i] < fc.uidsPresent[j] })
 		checkRoot(q, fc)
 		if fc.isFuncAtRoot {
 			return nil, errors.Errorf("uid_in function not allowed at root")