From 4d3f8abadaa98c6bb079d7bf23f7d6682525a771 Mon Sep 17 00:00:00 2001
From: minhaj-shakeel <minhaj@dgraph.io>
Date: Wed, 10 Mar 2021 17:24:11 +0530
Subject: [PATCH 01/21] fix order of entities query result

---
 graphql/resolve/query.go          | 13 +++--
 graphql/resolve/query_rewriter.go | 67 +++++++++++++++----------
 graphql/resolve/resolver.go       | 83 ++++++++++++++++++++++++++++++-
 3 files changed, 133 insertions(+), 30 deletions(-)

diff --git a/graphql/resolve/query.go b/graphql/resolve/query.go
index 497c3e49e2b..c12079ec330 100644
--- a/graphql/resolve/query.go
+++ b/graphql/resolve/query.go
@@ -57,13 +57,19 @@ func (qr QueryResolverFunc) Resolve(ctx context.Context, query schema.Query) *Re
 // 1) rewrite the query using qr (return error if failed)
 // 2) execute the rewritten query with ex (return error if failed)
 func NewQueryResolver(qr QueryRewriter, ex DgraphExecutor) QueryResolver {
-	return &queryResolver{queryRewriter: qr, executor: ex}
+	return &queryResolver{queryRewriter: qr, executor: ex, resultCompleter: CompletionFunc(noopCompletion)}
+}
+
+// NewEntitiesQueryResolver ...
+func NewEntitiesQueryResolver(qr QueryRewriter, ex DgraphExecutor) QueryResolver {
+	return &queryResolver{queryRewriter: qr, executor: ex, resultCompleter: CompletionFunc(entitiesCompletion)}
 }
 
 // a queryResolver can resolve a single GraphQL query field.
 type queryResolver struct {
-	queryRewriter QueryRewriter
-	executor      DgraphExecutor
+	queryRewriter   QueryRewriter
+	executor        DgraphExecutor
+	resultCompleter ResultCompleter
 }
 
 func (qr *queryResolver) Resolve(ctx context.Context, query schema.Query) *Resolved {
@@ -82,6 +88,7 @@ func (qr *queryResolver) Resolve(ctx context.Context, query schema.Query) *Resol
 	defer timer.Stop()
 
 	resolved := qr.rewriteAndExecute(ctx, query)
+	qr.resultCompleter.Complete(ctx, resolved)
 	resolverTrace.Dgraph = resolved.Extensions.Tracing.Execution.Resolvers[0].Dgraph
 	resolved.Extensions.Tracing.Execution.Resolvers[0] = resolverTrace
 	return resolved
diff --git a/graphql/resolve/query_rewriter.go b/graphql/resolve/query_rewriter.go
index 68e6bb18625..b22410118ac 100644
--- a/graphql/resolve/query_rewriter.go
+++ b/graphql/resolve/query_rewriter.go
@@ -160,28 +160,11 @@ func (qr *queryRewriter) Rewrite(
 	}
 }
 
-// entitiesQuery rewrites the Apollo `_entities` Query which is sent from the Apollo gateway to a DQL query.
-// This query is sent to the Dgraph service to resolve types `extended` and defined by this service.
-func entitiesQuery(field schema.Query, authRw *authRewriter) ([]*gql.GraphQuery, error) {
-
-	// Input Argument to the Query is a List of "__typename" and "keyField" pair.
-	// For this type Extension:-
-	// 	extend type Product @key(fields: "upc") {
-	// 		upc: String @external
-	// 		reviews: [Review]
-	// 	}
-	// Input to the Query will be
-	// "_representations": [
-	// 		{
-	// 		  "__typename": "Product",
-	// 	 	 "upc": "B00005N5PF"
-	// 		},
-	// 		...
-	//   ]
-
+// parseReporesentationsArgument parses the "_representations" argument in the "_entities" query.
+func parseRepresentationsArgument(field schema.Query) (string, string, bool, []interface{}, error) {
 	representations, ok := field.ArgValue("representations").([]interface{})
 	if !ok {
-		return nil, fmt.Errorf("Error parsing `representations` argument")
+		return "", "", false, make([]interface{}, 0), fmt.Errorf("Error parsing `representations` argument")
 	}
 	typeNames := make(map[string]bool)
 	keyFieldValueList := make([]interface{}, 0)
@@ -191,30 +174,30 @@ func entitiesQuery(field schema.Query, authRw *authRewriter) ([]*gql.GraphQuery,
 	for i, rep := range representations {
 		representation, ok := rep.(map[string]interface{})
 		if !ok {
-			return nil, fmt.Errorf("Error parsing in %dth item in the `_representations` argument", i)
+			return "", "", false, make([]interface{}, 0), fmt.Errorf("Error parsing in %dth item in the `_representations` argument", i)
 		}
 
 		typename, ok := representation["__typename"].(string)
 		if !ok {
-			return nil, fmt.Errorf("Unable to extract __typename from %dth item in the `_representations` argument", i)
+			return "", "", false, make([]interface{}, 0), fmt.Errorf("Unable to extract __typename from %dth item in the `_representations` argument", i)
 		}
 
 		// Store all the typeNames into an map to perfrom validation at last.
 		typeNames[typename] = true
 		keyFieldName, keyFieldIsID, err = field.KeyField(typename)
 		if err != nil {
-			return nil, err
+			return "", "", false, make([]interface{}, 0), err
 		}
 		keyFieldValue, ok := representation[keyFieldName]
 		if !ok {
-			return nil, fmt.Errorf("Unable to extract value for key field `%s` from %dth item in the `_representations` argument", keyFieldName, i)
+			return "", "", false, make([]interface{}, 0), fmt.Errorf("Unable to extract value for key field `%s` from %dth item in the `_representations` argument", keyFieldName, i)
 		}
 		keyFieldValueList = append(keyFieldValueList, keyFieldValue)
 	}
 
 	// Return error if there was no typename extracted from the `_representations` argument.
 	if len(typeNames) == 0 {
-		return nil, fmt.Errorf("Expect one typename in `_representations` argument, got none")
+		return "", "", false, make([]interface{}, 0), fmt.Errorf("Expect one typename in `_representations` argument, got none")
 	}
 
 	// Since we have restricted that all the typeNames for the inputs in the
@@ -227,7 +210,7 @@ func entitiesQuery(field schema.Query, authRw *authRewriter) ([]*gql.GraphQuery,
 			keys[i] = k
 			i++
 		}
-		return nil, fmt.Errorf("Expected only one unique typename in `_representations` argument, got: %v", keys)
+		return "", "", false, make([]interface{}, 0), fmt.Errorf("Expected only one unique typename in `_representations` argument, got: %v", keys)
 	}
 
 	var typeName string
@@ -235,6 +218,33 @@ func entitiesQuery(field schema.Query, authRw *authRewriter) ([]*gql.GraphQuery,
 		typeName = k
 	}
 
+	return keyFieldName, typeName, keyFieldIsID, keyFieldValueList, nil
+}
+
+// entitiesQuery rewrites the Apollo `_entities` Query which is sent from the Apollo gateway to a DQL query.
+// This query is sent to the Dgraph service to resolve types `extended` and defined by this service.
+func entitiesQuery(field schema.Query, authRw *authRewriter) ([]*gql.GraphQuery, error) {
+
+	// Input Argument to the Query is a List of "__typename" and "keyField" pair.
+	// For this type Extension:-
+	// 	extend type Product @key(fields: "upc") {
+	// 		upc: String @external
+	// 		reviews: [Review]
+	// 	}
+	// Input to the Query will be
+	// "_representations": [
+	// 		{
+	// 		  "__typename": "Product",
+	// 	 	 "upc": "B00005N5PF"
+	// 		},
+	// 		...
+	//   ]
+
+	keyFieldName, typeName, keyFieldIsID, keyFieldValueList, err := parseRepresentationsArgument(field)
+	if err != nil {
+		return nil, err
+	}
+
 	typeDefn := field.BuildType(typeName)
 	rbac := authRw.evaluateStaticRules(typeDefn)
 
@@ -273,6 +283,11 @@ func entitiesQuery(field schema.Query, authRw *authRewriter) ([]*gql.GraphQuery,
 	// 	}
 	addTypeFilter(dgQuery, typeDefn)
 
+	// Add the  ascending Order of the keyField in the query.
+	// The result will be converted into the exact in the resultCompletion step.
+	dgQuery.Order = append(dgQuery.Order,
+		&pb.Order{Attr: typeDefn.DgraphPredicate(keyFieldName)})
+
 	selectionAuth := addSelectionSetFrom(dgQuery, field, authRw)
 	addUID(dgQuery)
 
diff --git a/graphql/resolve/resolver.go b/graphql/resolve/resolver.go
index 78a80a3d53a..a18df01c676 100644
--- a/graphql/resolve/resolver.go
+++ b/graphql/resolve/resolver.go
@@ -20,6 +20,7 @@ import (
 	"context"
 	"encoding/json"
 	"net/http"
+	"sort"
 	"strings"
 	"sync"
 	"time"
@@ -84,6 +85,11 @@ type ResolverFactory interface {
 	WithSchemaIntrospection() ResolverFactory
 }
 
+// A ResultCompleter ....
+type ResultCompleter interface {
+	Complete(ctx context.Context, resolved *Resolved)
+}
+
 // RequestResolver can process GraphQL requests and write GraphQL JSON responses.
 // A schema.Request may contain any number of queries or mutations (never both).
 // RequestResolver.Resolve() resolves all of them by finding the resolved answers
@@ -142,6 +148,15 @@ type Resolved struct {
 	Extensions *schema.Extensions
 }
 
+// CompletionFunc is an adapter that allows us to compose completions and build a
+// ResultCompleter from a function.  Based on the http.HandlerFunc pattern.
+type CompletionFunc func(ctx context.Context, resolved *Resolved)
+
+// Complete calls cf(ctx, resolved)
+func (cf CompletionFunc) Complete(ctx context.Context, resolved *Resolved) {
+	cf(ctx, resolved)
+}
+
 // NewDgraphExecutor builds a DgraphExecutor for proxying requests through dgraph.
 func NewDgraphExecutor() DgraphExecutor {
 	return newDgraphExecutor(&dgraph.DgraphEx{})
@@ -220,13 +235,18 @@ func (rf *resolverFactory) WithConventionResolvers(
 	queries := append(s.Queries(schema.GetQuery), s.Queries(schema.FilterQuery)...)
 	queries = append(queries, s.Queries(schema.PasswordQuery)...)
 	queries = append(queries, s.Queries(schema.AggregateQuery)...)
-	queries = append(queries, s.Queries(schema.EntitiesQuery)...)
 	for _, q := range queries {
 		rf.WithQueryResolver(q, func(q schema.Query) QueryResolver {
 			return NewQueryResolver(fns.Qrw, fns.Ex)
 		})
 	}
 
+	for _, q := range s.Queries(schema.EntitiesQuery) {
+		rf.WithQueryResolver(q, func(q schema.Query) QueryResolver {
+			return NewEntitiesQueryResolver(fns.Qrw, fns.Ex)
+		})
+	}
+
 	for _, q := range s.Queries(schema.HTTPQuery) {
 		rf.WithQueryResolver(q, func(q schema.Query) QueryResolver {
 			return NewHTTPQueryResolver(nil)
@@ -302,6 +322,67 @@ func NewResolverFactory(
 	}
 }
 
+// entitiesCompletion transform the result of the `_entities` query.
+// It changes the order of the result to the order of keyField in the
+// `_representations` argument.
+func entitiesCompletion(ctx context.Context, resolved *Resolved) {
+	// return id Data is not present
+	if len(resolved.Data) == 0 {
+		return
+	}
+
+	var data map[string][]interface{}
+	err := schema.Unmarshal(resolved.Data, &data)
+	if err != nil {
+		resolved.Err = schema.AppendGQLErrs(resolved.Err, err)
+		return
+	}
+
+	// fetch the keyFieldValueList from the query arguments.
+	_, _, _, keyFieldValueList, err := parseRepresentationsArgument(resolved.Field.(schema.Query))
+	if err != nil {
+		resolved.Err = schema.AppendGQLErrs(resolved.Err, err)
+		return
+	}
+
+	// store the index of the keyField Values present in the argument in a map.
+	indexMap := make(map[interface{}]int)
+	for i, key := range keyFieldValueList {
+		indexMap[key] = i
+	}
+
+	// sort the keyField Values in the ascending order.
+	sort.Slice(keyFieldValueList, func(i, j int) bool {
+		return keyFieldValueList[i].(string) < keyFieldValueList[j].(string)
+	})
+
+	// create the new output according to the index of the keyFields present in the argument.
+	output := make([]interface{}, len(keyFieldValueList))
+	for i, key := range keyFieldValueList {
+		output[indexMap[key]] = []interface{}(data["_entities"])[i]
+	}
+
+	// replace the result obtained from the dgraph and marshal back.
+	data["_entities"] = output
+	resolved.Data, err = json.Marshal(data)
+	if err != nil {
+		resolved.Err = schema.AppendGQLErrs(resolved.Err, err)
+	}
+
+}
+
+// noopCompletion just passes back it's result and err arguments
+func noopCompletion(ctx context.Context, resolved *Resolved) {}
+
+// StdQueryCompletion is the completion steps that get run for queries
+func StdQueryCompletion() CompletionFunc {
+	return noopCompletion
+}
+
+func entitiesQueryCompletion() CompletionFunc {
+	return entitiesCompletion
+}
+
 func (rf *resolverFactory) queryResolverFor(query schema.Query) QueryResolver {
 	rf.RLock()
 	defer rf.RUnlock()

From 1a4f5673d62c330b65d5c979f3ee2d60b58c6d06 Mon Sep 17 00:00:00 2001
From: minhaj-shakeel <minhaj@dgraph.io>
Date: Wed, 10 Mar 2021 17:29:44 +0530
Subject: [PATCH 02/21] remove stale func

---
 graphql/resolve/resolver.go | 5 -----
 1 file changed, 5 deletions(-)

diff --git a/graphql/resolve/resolver.go b/graphql/resolve/resolver.go
index a18df01c676..5ced9aad8bb 100644
--- a/graphql/resolve/resolver.go
+++ b/graphql/resolve/resolver.go
@@ -374,11 +374,6 @@ func entitiesCompletion(ctx context.Context, resolved *Resolved) {
 // noopCompletion just passes back it's result and err arguments
 func noopCompletion(ctx context.Context, resolved *Resolved) {}
 
-// StdQueryCompletion is the completion steps that get run for queries
-func StdQueryCompletion() CompletionFunc {
-	return noopCompletion
-}
-
 func entitiesQueryCompletion() CompletionFunc {
 	return entitiesCompletion
 }

From 634cb65979d7100922dffd912d85ee035abd23e1 Mon Sep 17 00:00:00 2001
From: minhaj-shakeel <minhaj@dgraph.io>
Date: Wed, 10 Mar 2021 17:59:11 +0530
Subject: [PATCH 03/21] add test case

---
 graphql/e2e/common/query.go | 59 +++++++++++++++++++++++++++----------
 1 file changed, 44 insertions(+), 15 deletions(-)

diff --git a/graphql/e2e/common/query.go b/graphql/e2e/common/query.go
index a6a3892790e..ffdb0c3d410 100644
--- a/graphql/e2e/common/query.go
+++ b/graphql/e2e/common/query.go
@@ -383,8 +383,8 @@ func allPosts(t *testing.T) []*post {
 
 func entitiesQuery(t *testing.T) {
 	addSpaceShipParams := &GraphQLParams{
-		Query: `mutation addSpaceShip($id1: String!, $missionId1: String! ) {
-			addSpaceShip(input: [{id: $id1, missions: [{id: $missionId1, designation: "Apollo1"}]} ]) {
+		Query: `mutation addSpaceShip($id1: String!, $id2: String!, $id3: String!, $id4: String! ) {
+			addSpaceShip(input: [{id: $id1, missions: [{id: "Mission1", designation: "Apollo1"}]},{id: $id2, missions: [{id: "Mission2", designation: "Apollo2"}]},{id: $id3, missions: [{id: "Mission3", designation: "Apollo3"}]}, {id: $id4, missions: [{id: "Mission4", designation: "Apollo4"}]}]){
 				spaceShip {
 					id
 					missions {
@@ -395,8 +395,10 @@ func entitiesQuery(t *testing.T) {
 			}
 		}`,
 		Variables: map[string]interface{}{
-			"id1":        "SpaceShip1",
-			"missionId1": "Mission1",
+			"id1": "SpaceShip1",
+			"id2": "SpaceShip2",
+			"id3": "SpaceShip3",
+			"id4": "SpaceShip4",
 		},
 	}
 
@@ -404,8 +406,8 @@ func entitiesQuery(t *testing.T) {
 	RequireNoGQLErrors(t, gqlResponse)
 
 	entitiesQueryParams := &GraphQLParams{
-		Query: `query _entities($typeName: String!, $id1: String!){
-			_entities(representations: [{__typename: $typeName, id: $id1}]) {
+		Query: `query _entities($typeName: String!, $id1: String!, $id2: String!, $id3: String!, $id4: String!){
+			_entities(representations: [{__typename: $typeName, id: $id4},{__typename: $typeName, id: $id2},{__typename: $typeName, id: $id1},{__typename: $typeName, id: $id3} ]) {
 				... on SpaceShip {
 					missions(order: {asc: id}){
 						id
@@ -417,6 +419,9 @@ func entitiesQuery(t *testing.T) {
 		Variables: map[string]interface{}{
 			"typeName": "SpaceShip",
 			"id1":      "SpaceShip1",
+			"id2":      "SpaceShip2",
+			"id3":      "SpaceShip3",
+			"id4":      "SpaceShip4",
 		},
 	}
 
@@ -425,15 +430,39 @@ func entitiesQuery(t *testing.T) {
 
 	expectedJSON := `{
 		"_entities": [
-		  {
-			"missions": [
-			  {
-				"id": "Mission1",
-				"designation": "Apollo1"
-			  }
-			]
-		  }
-		]
+      {
+        "missions": [
+          {
+            "designation": "Apollo4",
+            "id": "Mission4"
+          }
+        ]
+      },
+      {
+        "missions": [
+          {
+            "designation": "Apollo2",
+            "id": "Mission2"
+          }
+        ]
+      },
+      {
+        "missions": [
+          {
+            "designation": "Apollo1",
+            "id": "Mission1"
+          }
+        ]
+      },
+      {
+        "missions": [
+          {
+            "designation": "Apollo3",
+            "id": "Mission3"
+          }
+        ]
+      }
+    ]
 	  }`
 
 	testutil.CompareJSON(t, expectedJSON, string(entitiesResp.Data))

From 0a7e9f99ea779cb4049b41acad0c101538c14ba0 Mon Sep 17 00:00:00 2001
From: minhaj-shakeel <minhaj@dgraph.io>
Date: Wed, 10 Mar 2021 18:07:06 +0530
Subject: [PATCH 04/21] fix delete in e2e test

---
 graphql/e2e/common/query.go | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/graphql/e2e/common/query.go b/graphql/e2e/common/query.go
index ffdb0c3d410..a0ba8406c0b 100644
--- a/graphql/e2e/common/query.go
+++ b/graphql/e2e/common/query.go
@@ -467,10 +467,10 @@ func entitiesQuery(t *testing.T) {
 
 	testutil.CompareJSON(t, expectedJSON, string(entitiesResp.Data))
 
-	spaceShipDeleteFilter := map[string]interface{}{"id": map[string]interface{}{"in": []string{"SpaceShip1"}}}
+	spaceShipDeleteFilter := map[string]interface{}{"id": map[string]interface{}{"in": []string{"SpaceShip1", "SpaceShip2", "SpaceShip3", "SpaceShip4"}}}
 	DeleteGqlType(t, "SpaceShip", spaceShipDeleteFilter, 1, nil)
 
-	missionDeleteFilter := map[string]interface{}{"id": map[string]interface{}{"in": []string{"Mission1"}}}
+	missionDeleteFilter := map[string]interface{}{"id": map[string]interface{}{"in": []string{"Mission1", "Mission2", "Mission3", "Mission4"}}}
 	DeleteGqlType(t, "Mission", missionDeleteFilter, 1, nil)
 
 }

From 9b7d6bab21f62423028e6bb4f652754e50f18790 Mon Sep 17 00:00:00 2001
From: minhaj-shakeel <minhaj@dgraph.io>
Date: Wed, 10 Mar 2021 18:19:23 +0530
Subject: [PATCH 05/21] fix old entities query rewriting test

---
 graphql/resolve/auth_query_test.yaml | 12 ++++++------
 graphql/resolve/query_test.yaml      |  4 ++--
 2 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/graphql/resolve/auth_query_test.yaml b/graphql/resolve/auth_query_test.yaml
index ee3875c9479..32f7ef0c342 100644
--- a/graphql/resolve/auth_query_test.yaml
+++ b/graphql/resolve/auth_query_test.yaml
@@ -1887,14 +1887,14 @@
     USER: "user"
   dgquery: |-
     query {
-      _entities(func: uid(_EntityRoot)) {
+      _entities(func: uid(_EntityRoot), orderasc: Mission.id) {
         dgraph.type
         Mission.id : Mission.id
         Mission.designation : Mission.designation
         Mission.startDate : Mission.startDate
         dgraph.uid : uid
       }
-      _EntityRoot as var(func: uid(Mission1)) @filter(uid(MissionAuth2))
+      _EntityRoot as var(func: uid(Mission1), orderasc: Mission.id) @filter(uid(MissionAuth2))
       Mission1 as var(func: eq(Mission.id, "0x1", "0x2", "0x3")) @filter(type(Mission))
       MissionAuth2 as var(func: uid(Mission1)) @filter(eq(Mission.supervisorName, "user")) @cascade {
         Mission.id : Mission.id
@@ -1916,7 +1916,7 @@
     USER: "user"
   dgquery: |-
     query {
-      _entities(func: uid(_EntityRoot)) {
+      _entities(func: uid(_EntityRoot), orderasc: Astronaut.id) {
         dgraph.type
         Astronaut.missions : Astronaut.missions @filter(uid(Mission1)) {
           Mission.designation : Mission.designation
@@ -1924,7 +1924,7 @@
         }
         dgraph.uid : uid
       }
-      _EntityRoot as var(func: uid(Astronaut4))
+      _EntityRoot as var(func: uid(Astronaut4), orderasc: Astronaut.id)
       Astronaut4 as var(func: eq(Astronaut.id, "0x1", "0x2", "0x3")) @filter(type(Astronaut))
       var(func: uid(_EntityRoot)) {
         Mission2 as Astronaut.missions
@@ -1968,11 +1968,11 @@
     ROLE: "admin"
   dgquery: |-
     query {
-      _entities(func: uid(_EntityRoot)) {
+      _entities(func: uid(_EntityRoot), orderasc: Astronaut.id) {
         dgraph.type
         dgraph.uid : uid
       }
-      _EntityRoot as var(func: uid(Astronaut3))
+      _EntityRoot as var(func: uid(Astronaut3), orderasc: Astronaut.id)
       Astronaut3 as var(func: eq(Astronaut.id, "0x1", "0x2", "0x3")) @filter(type(Astronaut))
     }
 
diff --git a/graphql/resolve/query_test.yaml b/graphql/resolve/query_test.yaml
index 264ba9be54f..9d4a51e395e 100644
--- a/graphql/resolve/query_test.yaml
+++ b/graphql/resolve/query_test.yaml
@@ -3298,7 +3298,7 @@
     }
   dgquery:  |-
     query {
-      _entities(func: eq(Astronaut.id, "0x1", "0x2")) @filter(type(Astronaut)) {
+      _entities(func: eq(Astronaut.id, "0x1", "0x2"), orderasc: Astronaut.id) @filter(type(Astronaut)) {
         dgraph.type
         Astronaut.missions : Astronaut.missions {
           Mission.designation : Mission.designation
@@ -3322,7 +3322,7 @@
     }
   dgquery:  |-
     query {
-      _entities(func: eq(SpaceShip.id, "0x1", "0x2")) @filter(type(SpaceShip)) {
+      _entities(func: eq(SpaceShip.id, "0x1", "0x2"), orderasc: SpaceShip.id) @filter(type(SpaceShip)) {
         dgraph.type
         SpaceShip.missions : SpaceShip.missions {
           Mission.designation : Mission.designation

From 83fe1bfe98dd04bb2080b3fd83a3e0287a6caf99 Mon Sep 17 00:00:00 2001
From: minhaj-shakeel <minhaj@dgraph.io>
Date: Wed, 10 Mar 2021 18:47:51 +0530
Subject: [PATCH 06/21] fix deletion

---
 graphql/e2e/common/query.go | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/graphql/e2e/common/query.go b/graphql/e2e/common/query.go
index a0ba8406c0b..c0fdfc1b317 100644
--- a/graphql/e2e/common/query.go
+++ b/graphql/e2e/common/query.go
@@ -468,10 +468,10 @@ func entitiesQuery(t *testing.T) {
 	testutil.CompareJSON(t, expectedJSON, string(entitiesResp.Data))
 
 	spaceShipDeleteFilter := map[string]interface{}{"id": map[string]interface{}{"in": []string{"SpaceShip1", "SpaceShip2", "SpaceShip3", "SpaceShip4"}}}
-	DeleteGqlType(t, "SpaceShip", spaceShipDeleteFilter, 1, nil)
+	DeleteGqlType(t, "SpaceShip", spaceShipDeleteFilter, 4, nil)
 
 	missionDeleteFilter := map[string]interface{}{"id": map[string]interface{}{"in": []string{"Mission1", "Mission2", "Mission3", "Mission4"}}}
-	DeleteGqlType(t, "Mission", missionDeleteFilter, 1, nil)
+	DeleteGqlType(t, "Mission", missionDeleteFilter, 4, nil)
 
 }
 

From af0adef9a078797cf21446154d2664eab76cca95 Mon Sep 17 00:00:00 2001
From: minhaj-shakeel <minhaj@dgraph.io>
Date: Wed, 10 Mar 2021 20:41:09 +0530
Subject: [PATCH 07/21] address commebts

---
 graphql/e2e/common/query.go       |  3 +-
 graphql/resolve/query.go          |  6 ++--
 graphql/resolve/query_rewriter.go | 50 +++++++++++++++++++------------
 graphql/resolve/resolver.go       | 27 ++++++++---------
 4 files changed, 50 insertions(+), 36 deletions(-)

diff --git a/graphql/e2e/common/query.go b/graphql/e2e/common/query.go
index c0fdfc1b317..dc6d38e1808 100644
--- a/graphql/e2e/common/query.go
+++ b/graphql/e2e/common/query.go
@@ -37,6 +37,7 @@ import (
 
 	"github.com/google/go-cmp/cmp/cmpopts"
 
+	"github.com/dgraph-io/dgraph/graphql/e2e/common"
 	"github.com/dgraph-io/dgraph/graphql/schema"
 
 	"github.com/dgraph-io/dgraph/testutil"
@@ -465,7 +466,7 @@ func entitiesQuery(t *testing.T) {
     ]
 	  }`
 
-	testutil.CompareJSON(t, expectedJSON, string(entitiesResp.Data))
+	common.JSONEqGraphQL(t, expectedJSON, string(entitiesResp.Data))
 
 	spaceShipDeleteFilter := map[string]interface{}{"id": map[string]interface{}{"in": []string{"SpaceShip1", "SpaceShip2", "SpaceShip3", "SpaceShip4"}}}
 	DeleteGqlType(t, "SpaceShip", spaceShipDeleteFilter, 4, nil)
diff --git a/graphql/resolve/query.go b/graphql/resolve/query.go
index c12079ec330..505297b85ea 100644
--- a/graphql/resolve/query.go
+++ b/graphql/resolve/query.go
@@ -56,13 +56,15 @@ func (qr QueryResolverFunc) Resolve(ctx context.Context, query schema.Query) *Re
 // NewQueryResolver creates a new query resolver.  The resolver runs the pipeline:
 // 1) rewrite the query using qr (return error if failed)
 // 2) execute the rewritten query with ex (return error if failed)
+// 3) process the result with rc
 func NewQueryResolver(qr QueryRewriter, ex DgraphExecutor) QueryResolver {
 	return &queryResolver{queryRewriter: qr, executor: ex, resultCompleter: CompletionFunc(noopCompletion)}
 }
 
-// NewEntitiesQueryResolver ...
+// NewEntitiesQueryResolver creates a new query resolver for `_entities` query.
+// It is introduced because result completion works little different for `_entities` query.
 func NewEntitiesQueryResolver(qr QueryRewriter, ex DgraphExecutor) QueryResolver {
-	return &queryResolver{queryRewriter: qr, executor: ex, resultCompleter: CompletionFunc(entitiesCompletion)}
+	return &queryResolver{queryRewriter: qr, executor: ex, resultCompleter: CompletionFunc(entitiesQueryCompletion)}
 }
 
 // a queryResolver can resolve a single GraphQL query field.
diff --git a/graphql/resolve/query_rewriter.go b/graphql/resolve/query_rewriter.go
index b22410118ac..88506dc5315 100644
--- a/graphql/resolve/query_rewriter.go
+++ b/graphql/resolve/query_rewriter.go
@@ -160,11 +160,18 @@ func (qr *queryRewriter) Rewrite(
 	}
 }
 
+type parsedRepresentations struct {
+	keyFieldName      string
+	typeName          string
+	keyFieldIsID      bool
+	keyFieldValueList []interface{}
+}
+
 // parseReporesentationsArgument parses the "_representations" argument in the "_entities" query.
-func parseRepresentationsArgument(field schema.Query) (string, string, bool, []interface{}, error) {
+func parseRepresentationsArgument(field schema.Query) (*parsedRepresentations, error) {
 	representations, ok := field.ArgValue("representations").([]interface{})
 	if !ok {
-		return "", "", false, make([]interface{}, 0), fmt.Errorf("Error parsing `representations` argument")
+		return nil, fmt.Errorf("Error parsing `representations` argument")
 	}
 	typeNames := make(map[string]bool)
 	keyFieldValueList := make([]interface{}, 0)
@@ -174,30 +181,30 @@ func parseRepresentationsArgument(field schema.Query) (string, string, bool, []i
 	for i, rep := range representations {
 		representation, ok := rep.(map[string]interface{})
 		if !ok {
-			return "", "", false, make([]interface{}, 0), fmt.Errorf("Error parsing in %dth item in the `_representations` argument", i)
+			return nil, fmt.Errorf("Error parsing in %dth item in the `_representations` argument", i)
 		}
 
 		typename, ok := representation["__typename"].(string)
 		if !ok {
-			return "", "", false, make([]interface{}, 0), fmt.Errorf("Unable to extract __typename from %dth item in the `_representations` argument", i)
+			return nil, fmt.Errorf("Unable to extract __typename from %dth item in the `_representations` argument", i)
 		}
 
 		// Store all the typeNames into an map to perfrom validation at last.
 		typeNames[typename] = true
 		keyFieldName, keyFieldIsID, err = field.KeyField(typename)
 		if err != nil {
-			return "", "", false, make([]interface{}, 0), err
+			return nil, err
 		}
 		keyFieldValue, ok := representation[keyFieldName]
 		if !ok {
-			return "", "", false, make([]interface{}, 0), fmt.Errorf("Unable to extract value for key field `%s` from %dth item in the `_representations` argument", keyFieldName, i)
+			return nil, fmt.Errorf("Unable to extract value for key field `%s` from %dth item in the `_representations` argument", keyFieldName, i)
 		}
 		keyFieldValueList = append(keyFieldValueList, keyFieldValue)
 	}
 
 	// Return error if there was no typename extracted from the `_representations` argument.
 	if len(typeNames) == 0 {
-		return "", "", false, make([]interface{}, 0), fmt.Errorf("Expect one typename in `_representations` argument, got none")
+		return nil, fmt.Errorf("Expect one typename in `_representations` argument, got none")
 	}
 
 	// Since we have restricted that all the typeNames for the inputs in the
@@ -210,7 +217,7 @@ func parseRepresentationsArgument(field schema.Query) (string, string, bool, []i
 			keys[i] = k
 			i++
 		}
-		return "", "", false, make([]interface{}, 0), fmt.Errorf("Expected only one unique typename in `_representations` argument, got: %v", keys)
+		return nil, fmt.Errorf("Expected only one unique typename in `_representations` argument, got: %v", keys)
 	}
 
 	var typeName string
@@ -218,7 +225,11 @@ func parseRepresentationsArgument(field schema.Query) (string, string, bool, []i
 		typeName = k
 	}
 
-	return keyFieldName, typeName, keyFieldIsID, keyFieldValueList, nil
+	return &parsedRepresentations{
+		keyFieldName:      keyFieldName,
+		typeName:          typeName,
+		keyFieldIsID:      keyFieldIsID,
+		keyFieldValueList: keyFieldValueList}, nil
 }
 
 // entitiesQuery rewrites the Apollo `_entities` Query which is sent from the Apollo gateway to a DQL query.
@@ -240,12 +251,12 @@ func entitiesQuery(field schema.Query, authRw *authRewriter) ([]*gql.GraphQuery,
 	// 		...
 	//   ]
 
-	keyFieldName, typeName, keyFieldIsID, keyFieldValueList, err := parseRepresentationsArgument(field)
+	parsedRepr, err := parseRepresentationsArgument(field)
 	if err != nil {
 		return nil, err
 	}
 
-	typeDefn := field.BuildType(typeName)
+	typeDefn := field.BuildType(parsedRepr.typeName)
 	rbac := authRw.evaluateStaticRules(typeDefn)
 
 	dgQuery := &gql.GraphQuery{
@@ -271,10 +282,16 @@ func entitiesQuery(field schema.Query, authRw *authRewriter) ([]*gql.GraphQuery,
 	// If the key field is of ID type and is not an external field
 	// then we query it using the `uid` otherwise we treat it as string
 	// and query using `eq` function.
-	if keyFieldIsID && !typeDefn.Field(keyFieldName).IsExternal() {
-		addUIDFunc(dgQuery, convertIDs(keyFieldValueList))
+	// We also don't need to add Order to the query as the results are
+	// automatically returned in the ascending order of the uids.
+	if parsedRepr.keyFieldIsID && !typeDefn.Field(parsedRepr.keyFieldName).IsExternal() {
+		addUIDFunc(dgQuery, convertIDs(parsedRepr.keyFieldValueList))
 	} else {
-		addEqFunc(dgQuery, typeDefn.DgraphPredicate(keyFieldName), keyFieldValueList)
+		addEqFunc(dgQuery, typeDefn.DgraphPredicate(parsedRepr.keyFieldName), parsedRepr.keyFieldValueList)
+		// Add the  ascending Order of the keyField in the query.
+		// The result will be converted into the exact in the resultCompletion step.
+		dgQuery.Order = append(dgQuery.Order,
+			&pb.Order{Attr: typeDefn.DgraphPredicate(parsedRepr.keyFieldName)})
 	}
 	// AddTypeFilter in as the Filter to the Root the Query.
 	// Query will be like :-
@@ -283,11 +300,6 @@ func entitiesQuery(field schema.Query, authRw *authRewriter) ([]*gql.GraphQuery,
 	// 	}
 	addTypeFilter(dgQuery, typeDefn)
 
-	// Add the  ascending Order of the keyField in the query.
-	// The result will be converted into the exact in the resultCompletion step.
-	dgQuery.Order = append(dgQuery.Order,
-		&pb.Order{Attr: typeDefn.DgraphPredicate(keyFieldName)})
-
 	selectionAuth := addSelectionSetFrom(dgQuery, field, authRw)
 	addUID(dgQuery)
 
diff --git a/graphql/resolve/resolver.go b/graphql/resolve/resolver.go
index 5ced9aad8bb..ed0a6ec439f 100644
--- a/graphql/resolve/resolver.go
+++ b/graphql/resolve/resolver.go
@@ -85,7 +85,8 @@ type ResolverFactory interface {
 	WithSchemaIntrospection() ResolverFactory
 }
 
-// A ResultCompleter ....
+// A ResultCompleter can take a []byte slice representing an intermediate result
+// in resolving field and applies a completion step.
 type ResultCompleter interface {
 	Complete(ctx context.Context, resolved *Resolved)
 }
@@ -325,8 +326,8 @@ func NewResolverFactory(
 // entitiesCompletion transform the result of the `_entities` query.
 // It changes the order of the result to the order of keyField in the
 // `_representations` argument.
-func entitiesCompletion(ctx context.Context, resolved *Resolved) {
-	// return id Data is not present
+func entitiesQueryCompletion(ctx context.Context, resolved *Resolved) {
+	// return if Data is not present
 	if len(resolved.Data) == 0 {
 		return
 	}
@@ -339,27 +340,29 @@ func entitiesCompletion(ctx context.Context, resolved *Resolved) {
 	}
 
 	// fetch the keyFieldValueList from the query arguments.
-	_, _, _, keyFieldValueList, err := parseRepresentationsArgument(resolved.Field.(schema.Query))
+	repr, err := parseRepresentationsArgument(resolved.Field.(schema.Query))
 	if err != nil {
 		resolved.Err = schema.AppendGQLErrs(resolved.Err, err)
 		return
 	}
 
 	// store the index of the keyField Values present in the argument in a map.
+	// key in the map as there are multiple types like String, Int, Int64 allowed as @id.
 	indexMap := make(map[interface{}]int)
-	for i, key := range keyFieldValueList {
+	for i, key := range repr.keyFieldValueList {
 		indexMap[key] = i
 	}
 
 	// sort the keyField Values in the ascending order.
-	sort.Slice(keyFieldValueList, func(i, j int) bool {
-		return keyFieldValueList[i].(string) < keyFieldValueList[j].(string)
+	sort.Slice(repr.keyFieldValueList, func(i, j int) bool {
+		return repr.keyFieldValueList[i].(string) < repr.keyFieldValueList[j].(string)
 	})
 
 	// create the new output according to the index of the keyFields present in the argument.
-	output := make([]interface{}, len(keyFieldValueList))
-	for i, key := range keyFieldValueList {
-		output[indexMap[key]] = []interface{}(data["_entities"])[i]
+	output := make([]interface{}, len(repr.keyFieldValueList))
+	entitiesQryResp := data["_entities"]
+	for i, key := range repr.keyFieldValueList {
+		output[indexMap[key]] = entitiesQryResp[i]
 	}
 
 	// replace the result obtained from the dgraph and marshal back.
@@ -374,10 +377,6 @@ func entitiesCompletion(ctx context.Context, resolved *Resolved) {
 // noopCompletion just passes back it's result and err arguments
 func noopCompletion(ctx context.Context, resolved *Resolved) {}
 
-func entitiesQueryCompletion() CompletionFunc {
-	return entitiesCompletion
-}
-
 func (rf *resolverFactory) queryResolverFor(query schema.Query) QueryResolver {
 	rf.RLock()
 	defer rf.RUnlock()

From e9593154e2d31a4d6f97835eea47cae40d1c6653 Mon Sep 17 00:00:00 2001
From: minhaj-shakeel <minhaj@dgraph.io>
Date: Wed, 10 Mar 2021 20:44:20 +0530
Subject: [PATCH 08/21] fix JSON eq

---
 graphql/e2e/common/query.go | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/graphql/e2e/common/query.go b/graphql/e2e/common/query.go
index dc6d38e1808..38d350599c8 100644
--- a/graphql/e2e/common/query.go
+++ b/graphql/e2e/common/query.go
@@ -37,7 +37,6 @@ import (
 
 	"github.com/google/go-cmp/cmp/cmpopts"
 
-	"github.com/dgraph-io/dgraph/graphql/e2e/common"
 	"github.com/dgraph-io/dgraph/graphql/schema"
 
 	"github.com/dgraph-io/dgraph/testutil"
@@ -466,7 +465,7 @@ func entitiesQuery(t *testing.T) {
     ]
 	  }`
 
-	common.JSONEqGraphQL(t, expectedJSON, string(entitiesResp.Data))
+	JSONEqGraphQL(t, expectedJSON, string(entitiesResp.Data))
 
 	spaceShipDeleteFilter := map[string]interface{}{"id": map[string]interface{}{"in": []string{"SpaceShip1", "SpaceShip2", "SpaceShip3", "SpaceShip4"}}}
 	DeleteGqlType(t, "SpaceShip", spaceShipDeleteFilter, 4, nil)

From 4ea6765e25a3649bf04aa6829724cfe2ff3c8014 Mon Sep 17 00:00:00 2001
From: minhaj-shakeel <minhaj@dgraph.io>
Date: Thu, 11 Mar 2021 15:51:46 +0530
Subject: [PATCH 09/21] fix JSON eq

---
 graphql/e2e/common/query.go | 37 +------------------------------------
 graphql/resolve/resolver.go |  3 ++-
 2 files changed, 3 insertions(+), 37 deletions(-)

diff --git a/graphql/e2e/common/query.go b/graphql/e2e/common/query.go
index 38d350599c8..d9963bb4f33 100644
--- a/graphql/e2e/common/query.go
+++ b/graphql/e2e/common/query.go
@@ -428,42 +428,7 @@ func entitiesQuery(t *testing.T) {
 	entitiesResp := entitiesQueryParams.ExecuteAsPost(t, GraphqlURL)
 	RequireNoGQLErrors(t, entitiesResp)
 
-	expectedJSON := `{
-		"_entities": [
-      {
-        "missions": [
-          {
-            "designation": "Apollo4",
-            "id": "Mission4"
-          }
-        ]
-      },
-      {
-        "missions": [
-          {
-            "designation": "Apollo2",
-            "id": "Mission2"
-          }
-        ]
-      },
-      {
-        "missions": [
-          {
-            "designation": "Apollo1",
-            "id": "Mission1"
-          }
-        ]
-      },
-      {
-        "missions": [
-          {
-            "designation": "Apollo3",
-            "id": "Mission3"
-          }
-        ]
-      }
-    ]
-	  }`
+	expectedJSON := `{"_entities":[{"missions":[{"designation":"Apollo4","id":"Mission4"}]},{"missions":[{"designation":"Apollo2","id":"Mission2"}]},{"missions":[{"designation":"Apollo1","id":"Mission1"}]},{"missions":[{"designation":"Apollo3","id":"Mission3"}]}]}`
 
 	JSONEqGraphQL(t, expectedJSON, string(entitiesResp.Data))
 
diff --git a/graphql/resolve/resolver.go b/graphql/resolve/resolver.go
index ed0a6ec439f..6a1e1ad42f4 100644
--- a/graphql/resolve/resolver.go
+++ b/graphql/resolve/resolver.go
@@ -361,7 +361,8 @@ func entitiesQueryCompletion(ctx context.Context, resolved *Resolved) {
 	// create the new output according to the index of the keyFields present in the argument.
 	output := make([]interface{}, len(repr.keyFieldValueList))
 	entitiesQryResp := data["_entities"]
-	for i, key := range repr.keyFieldValueList {
+	for i := 0; i < len(entitiesQryResp); i++ {
+		key := repr.keyFieldValueList[i]
 		output[indexMap[key]] = entitiesQryResp[i]
 	}
 

From 5bae8256690463a5cdf118a8c03edf9267639947 Mon Sep 17 00:00:00 2001
From: minhaj-shakeel <minhaj@dgraph.io>
Date: Thu, 11 Mar 2021 17:08:30 +0530
Subject: [PATCH 10/21] handle duplicate keys

---
 graphql/resolve/resolver.go | 51 +++++++++++++++++++++++++++++--------
 1 file changed, 40 insertions(+), 11 deletions(-)

diff --git a/graphql/resolve/resolver.go b/graphql/resolve/resolver.go
index 6a1e1ad42f4..1d40b55f26c 100644
--- a/graphql/resolve/resolver.go
+++ b/graphql/resolve/resolver.go
@@ -348,22 +348,51 @@ func entitiesQueryCompletion(ctx context.Context, resolved *Resolved) {
 
 	// store the index of the keyField Values present in the argument in a map.
 	// key in the map as there are multiple types like String, Int, Int64 allowed as @id.
-	indexMap := make(map[interface{}]int)
+	// There could be duplicate keys in the representations so the value of map is a list
+	// of integers containing all the indices for a key.
+	indexMap := make(map[interface{}][]int)
+	uniqueKeyList := make([]interface{}, 0)
 	for i, key := range repr.keyFieldValueList {
-		indexMap[key] = i
-	}
-
-	// sort the keyField Values in the ascending order.
-	sort.Slice(repr.keyFieldValueList, func(i, j int) bool {
-		return repr.keyFieldValueList[i].(string) < repr.keyFieldValueList[j].(string)
+		indexMap[key] = append(indexMap[key], i)
+	}
+
+	// Create a list containing unique keys and then sort in ascending order because this
+	// will be the order in which the data is received.
+	// for eg: for keys: {1, 2, 4, 1, 3} is converted into {1, 2, 4, 3} and then {1, 2, 3, 4}
+	// this will be the order of received data from the dgraph.
+	for k := range indexMap {
+		uniqueKeyList = append(uniqueKeyList, k)
+	}
+	sort.Slice(uniqueKeyList, func(i, j int) bool {
+		switch uniqueKeyList[i].(type) {
+		case int64:
+			return uniqueKeyList[i].(int64) < uniqueKeyList[j].(int64)
+		case float64:
+			return uniqueKeyList[i].(float64) < uniqueKeyList[j].(float64)
+		case string:
+			return uniqueKeyList[i].(string) < uniqueKeyList[j].(string)
+		case json.Number:
+			return uniqueKeyList[i].(json.Number) < uniqueKeyList[j].(json.Number)
+		}
+		return false
 	})
 
 	// create the new output according to the index of the keyFields present in the argument.
-	output := make([]interface{}, len(repr.keyFieldValueList))
 	entitiesQryResp := data["_entities"]
-	for i := 0; i < len(entitiesQryResp); i++ {
-		key := repr.keyFieldValueList[i]
-		output[indexMap[key]] = entitiesQryResp[i]
+
+	// if `entitiesQueryResp` contains less number of elements than the number of unique keys
+	// which is because the object related to certain key is not present in the dgraph.
+	// This will end into an error at the Gateway, so no need to order the result here.
+	if len(entitiesQryResp) < len(uniqueKeyList) {
+		return
+	}
+
+	// Reorder the output response according to the order of the keys in the representations argument.
+	output := make([]interface{}, len(repr.keyFieldValueList))
+	for i, key := range uniqueKeyList {
+		for _, idx := range indexMap[key] {
+			output[idx] = entitiesQryResp[i]
+		}
 	}
 
 	// replace the result obtained from the dgraph and marshal back.

From 304d35b19034b831ebac3241f27e78c1a275b13b Mon Sep 17 00:00:00 2001
From: minhaj-shakeel <minhaj@dgraph.io>
Date: Thu, 11 Mar 2021 17:13:49 +0530
Subject: [PATCH 11/21] add duplicate keys

---
 graphql/e2e/common/query.go | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/graphql/e2e/common/query.go b/graphql/e2e/common/query.go
index d9963bb4f33..8caaf75e4e1 100644
--- a/graphql/e2e/common/query.go
+++ b/graphql/e2e/common/query.go
@@ -407,7 +407,7 @@ func entitiesQuery(t *testing.T) {
 
 	entitiesQueryParams := &GraphQLParams{
 		Query: `query _entities($typeName: String!, $id1: String!, $id2: String!, $id3: String!, $id4: String!){
-			_entities(representations: [{__typename: $typeName, id: $id4},{__typename: $typeName, id: $id2},{__typename: $typeName, id: $id1},{__typename: $typeName, id: $id3} ]) {
+			_entities(representations: [{__typename: $typeName, id: $id4},{__typename: $typeName, id: $id2},{__typename: $typeName, id: $id1},{__typename: $typeName, id: $id3},{__typename: $typeName, id: $id1}]) {
 				... on SpaceShip {
 					missions(order: {asc: id}){
 						id
@@ -428,7 +428,7 @@ func entitiesQuery(t *testing.T) {
 	entitiesResp := entitiesQueryParams.ExecuteAsPost(t, GraphqlURL)
 	RequireNoGQLErrors(t, entitiesResp)
 
-	expectedJSON := `{"_entities":[{"missions":[{"designation":"Apollo4","id":"Mission4"}]},{"missions":[{"designation":"Apollo2","id":"Mission2"}]},{"missions":[{"designation":"Apollo1","id":"Mission1"}]},{"missions":[{"designation":"Apollo3","id":"Mission3"}]}]}`
+	expectedJSON := `{"_entities":[{"missions":[{"designation":"Apollo4","id":"Mission4"}]},{"missions":[{"designation":"Apollo2","id":"Mission2"}]},{"missions":[{"designation":"Apollo1","id":"Mission1"}]},{"missions":[{"designation":"Apollo3","id":"Mission3"}]},{"missions":[{"designation":"Apollo1","id":"Mission1"}]}]}`
 
 	JSONEqGraphQL(t, expectedJSON, string(entitiesResp.Data))
 

From c9448db4d31d9de0682aea399ef22a63bc83c6c8 Mon Sep 17 00:00:00 2001
From: minhaj-shakeel <minhaj@dgraph.io>
Date: Thu, 11 Mar 2021 17:38:06 +0530
Subject: [PATCH 12/21] add test case for entities query with keyfield of type
 Int

---
 graphql/e2e/common/common.go      |  3 +-
 graphql/e2e/common/query.go       | 61 ++++++++++++++++++++++++++++++-
 graphql/e2e/normal/schema.graphql |  5 +++
 3 files changed, 67 insertions(+), 2 deletions(-)

diff --git a/graphql/e2e/common/common.go b/graphql/e2e/common/common.go
index 5d8668af9d4..40f5d6d2de7 100644
--- a/graphql/e2e/common/common.go
+++ b/graphql/e2e/common/common.go
@@ -797,7 +797,8 @@ func RunAll(t *testing.T) {
 	t.Run("query only typename", queryOnlyTypename)
 	t.Run("query nested only typename", querynestedOnlyTypename)
 	t.Run("test onlytypename for interface types", onlytypenameForInterface)
-	t.Run("entitites Query on extended type", entitiesQuery)
+	t.Run("entitites Query on extended type with key field of type String", entitiesQueryWithKeyFieldOfTypeString)
+	t.Run("entitites Query on extended type with key field of type Int", entitiesQueryWithKeyFieldOfTypeInt)
 
 	t.Run("get state by xid", getStateByXid)
 	t.Run("get state without args", getStateWithoutArgs)
diff --git a/graphql/e2e/common/query.go b/graphql/e2e/common/query.go
index 8caaf75e4e1..f20caab7a2a 100644
--- a/graphql/e2e/common/query.go
+++ b/graphql/e2e/common/query.go
@@ -381,7 +381,7 @@ func allPosts(t *testing.T) []*post {
 	return result.QueryPost
 }
 
-func entitiesQuery(t *testing.T) {
+func entitiesQueryWithKeyFieldOfTypeString(t *testing.T) {
 	addSpaceShipParams := &GraphQLParams{
 		Query: `mutation addSpaceShip($id1: String!, $id2: String!, $id3: String!, $id4: String! ) {
 			addSpaceShip(input: [{id: $id1, missions: [{id: "Mission1", designation: "Apollo1"}]},{id: $id2, missions: [{id: "Mission2", designation: "Apollo2"}]},{id: $id3, missions: [{id: "Mission3", designation: "Apollo3"}]}, {id: $id4, missions: [{id: "Mission4", designation: "Apollo4"}]}]){
@@ -440,6 +440,65 @@ func entitiesQuery(t *testing.T) {
 
 }
 
+func entitiesQueryWithKeyFieldOfTypeInt(t *testing.T) {
+	addPlanetParams := &GraphQLParams{
+		Query: `mutation addPlanet($id1: String!, $id2: String!, $id3: String!, $id4: String! ) {
+			addPlanet(input: [{id: $id1, missions: [{id: "Mission1", designation: "Apollo1"}]},{id: $id2, missions: [{id: "Mission2", designation: "Apollo2"}]},{id: $id3, missions: [{id: "Mission3", designation: "Apollo3"}]}, {id: $id4, missions: [{id: "Mission4", designation: "Apollo4"}]}]){
+				planet {
+					id
+					missions {
+						id
+						designation
+					}
+				}
+			}
+		}`,
+		Variables: map[string]interface{}{
+			"id1": "1",
+			"id2": "2",
+			"id3": "3",
+			"id4": "4",
+		},
+	}
+
+	gqlResponse := addPlanetParams.ExecuteAsPost(t, GraphqlURL)
+	RequireNoGQLErrors(t, gqlResponse)
+
+	entitiesQueryParams := &GraphQLParams{
+		Query: `query _entities($typeName: String!, $id1: String!, $id2: String!, $id3: String!, $id4: String!){
+			_entities(representations: [{__typename: $typeName, id: $id4},{__typename: $typeName, id: $id2},{__typename: $typeName, id: $id1},{__typename: $typeName, id: $id3},{__typename: $typeName, id: $id1}]) {
+				... on Planet {
+					missions(order: {asc: id}){
+						id
+						designation
+					}
+				}
+			}
+		}`,
+		Variables: map[string]interface{}{
+			"typeName": "Planet",
+			"id1":      1,
+			"id2":      2,
+			"id3":      3,
+			"id4":      4,
+		},
+	}
+
+	entitiesResp := entitiesQueryParams.ExecuteAsPost(t, GraphqlURL)
+	RequireNoGQLErrors(t, entitiesResp)
+
+	expectedJSON := `{"_entities":[{"missions":[{"designation":"Apollo4","id":"Mission4"}]},{"missions":[{"designation":"Apollo2","id":"Mission2"}]},{"missions":[{"designation":"Apollo1","id":"Mission1"}]},{"missions":[{"designation":"Apollo3","id":"Mission3"}]},{"missions":[{"designation":"Apollo1","id":"Mission1"}]}]}`
+
+	JSONEqGraphQL(t, expectedJSON, string(entitiesResp.Data))
+
+	planetDeleteFilter := map[string]interface{}{"id": map[string]interface{}{"in": []int{1, 2, 3, 4}}}
+	DeleteGqlType(t, "Planet", planetDeleteFilter, 4, nil)
+
+	missionDeleteFilter := map[string]interface{}{"id": map[string]interface{}{"in": []string{"Mission1", "Mission2", "Mission3", "Mission4"}}}
+	DeleteGqlType(t, "Mission", missionDeleteFilter, 4, nil)
+
+}
+
 func inFilterOnString(t *testing.T) {
 	addStateParams := &GraphQLParams{
 		Query: `mutation addState($name1: String!, $code1: String!, $name2: String!, $code2: String! ) {
diff --git a/graphql/e2e/normal/schema.graphql b/graphql/e2e/normal/schema.graphql
index 46b40225a99..521712dbe5d 100644
--- a/graphql/e2e/normal/schema.graphql
+++ b/graphql/e2e/normal/schema.graphql
@@ -345,6 +345,11 @@ type SpaceShip @key(fields: "id") @extends {
     missions: [Mission]
 }
 
+type Planet @key(fields: "id") @extends {
+        id: Int! @id @external
+        missions: [Mission]
+}
+
 type Region {
         id: String! @id
         name: String!

From ee3e899840ae33566f69158b8bb1a78d94741669 Mon Sep 17 00:00:00 2001
From: minhaj-shakeel <minhaj@dgraph.io>
Date: Thu, 11 Mar 2021 18:14:51 +0530
Subject: [PATCH 13/21] fix test case for int

---
 graphql/e2e/common/query.go             | 10 ++--------
 graphql/e2e/normal/schema_response.json | 11 +++++++++++
 2 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/graphql/e2e/common/query.go b/graphql/e2e/common/query.go
index f20caab7a2a..bf06906a6bd 100644
--- a/graphql/e2e/common/query.go
+++ b/graphql/e2e/common/query.go
@@ -442,8 +442,8 @@ func entitiesQueryWithKeyFieldOfTypeString(t *testing.T) {
 
 func entitiesQueryWithKeyFieldOfTypeInt(t *testing.T) {
 	addPlanetParams := &GraphQLParams{
-		Query: `mutation addPlanet($id1: String!, $id2: String!, $id3: String!, $id4: String! ) {
-			addPlanet(input: [{id: $id1, missions: [{id: "Mission1", designation: "Apollo1"}]},{id: $id2, missions: [{id: "Mission2", designation: "Apollo2"}]},{id: $id3, missions: [{id: "Mission3", designation: "Apollo3"}]}, {id: $id4, missions: [{id: "Mission4", designation: "Apollo4"}]}]){
+		Query: `mutation {
+			addPlanet(input: [{id: 1, missions: [{id: "Mission1", designation: "Apollo1"}]},{id: 2, missions: [{id: "Mission2", designation: "Apollo2"}]},{id: 3, missions: [{id: "Mission3", designation: "Apollo3"}]}, {id: 4, missions: [{id: "Mission4", designation: "Apollo4"}]}]){
 				planet {
 					id
 					missions {
@@ -453,12 +453,6 @@ func entitiesQueryWithKeyFieldOfTypeInt(t *testing.T) {
 				}
 			}
 		}`,
-		Variables: map[string]interface{}{
-			"id1": "1",
-			"id2": "2",
-			"id3": "3",
-			"id4": "4",
-		},
 	}
 
 	gqlResponse := addPlanetParams.ExecuteAsPost(t, GraphqlURL)
diff --git a/graphql/e2e/normal/schema_response.json b/graphql/e2e/normal/schema_response.json
index 490d2604cf2..cdd3bb4d54a 100644
--- a/graphql/e2e/normal/schema_response.json
+++ b/graphql/e2e/normal/schema_response.json
@@ -1359,6 +1359,17 @@
       ],
       "name": "Zoo"
     },
+    {
+      "fields": [
+        {
+          "name": "Planet.missions"
+        },
+        {
+          "name": "Planet.id"
+        }
+      ],
+      "name": "Planet"
+    },
     {
       "fields": [
         {

From 5953106ae0955b9b1d6eed460a1751e43815db5b Mon Sep 17 00:00:00 2001
From: minhaj-shakeel <minhaj@dgraph.io>
Date: Thu, 11 Mar 2021 18:23:20 +0530
Subject: [PATCH 14/21] fix test case for int

---
 graphql/e2e/common/query.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/graphql/e2e/common/query.go b/graphql/e2e/common/query.go
index bf06906a6bd..5ebbd443793 100644
--- a/graphql/e2e/common/query.go
+++ b/graphql/e2e/common/query.go
@@ -459,7 +459,7 @@ func entitiesQueryWithKeyFieldOfTypeInt(t *testing.T) {
 	RequireNoGQLErrors(t, gqlResponse)
 
 	entitiesQueryParams := &GraphQLParams{
-		Query: `query _entities($typeName: String!, $id1: String!, $id2: String!, $id3: String!, $id4: String!){
+		Query: `query _entities($typeName: String!, $id1: Int!, $id2: Int!, $id3: Int!, $id4: Int!){
 			_entities(representations: [{__typename: $typeName, id: $id4},{__typename: $typeName, id: $id2},{__typename: $typeName, id: $id1},{__typename: $typeName, id: $id3},{__typename: $typeName, id: $id1}]) {
 				... on Planet {
 					missions(order: {asc: id}){

From 3cf3dab337692e8bb6528681927bd74adde8609f Mon Sep 17 00:00:00 2001
From: minhaj-shakeel <minhaj@dgraph.io>
Date: Thu, 11 Mar 2021 18:27:32 +0530
Subject: [PATCH 15/21] fix response test

---
 graphql/e2e/normal/schema_response.json | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/graphql/e2e/normal/schema_response.json b/graphql/e2e/normal/schema_response.json
index cdd3bb4d54a..0daa296e919 100644
--- a/graphql/e2e/normal/schema_response.json
+++ b/graphql/e2e/normal/schema_response.json
@@ -593,6 +593,20 @@
       "type": "uid",
       "list": true
     },
+    {
+      "predicate": "Planet.id",
+      "type": "int",
+      "index": true,
+      "tokenizer": [
+        "hash"
+      ],
+      "upsert": true
+    },
+    {
+      "predicate": "Planet.missions",
+      "type": "uid",
+      "list": true
+    },
     {
       "predicate": "Starship.length",
       "type": "float"

From 91a4d8f184e8d03ca410ce24bbeb957052bc3192 Mon Sep 17 00:00:00 2001
From: minhaj-shakeel <minhaj@dgraph.io>
Date: Thu, 11 Mar 2021 20:11:53 +0530
Subject: [PATCH 16/21] fix directive e2e test

---
 graphql/e2e/directives/schema.graphql | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/graphql/e2e/directives/schema.graphql b/graphql/e2e/directives/schema.graphql
index 7f06c40c698..8904c99edf3 100644
--- a/graphql/e2e/directives/schema.graphql
+++ b/graphql/e2e/directives/schema.graphql
@@ -327,6 +327,11 @@ type SpaceShip @key(fields: "id") @extends {
     missions: [Mission]
 }
 
+type Planet @key(fields: "id") @extends {
+        id: Int! @id @external
+        missions: [Mission]
+}
+
 type Region {
     id: String! @id
     name: String!

From 794628270f19070376ad6243c9f0fef582cd9f50 Mon Sep 17 00:00:00 2001
From: minhaj-shakeel <minhaj@dgraph.io>
Date: Thu, 11 Mar 2021 21:06:57 +0530
Subject: [PATCH 17/21] fix comparison for json.Number

---
 graphql/e2e/directives/schema_response.json | 25 ++++++++++++++++++++
 graphql/resolve/resolver.go                 | 26 +++++++++++++++------
 2 files changed, 44 insertions(+), 7 deletions(-)

diff --git a/graphql/e2e/directives/schema_response.json b/graphql/e2e/directives/schema_response.json
index 94f3e648411..3a39cf16b4a 100644
--- a/graphql/e2e/directives/schema_response.json
+++ b/graphql/e2e/directives/schema_response.json
@@ -432,6 +432,20 @@
       "type": "uid",
       "list": true
     },
+    {
+      "predicate": "Planet.id",
+      "type": "int",
+      "index": true,
+      "tokenizer": [
+        "hash"
+      ],
+      "upsert": true
+    },
+    {
+      "predicate": "Planet.missions",
+      "type": "uid",
+      "list": true
+    },
     {
       "predicate": "State.capital",
       "type": "string"
@@ -1137,6 +1151,17 @@
       ],
       "name": "SpaceShip"
     },
+    {
+      "fields": [
+        {
+          "name": "Planet.id"
+        },
+        {
+          "name": "Planet.missions"
+        }
+      ],
+      "name": "Planet"
+    },
     {
       "fields": [
         {
diff --git a/graphql/resolve/resolver.go b/graphql/resolve/resolver.go
index 1d40b55f26c..d406e5d8ffb 100644
--- a/graphql/resolve/resolver.go
+++ b/graphql/resolve/resolver.go
@@ -346,10 +346,13 @@ func entitiesQueryCompletion(ctx context.Context, resolved *Resolved) {
 		return
 	}
 
+	typeDefn := resolved.Field.(schema.Query).BuildType(repr.typeName)
+	keyFieldType := typeDefn.Field(repr.keyFieldName).Type().Name()
+
 	// store the index of the keyField Values present in the argument in a map.
-	// key in the map as there are multiple types like String, Int, Int64 allowed as @id.
-	// There could be duplicate keys in the representations so the value of map is a list
-	// of integers containing all the indices for a key.
+	// key in the map is of type interface because there are multiple types like String,
+	// Int, Int64 allowed as @id. There could be duplicate keys in the representations
+	// so the value of map is a list of integers containing all the indices for a key.
 	indexMap := make(map[interface{}][]int)
 	uniqueKeyList := make([]interface{}, 0)
 	for i, key := range repr.keyFieldValueList {
@@ -365,14 +368,23 @@ func entitiesQueryCompletion(ctx context.Context, resolved *Resolved) {
 	}
 	sort.Slice(uniqueKeyList, func(i, j int) bool {
 		switch uniqueKeyList[i].(type) {
+		case string:
+			return uniqueKeyList[i].(string) < uniqueKeyList[j].(string)
+		case json.Number:
+			switch keyFieldType {
+			case "Int":
+				val1, _ := uniqueKeyList[i].(json.Number).Int64()
+				val2, _ := uniqueKeyList[j].(json.Number).Int64()
+				return val1 < val2
+			case "Float":
+				val1, _ := uniqueKeyList[i].(json.Number).Float64()
+				val2, _ := uniqueKeyList[j].(json.Number).Float64()
+				return val1 < val2
+			}
 		case int64:
 			return uniqueKeyList[i].(int64) < uniqueKeyList[j].(int64)
 		case float64:
 			return uniqueKeyList[i].(float64) < uniqueKeyList[j].(float64)
-		case string:
-			return uniqueKeyList[i].(string) < uniqueKeyList[j].(string)
-		case json.Number:
-			return uniqueKeyList[i].(json.Number) < uniqueKeyList[j].(json.Number)
 		}
 		return false
 	})

From 8d84991eb44aae1eed867733a105e398effb42fb Mon Sep 17 00:00:00 2001
From: minhaj-shakeel <minhaj@dgraph.io>
Date: Thu, 11 Mar 2021 21:25:45 +0530
Subject: [PATCH 18/21] fix tokenizer

---
 graphql/e2e/directives/schema_response.json | 2 +-
 graphql/e2e/normal/schema_response.json     | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/graphql/e2e/directives/schema_response.json b/graphql/e2e/directives/schema_response.json
index 3a39cf16b4a..b5e67d865a6 100644
--- a/graphql/e2e/directives/schema_response.json
+++ b/graphql/e2e/directives/schema_response.json
@@ -437,7 +437,7 @@
       "type": "int",
       "index": true,
       "tokenizer": [
-        "hash"
+        "int"
       ],
       "upsert": true
     },
diff --git a/graphql/e2e/normal/schema_response.json b/graphql/e2e/normal/schema_response.json
index 0daa296e919..443ff83b7df 100644
--- a/graphql/e2e/normal/schema_response.json
+++ b/graphql/e2e/normal/schema_response.json
@@ -598,7 +598,7 @@
       "type": "int",
       "index": true,
       "tokenizer": [
-        "hash"
+        "int"
       ],
       "upsert": true
     },

From 35465acff990c24e2d23f4119349f1350d86ad3b Mon Sep 17 00:00:00 2001
From: minhaj-shakeel <minhaj@dgraph.io>
Date: Thu, 11 Mar 2021 23:34:33 +0530
Subject: [PATCH 19/21] optimize explicit typecasting

---
 graphql/resolve/resolver.go | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/graphql/resolve/resolver.go b/graphql/resolve/resolver.go
index d406e5d8ffb..2f44ddaf8ae 100644
--- a/graphql/resolve/resolver.go
+++ b/graphql/resolve/resolver.go
@@ -367,24 +367,24 @@ func entitiesQueryCompletion(ctx context.Context, resolved *Resolved) {
 		uniqueKeyList = append(uniqueKeyList, k)
 	}
 	sort.Slice(uniqueKeyList, func(i, j int) bool {
-		switch uniqueKeyList[i].(type) {
+		switch val := uniqueKeyList[i].(type) {
 		case string:
-			return uniqueKeyList[i].(string) < uniqueKeyList[j].(string)
+			return val < uniqueKeyList[j].(string)
 		case json.Number:
 			switch keyFieldType {
 			case "Int":
-				val1, _ := uniqueKeyList[i].(json.Number).Int64()
+				val1, _ := val.Int64()
 				val2, _ := uniqueKeyList[j].(json.Number).Int64()
 				return val1 < val2
 			case "Float":
-				val1, _ := uniqueKeyList[i].(json.Number).Float64()
+				val1, _ := val.Float64()
 				val2, _ := uniqueKeyList[j].(json.Number).Float64()
 				return val1 < val2
 			}
 		case int64:
-			return uniqueKeyList[i].(int64) < uniqueKeyList[j].(int64)
+			return val < uniqueKeyList[j].(int64)
 		case float64:
-			return uniqueKeyList[i].(float64) < uniqueKeyList[j].(float64)
+			return val < uniqueKeyList[j].(float64)
 		}
 		return false
 	})

From 2da1d8a4948c9ce8492009b3bcf698d77b44bb33 Mon Sep 17 00:00:00 2001
From: minhaj-shakeel <minhaj@dgraph.io>
Date: Fri, 12 Mar 2021 11:50:12 +0530
Subject: [PATCH 20/21] address comment

---
 graphql/resolve/resolver.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/graphql/resolve/resolver.go b/graphql/resolve/resolver.go
index 2f44ddaf8ae..2836e1aea13 100644
--- a/graphql/resolve/resolver.go
+++ b/graphql/resolve/resolver.go
@@ -372,7 +372,7 @@ func entitiesQueryCompletion(ctx context.Context, resolved *Resolved) {
 			return val < uniqueKeyList[j].(string)
 		case json.Number:
 			switch keyFieldType {
-			case "Int":
+			case "Int", "Int64":
 				val1, _ := val.Int64()
 				val2, _ := uniqueKeyList[j].(json.Number).Int64()
 				return val1 < val2

From 76dd73418f6a5f4ef87bc8ab0b631fe6af29c90f Mon Sep 17 00:00:00 2001
From: minhaj-shakeel <minhaj@dgraph.io>
Date: Fri, 12 Mar 2021 12:10:06 +0530
Subject: [PATCH 21/21] fix auth test

---
 graphql/resolve/auth_query_test.yaml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/graphql/resolve/auth_query_test.yaml b/graphql/resolve/auth_query_test.yaml
index 4d787a667fd..65be3d5837c 100644
--- a/graphql/resolve/auth_query_test.yaml
+++ b/graphql/resolve/auth_query_test.yaml
@@ -1897,7 +1897,7 @@
       _EntityRoot as var(func: uid(Mission_1), orderasc: Mission.id) @filter(uid(Mission_Auth2))
       Mission_1 as var(func: eq(Mission.id, "0x1", "0x2", "0x3")) @filter(type(Mission))
       Mission_Auth2 as var(func: uid(Mission_1)) @filter(eq(Mission.supervisorName, "user")) @cascade {
-      Mission.id : Mission.id
+        Mission.id : Mission.id
       }
     }
 - name: "Entities query with top level RBAC rule true and level 1 query auth rule"