Allow the result of count(uid) to be assigned to a var. (#2947)

Allow expressions of the form "f as count(uid)". When resolving the value of such a variable inside a math block it gets treated as a constant.
hypermodeinc · Jan 30, 2019 · 64c681a · 64c681a
1 parent 553d76d
commit 64c681a
Show file tree

Hide file tree

Showing 4 changed files with 76 additions and 14 deletions.
diff --git a/gql/parser.go b/gql/parser.go
@@ -2641,11 +2641,9 @@ func godeep(it *lex.ItemIterator, gq *GraphQuery) error {
 						goto Fall
 					}
 
-					if varName != "" {
-						return it.Errorf("Cannot assign variable to count()")
-					}
 					count = notSeen
 					gq.UidCount = true
+					gq.Var = varName
 					if alias != "" {
 						gq.UidCountAlias = alias
 					}

diff --git a/gql/parser_test.go b/gql/parser_test.go
@@ -3663,17 +3663,6 @@ func TestCountAtRootErr(t *testing.T) {
 }
 
 func TestCountAtRootErr2(t *testing.T) {
-	query := `{
-		me(func: uid( 1)) {
-			a as count(uid)
-		}
-	}`
-	_, err := Parse(Request{Str: query})
-	require.Error(t, err)
-	require.Contains(t, err.Error(), "Cannot assign variable to count()")
-}
-
-func TestCountAtRootErr3(t *testing.T) {
 	query := `{
 		me(func: uid( 1)) {
 			count()

diff --git a/query/query.go b/query/query.go
@@ -1227,6 +1227,14 @@ func (sg *SubGraph) transformVars(doneVars map[string]varValue, path []*SubGraph
 		if err != nil {
 			return err
 		}
+
+		// This is the result of setting the result of count(uid) to a variable.
+		// Treat this value as a constant.
+		if val, ok := newMap[math.MaxUint64]; ok && len(newMap) == 1 {
+			mt.Const = val
+			continue
+		}
+
 		mt.Val = newMap
 	}
 	return nil
@@ -1488,6 +1496,17 @@ func (sg *SubGraph) populateUidValVar(doneVars map[string]varValue, sgPath []*Su
 			}
 			doneVars[sg.Params.Var].Vals[uid] = val
 		}
+	} else if sg.Params.uidCount {
+		doneVars[sg.Params.Var] = varValue{
+			Vals: make(map[uint64]types.Val),
+			path: sgPath,
+		}
+
+		val := types.Val{
+			Tid:   types.IntID,
+			Value: int64(len(sg.DestUIDs.Uids)),
+		}
+		doneVars[sg.Params.Var].Vals[math.MaxUint64] = val
 	} else if len(sg.DestUIDs.Uids) != 0 || (sg.Attr == "uid" && sg.SrcUIDs != nil) {
 		// Uid variable could be defined using uid or a predicate.
 		uids := sg.DestUIDs

diff --git a/query/query0_test.go b/query/query0_test.go
@@ -1656,6 +1656,62 @@ func TestNestedFuncRoot4(t *testing.T) {
 	require.JSONEq(t, `{"data": {"me":[{"name":"Rick Grimes"},{"name":"Andrea"}]}}`, js)
 }
 
+func TestCountUidToVar(t *testing.T) {
+	query := `
+	{
+		var(func: has(school), first: 3) {
+			f as count(uid)
+		}
+
+		me(func: uid(1)) {
+			score: math(f)
+		}
+	}
+    `
+	js := processToFastJsonNoErr(t, query)
+	require.JSONEq(t, `{"data": {"me":[{"score": 3}]}}`, js)
+}
+
+func TestCountUidToVarMultiple(t *testing.T) {
+	query := `
+	{
+		var(func: has(school), first: 3) {
+			f as count(uid)
+		}
+
+		var(func: has(follow), first: 4) {
+			g as count(uid)
+		}
+
+		me(func: uid(1)) {
+			score: math(f + g)
+		}
+	}
+    `
+	js := processToFastJsonNoErr(t, query)
+	require.JSONEq(t, `{"data": {"me":[{"score": 7}]}}`, js)
+}
+
+func TestCountUidToVarCombinedWithNormalVar(t *testing.T) {
+	query := `
+	{
+		var(func: has(school), first: 3) {
+			f as count(uid)
+		}
+
+		var(func: has(follow)) {
+			g as count(path)
+		}
+
+		me(func: uid(1)) {
+			score: math(f + g)
+		}
+	}
+    `
+	js := processToFastJsonNoErr(t, query)
+	require.JSONEq(t, `{"data": {"me":[{"score": 5}]}}`, js)
+}
+
 var maxPendingCh chan uint64
 
 func TestMain(m *testing.M) {