Fix(GraphQL): Fix Upsert mutations in case of multiple upserts

graphql/admin/update_group.go

@@ -48,7 +48,7 @@ func (urw *updateGroupRewriter) Rewrite(
 		return nil, nil
 	}
 
-	upsertQuery := resolve.RewriteUpsertQueryFromMutation(m, nil, resolve.MutationQueryVar, "")
+	upsertQuery := resolve.RewriteUpsertQueryFromMutation(m, nil, resolve.MutationQueryVar, m.Name(), "")
 	srcUID := resolve.MutationQueryVarUID
 
 	var errSet, errDel error

graphql/e2e/common/mutation.go

@@ -5852,8 +5852,9 @@ func upsertMutationTests(t *testing.T) {
 	newCountry := addCountry(t, postExecutor)
 	// State should get added.
 	addStateParams := &GraphQLParams{
-		Query: `mutation addState($xcode: String!, $upsert: Boolean, $name: String!) {
-            addState(input: [{ xcode: $xcode, name: $name }], upsert: $upsert) {
+		Query: `mutation addState($xcode: String!, $upsert: Boolean, $name: String!, $xcode2: String!,
+					$name2: String!) {
+            addState(input: [{ xcode: $xcode, name: $name }, {xcode: $xcode2, name: $name2}], upsert: $upsert) {
                 state {
                     xcode
                     name
@@ -5866,6 +5867,8 @@ func upsertMutationTests(t *testing.T) {
 		Variables: map[string]interface{}{
 			"name":   "State1",
 			"xcode":  "S1",
+			"name2":  "State10",
+			"xcode2": "S10",
 			"upsert": true},
 	}
 
@@ -5874,18 +5877,26 @@ func upsertMutationTests(t *testing.T) {
 
 	addStateExpected := `{
         "addState": {
-            "state": [{
-                "xcode": "S1",
-                "name": "State1",
-				"country": null
-            }]
+            "state": [
+				{
+                	"xcode": "S1",
+                	"name": "State1",
+					"country": null
+            	},
+				{
+					"xcode": "S10",
+					"name": "State10",
+					"country": null
+				}]
         }
     }`
 	testutil.CompareJSON(t, addStateExpected, string(gqlResponse.Data))
 
 	// Add Mutation with Upsert: false should fail.
-	addStateParams.Query = `mutation addState($xcode: String!, $upsert: Boolean, $name: String!, $countryID: ID) {
-            addState(input: [{ xcode: $xcode, name: $name, country: {id: $countryID }}], upsert: $upsert) {
+	addStateParams.Query = `mutation addState($xcode: String!, $upsert: Boolean, $name: String!, $countryID: ID,
+				$xcode2: String!, $name2: String!) {
+            addState(input: [{ xcode: $xcode, name: $name, country: {id: $countryID }},
+							 { xcode: $xcode2, name: $name2}], upsert: $upsert) {
                 state {
                     xcode
                     name
@@ -5899,6 +5910,8 @@ func upsertMutationTests(t *testing.T) {
 		"upsert":    false,
 		"name":      "State2",
 		"xcode":     "S1",
+		"xcode2":    "S10",
+		"name2":     "NewState10",
 		"countryID": newCountry.ID,
 	}
 	gqlResponse = addStateParams.ExecuteAsPost(t, GraphqlURL)
@@ -5912,26 +5925,34 @@ func upsertMutationTests(t *testing.T) {
 		"upsert":    true,
 		"name":      "State2",
 		"xcode":     "S1",
+		"xcode2":    "S10",
+		"name2":     "NewState10",
 		"countryID": newCountry.ID,
 	}
 	gqlResponse = addStateParams.ExecuteAsPost(t, GraphqlURL)
 	RequireNoGQLErrors(t, gqlResponse)
 	addStateExpected = `{
         "addState": {
-            "state": [{
-                "xcode": "S1",
-                "name": "State2",
-				"country": {
-					"name": "Testland"
-				}
-            }]
+            "state": [
+				{
+                	"xcode": "S1",
+                	"name": "State2",
+					"country": {
+						"name": "Testland"
+					}
+            	},
+				{
+					"xcode": "S10",
+					"name": "NewState10",
+					"country": null
+				}]
         }
     }`
 	testutil.CompareJSON(t, addStateExpected, string(gqlResponse.Data))
 
 	// Clean Up
 	filter := map[string]interface{}{"id": []string{newCountry.ID}}
 	deleteCountry(t, filter, 1, nil)
-	filter = GetXidFilter("xcode", []interface{}{"S1"})
-	deleteState(t, filter, 1, nil)
+	filter = GetXidFilter("xcode", []interface{}{"S1", "S10"})
+	deleteState(t, filter, 2, nil)
 }

graphql/resolve/add_mutation_test.yaml

@@ -857,10 +857,10 @@
     }
   dgquerysec: |-
     query {
-      State1 as addState(func: uid(0x11)) @filter(type(State)) {
+      State1 as State1(func: uid(0x11)) @filter(type(State)) {
         uid
       }
-      State3 as addState(func: uid(0x13)) @filter(type(State)) {
+      State3 as State3(func: uid(0x13)) @filter(type(State)) {
         uid
       }
       var(func: uid(State1)) {
@@ -951,7 +951,7 @@
     }
   dgquerysec: |-
     query {
-      Book2 as addBook(func: uid(0x11)) @filter(type(Book)) {
+      Book2 as Book2(func: uid(0x11)) @filter(type(Book)) {
         uid
       }
     }
@@ -1000,7 +1000,7 @@
     }
   dgquerysec: |-
     query {
-      Book2 as addBook(func: uid(0x11)) @filter(type(Book)) {
+      Book2 as Book2(func: uid(0x11)) @filter(type(Book)) {
         uid
       }
     }
@@ -1058,7 +1058,7 @@
     }
   dgquerysec: |-
     query {
-      State1 as addState(func: uid(0x11)) @filter(type(State)) {
+      State1 as State1(func: uid(0x11)) @filter(type(State)) {
         uid
       }
       var(func: uid(State1)) {

graphql/resolve/auth_add_test.yaml

@@ -1251,7 +1251,7 @@
     }
   dgquerysec: |-
     query {
-      Tweets1 as addTweets(func: uid(TweetsRoot)) {
+      Tweets1 as Tweets1(func: uid(TweetsRoot)) {
         uid
       }
       TweetsRoot as var(func: uid(Tweets2))
@@ -1333,7 +1333,7 @@
     }
   dgquerysec: |-
     query {
-      State1 as addState(func: uid(StateRoot)) {
+      State1 as State1(func: uid(StateRoot)) {
         uid
       }
       StateRoot as var(func: uid(State3)) @filter(uid(StateAuth4))

graphql/resolve/mutation_rewriter.go

@@ -464,7 +464,7 @@ func (mrw *AddRewriter) Rewrite(
 			// State1 as addState(func: uid(0x11)) @filter(type(State)) {
 			// 		uid
 			// }
-			upsertQuery = append(upsertQuery, RewriteUpsertQueryFromMutation(m, authRw, upsertVar, idExistence[upsertVar])...)
+			upsertQuery = append(upsertQuery, RewriteUpsertQueryFromMutation(m, authRw, upsertVar, upsertVar, idExistence[upsertVar])...)
 			// Add upsert condition to ensure that the upsert takes place only when the node
 			// exists and has proper auth permission.
 			// Example condition:  cond: "@if(gt(len(State1), 0))"
@@ -585,7 +585,7 @@ func (urw *UpdateRewriter) Rewrite(
 	}
 	authRw.hasAuthRules = hasAuthRules(m.QueryField(), authRw)
 
-	upsertQuery := RewriteUpsertQueryFromMutation(m, authRw, MutationQueryVar, "")
+	upsertQuery := RewriteUpsertQueryFromMutation(m, authRw, MutationQueryVar, m.Name(), "")
 	srcUID := MutationQueryVarUID
 
 	if setArg == nil && delArg == nil {
@@ -719,6 +719,18 @@ func (mrw *AddRewriter) FromMutationResult(
 			fragment.(map[string]interface{})["uid"].(string), "_:")
 		val, ok := assigned[node]
 		if !ok {
+			// node was not part of assigned map. It is likely going to be part of Updated UIDs map.
+			// Extract and add any updated uids. This is done for upsert With Add Mutation.
+			// We extract out the variable name, eg. Project1 from uid(Project1)
+			uidVar := frag[0].fragment.(map[string]interface{})["uid"].(string)
+			uidVar = strings.TrimPrefix(uidVar, "uid(")
+			uidVar = strings.TrimSuffix(uidVar, ")")
+			updated := extractMutated(result, uidVar)
+			updateUids, err := extractUidsFromMutated(updated)
+			if err != nil {
+				return nil, err
+			}
+			uids = append(uids, updateUids...)
 			continue
 		}
 		uid, err := strconv.ParseUint(val, 0, 64)
@@ -732,24 +744,9 @@ func (mrw *AddRewriter) FromMutationResult(
 		uids = append(uids, uid)
 	}
 
-	// Extract and add any updated uids. This is done for upsert With Add Mutation.
-	// In this case, it may happen that no new node is created, but there may still
-	// be some updated nodes. We get these nodes over here and add to uid list.
-	mutated := extractMutated(result, mutation.Name())
-	if len(mutated) > 0 {
-		// This is the case of a conditional upsert where we should get uids from mutated.
-		for _, id := range mutated {
-			uid, err := strconv.ParseUint(id, 0, 64)
-			if err != nil {
-				return nil, schema.GQLWrapf(err,
-					"received %s as an updated uid from Dgraph, but couldn't parse it as "+
-						"uint64", id)
-			}
-			uids = append(uids, uid)
-		}
-	}
-
 	// Find out if its an upsert with Add mutation.
+	// In this case, it may happen that no new node is created, but there may still
+	// be some updated nodes. We don't throw an error in this case.
 	upsert := false
 	upsertVal := mutation.ArgValue(schema.UpsertArgName)
 	if upsertVal != nil {
@@ -797,18 +794,9 @@ func (urw *UpdateRewriter) FromMutationResult(
 
 	mutated := extractMutated(result, mutation.Name())
 
-	var uids []uint64
-	if len(mutated) > 0 {
-		// This is the case of a conditional upsert where we should get uids from mutated.
-		for _, id := range mutated {
-			uid, err := strconv.ParseUint(id, 0, 64)
-			if err != nil {
-				return nil, schema.GQLWrapf(err,
-					"received %s as an updated uid from Dgraph, but couldn't parse it as "+
-						"uint64", id)
-			}
-			uids = append(uids, uid)
-		}
+	uids, err := extractUidsFromMutated(mutated)
+	if err != nil {
+		return nil, err
 	}
 
 	customClaims, err := mutation.GetAuthMeta().ExtractCustomClaims(ctx)
@@ -838,10 +826,23 @@ func extractMutated(result map[string]interface{}, mutatedField string) []string
 			}
 		}
 	}
-
 	return mutated
 }
 
+func extractUidsFromMutated(mutated []string) ([]uint64, error) {
+	var ret []uint64
+	for _, id := range mutated {
+		uid, err := strconv.ParseUint(id, 0, 64)
+		if err != nil {
+			return ret, schema.GQLWrapf(err,
+				"received %s as an updated uid from Dgraph, but couldn't parse it as "+
+					"uint64", id)
+		}
+		ret = append(ret, uid)
+	}
+	return ret, nil
+}
+
 func addUpdateCondition(frags []*mutationFragment) {
 	for _, frag := range frags {
 		frag.conditions = append(frag.conditions, updateMutationCondition)
@@ -890,11 +891,12 @@ func RewriteUpsertQueryFromMutation(
 	m schema.Mutation,
 	authRw *authRewriter,
 	mutationQueryVar string,
+	queryAttribute string,
 	nodeID string) []*gql.GraphQuery {
 	// The query needs to assign the results to a variable, so that the mutation can use them.
 	dgQuery := []*gql.GraphQuery{{
 		Var:  mutationQueryVar,
-		Attr: m.Name(),
+		Attr: queryAttribute,
 	}}
 
 	rbac := authRw.evaluateStaticRules(m.MutatedType())
@@ -1021,7 +1023,7 @@ func (drw *deleteRewriter) Rewrite(
 	}
 	authRw.hasAuthRules = hasAuthRules(m.QueryField(), authRw)
 
-	dgQry := RewriteUpsertQueryFromMutation(m, authRw, MutationQueryVar, "")
+	dgQry := RewriteUpsertQueryFromMutation(m, authRw, MutationQueryVar, m.Name(), "")
 	qry := dgQry[0]
 
 	deletes := []interface{}{map[string]interface{}{"uid": "uid(x)"}}