graphql: Add extensions

graphql/admin/admin.go

@@ -480,7 +480,7 @@ func upsertEmptyGQLSchema() (*gqlSchema, error) {
 		},
 	}
 
-	assigned, result, err := resolve.AdminMutationExecutor().Mutate(context.Background(), qry,
+	assigned, result, _, err := resolve.AdminMutationExecutor().Mutate(context.Background(), qry,
 		mutations)
 	if err != nil {
 		return nil, err

graphql/admin/schema.go

@@ -96,22 +96,24 @@ func (asr *updateSchemaResolver) FromMutationResult(
 func (asr *updateSchemaResolver) Mutate(
 	ctx context.Context,
 	query *gql.GraphQuery,
-	mutations []*dgoapi.Mutation) (map[string]string, map[string]interface{}, error) {
-	assigned, result, err := asr.baseMutationExecutor.Mutate(ctx, query, mutations)
+	mutations []*dgoapi.Mutation) (map[string]string, map[string]interface{}, *schema.Extensions,
+	error) {
+	assigned, result, ext, err := asr.baseMutationExecutor.Mutate(ctx, query, mutations)
 	if err != nil {
-		return nil, nil, err
+		return nil, nil, ext, err
 	}
 
 	_, err = (&edgraph.Server{}).Alter(ctx, &dgoapi.Operation{Schema: asr.newDgraphSchema})
 	if err != nil {
-		return nil, nil, schema.GQLWrapf(err,
+		return nil, nil, ext, schema.GQLWrapf(err,
 			"succeeded in saving GraphQL schema but failed to alter Dgraph schema ")
 	}
 
-	return assigned, result, nil
+	return assigned, result, ext, nil
 }
 
-func (asr *updateSchemaResolver) Query(ctx context.Context, query *gql.GraphQuery) ([]byte, error) {
+func (asr *updateSchemaResolver) Query(ctx context.Context, query *gql.GraphQuery) ([]byte,
+	*schema.Extensions, error) {
 	return doQuery(asr.admin.schema, asr.mutation.QueryField())
 }
 
@@ -121,11 +123,12 @@ func (gsr *getSchemaResolver) Rewrite(ctx context.Context,
 	return nil, nil
 }
 
-func (gsr *getSchemaResolver) Query(ctx context.Context, query *gql.GraphQuery) ([]byte, error) {
+func (gsr *getSchemaResolver) Query(ctx context.Context, query *gql.GraphQuery) ([]byte,
+	*schema.Extensions, error) {
 	return doQuery(gsr.admin.schema, gsr.gqlQuery)
 }
 
-func doQuery(gql *gqlSchema, field schema.Field) ([]byte, error) {
+func doQuery(gql *gqlSchema, field schema.Field) ([]byte, *schema.Extensions, error) {
 
 	var buf bytes.Buffer
 	x.Check2(buf.WriteString(`{ "`))
@@ -155,7 +158,7 @@ func doQuery(gql *gqlSchema, field schema.Field) ([]byte, error) {
 	}
 	x.Check2(buf.WriteString("}]}"))
 
-	return buf.Bytes(), nil
+	return buf.Bytes(), nil, nil
 }
 
 func getSchemaInput(m schema.Mutation) (*updateGQLSchemaInput, error) {

graphql/dgraph/execute.go

@@ -31,8 +31,10 @@ import (
 	"github.com/dgraph-io/dgraph/x"
 )
 
+const touchedUidsKey = "_total"
+
 // Query is the underlying dgraph implementation of QueryExecutor.
-func Query(ctx context.Context, query *gql.GraphQuery) ([]byte, error) {
+func Query(ctx context.Context, query *gql.GraphQuery) ([]byte, *schema.Extensions, error) {
 	span := trace.FromContext(ctx)
 	stop := x.SpanTimer(span, "dgraph.Query")
 	defer stop()
@@ -50,21 +52,24 @@ func Query(ctx context.Context, query *gql.GraphQuery) ([]byte, error) {
 
 	ctx = context.WithValue(ctx, edgraph.IsGraphql, true)
 	resp, err := (&edgraph.Server{}).Query(ctx, req)
-	return resp.GetJson(), schema.GQLWrapf(err, "Dgraph query failed")
+	ext := &schema.Extensions{TouchedUids: resp.GetMetrics().GetNumUids()[touchedUidsKey]}
+
+	return resp.GetJson(), ext, schema.GQLWrapf(err, "Dgraph query failed")
 }
 
 // Mutate is the underlying dgraph implementation of MutationExecutor.
 func Mutate(
 	ctx context.Context,
 	query *gql.GraphQuery,
-	mutations []*dgoapi.Mutation) (map[string]string, map[string]interface{}, error) {
+	mutations []*dgoapi.Mutation) (map[string]string, map[string]interface{}, *schema.Extensions,
+	error) {
 
 	span := trace.FromContext(ctx)
 	stop := x.SpanTimer(span, "dgraph.Mutate")
 	defer stop()
 
 	if query == nil && len(mutations) == 0 {
-		return nil, nil, nil
+		return nil, nil, nil, nil
 	}
 
 	queryStr := AsString(query)
@@ -88,15 +93,16 @@ func Mutate(
 	ctx = context.WithValue(ctx, edgraph.IsGraphql, true)
 	resp, err := (&edgraph.Server{}).Query(ctx, req)
 	if err != nil {
-		return nil, nil, schema.GQLWrapf(err, "Dgraph mutation failed")
+		return nil, nil, nil, schema.GQLWrapf(err, "Dgraph mutation failed")
 	}
 
+	ext := &schema.Extensions{TouchedUids: resp.GetMetrics().GetNumUids()[touchedUidsKey]}
 	result := make(map[string]interface{})
 	if query != nil && len(resp.GetJson()) != 0 {
 		if err := json.Unmarshal(resp.GetJson(), &result); err != nil {
-			return nil, nil,
+			return nil, nil, ext,
 				schema.GQLWrapf(err, "Couldn't unmarshal response from Dgraph mutation")
 		}
 	}
-	return resp.GetUids(), result, schema.GQLWrapf(err, "Dgraph mutation failed")
+	return resp.GetUids(), result, ext, schema.GQLWrapf(err, "Dgraph mutation failed")
 }

graphql/e2e/common/common.go

@@ -265,6 +265,7 @@ func RunAll(t *testing.T) {
 	t.Run("query state by xid regex", queryStateByXidRegex)
 	t.Run("multiple operations", multipleOperations)
 	t.Run("query post with author", queryPostWithAuthor)
+	t.Run("queries have extensions", queriesHaveExtensions)
 
 	// mutation tests
 	t.Run("add mutation", addMutation)
@@ -296,6 +297,7 @@ func RunAll(t *testing.T) {
 	t.Run("empty delete", mutationEmptyDelete)
 	t.Run("password in mutation", passwordTest)
 	t.Run("duplicate xid in single mutation", deepMutationDuplicateXIDsSameObjectTest)
+	t.Run("mutations have extensions", mutationsHaveExtensions)
 
 	// error tests
 	t.Run("graphql completion on", graphQLCompletionOn)

graphql/e2e/common/error.go

@@ -265,17 +265,19 @@ func panicCatcher(t *testing.T) {
 
 type panicClient struct{}
 
-func (dg *panicClient) Query(ctx context.Context, query *gql.GraphQuery) ([]byte, error) {
+func (dg *panicClient) Query(ctx context.Context, query *gql.GraphQuery) ([]byte,
+	*schema.Extensions, error) {
 	x.Panic(errors.New(panicMsg))
-	return nil, nil
+	return nil, nil, nil
 }
 
 func (dg *panicClient) Mutate(
 	ctx context.Context,
 	query *gql.GraphQuery,
-	mutations []*dgoapi.Mutation) (map[string]string, map[string]interface{}, error) {
+	mutations []*dgoapi.Mutation) (map[string]string, map[string]interface{},
+	*schema.Extensions, error) {
 	x.Panic(errors.New(panicMsg))
-	return nil, nil, nil
+	return nil, nil, nil, nil
 }
 
 // clientInfoLogin check whether the client info(IP address) is propagated in the request.

graphql/e2e/common/mutation.go

@@ -3024,3 +3024,32 @@ func getXidFilter(xidKey string, xidVals []string) map[string]interface{} {
 
 	return filter
 }
+
+func mutationsHaveExtensions(t *testing.T) {
+	mutation := &GraphQLParams{
+		Query: `mutation {
+			addCategory(input: [{ name: "cat" }]) {
+				category {
+					id
+				}
+			}
+		}`,
+	}
+
+	touchedUidskey := "touched_uids"
+	gqlResponse := mutation.ExecuteAsPost(t, graphqlURL)
+	requireNoGQLErrors(t, gqlResponse)
+	require.Contains(t, gqlResponse.Extensions, touchedUidskey)
+	require.Greater(t, int(gqlResponse.Extensions[touchedUidskey].(float64)), 0)
+
+	// cleanup
+	var resp struct {
+		AddCategory struct {
+			Category []category
+		}
+	}
+	err := json.Unmarshal(gqlResponse.Data, &resp)
+	require.NoError(t, err)
+	deleteGqlType(t, "Category",
+		map[string]interface{}{"id": []string{resp.AddCategory.Category[0].ID}}, 1, nil)
+}

graphql/e2e/common/query.go

@@ -1424,3 +1424,19 @@ func queryPostWithAuthor(t *testing.T) {
 		`{"queryPost":[{"title":"Introducing GraphQL in Dgraph","author":{"name":"Ann Author"}}]}`,
 		string(gqlResponse.Data))
 }
+
+func queriesHaveExtensions(t *testing.T) {
+	query := &GraphQLParams{
+		Query: `query {
+			queryPost {
+				title
+			}
+		}`,
+	}
+
+	touchedUidskey := "touched_uids"
+	gqlResponse := query.ExecuteAsPost(t, graphqlURL)
+	requireNoGQLErrors(t, gqlResponse)
+	require.Contains(t, gqlResponse.Extensions, touchedUidskey)
+	require.Greater(t, int(gqlResponse.Extensions[touchedUidskey].(float64)), 0)
+}

graphql/e2e/subscription/subscription_test.go

@@ -17,6 +17,7 @@
 package subscription_test
 
 import (
+	"encoding/json"
 	"testing"
 	"time"
 
@@ -86,11 +87,19 @@ func TestSubscription(t *testing.T) {
 		  }`,
 	})
 	require.Nil(t, err)
+
 	res, err := subscriptionClient.RecvMsg()
+	require.NoError(t, err)
 
+	touchedUidskey := "touched_uids"
+	var subscriptionResp common.GraphQLResponse
+	err = json.Unmarshal(res, &subscriptionResp)
 	require.NoError(t, err)
+	require.Nil(t, subscriptionResp.Errors)
 
-	require.JSONEq(t, `{"data":{"getProduct":{"name":"sanitizer"}}}`, string(res))
+	require.JSONEq(t, `{"getProduct":{"name":"sanitizer"}}`, string(subscriptionResp.Data))
+	require.Contains(t, subscriptionResp.Extensions, touchedUidskey)
+	require.Greater(t, int(subscriptionResp.Extensions[touchedUidskey].(float64)), 0)
 
 	// Background indexing is happening so wait till it get indexed.
 	time.Sleep(time.Second * 2)
@@ -112,8 +121,17 @@ func TestSubscription(t *testing.T) {
 	res, err = subscriptionClient.RecvMsg()
 	require.NoError(t, err)
 
+	// makes sure that the we have a fresh instance to unmarshal to, otherwise there may be things
+	// from the previous unmarshal
+	subscriptionResp = common.GraphQLResponse{}
+	err = json.Unmarshal(res, &subscriptionResp)
+	require.NoError(t, err)
+	require.Nil(t, subscriptionResp.Errors)
+
 	// Check the latest update.
-	require.JSONEq(t, `{"data":{"getProduct":{"name":"mask"}}}`, string(res))
+	require.JSONEq(t, `{"getProduct":{"name":"mask"}}`, string(subscriptionResp.Data))
+	require.Contains(t, subscriptionResp.Extensions, touchedUidskey)
+	require.Greater(t, int(subscriptionResp.Extensions[touchedUidskey].(float64)), 0)
 
 	time.Sleep(2 * time.Second)
 	// Change schema to terminate subscription..

graphql/resolve/mutation.go

@@ -96,13 +96,14 @@ type MutationRewriter interface {
 // mutated map and any errors.
 type MutationExecutor interface {
 	// Mutate performs the actual mutation and returns a map of newly assigned nodes,
-	// a map of variable->[]uid from upsert mutations and any errors.  If an error
+	// a map of variable->[]uid from upsert mutations, extensions and any errors. If an error
 	// occurs, that indicates that the mutation failed in some way significant enough
-	// way as to not continue procissing this mutation or others in the same request.
+	// way as to not continue processing this mutation or others in the same request.
 	Mutate(
 		ctx context.Context,
 		query *gql.GraphQuery,
-		mutations []*dgoapi.Mutation) (map[string]string, map[string]interface{}, error)
+		mutations []*dgoapi.Mutation) (map[string]string, map[string]interface{},
+		*schema.Extensions, error)
 }
 
 // MutationResolverFunc is an adapter that allows to build a MutationResolver from
@@ -114,7 +115,8 @@ type MutationResolverFunc func(ctx context.Context, mutation schema.Mutation) (*
 type MutationExecutionFunc func(
 	ctx context.Context,
 	query *gql.GraphQuery,
-	mutations []*dgoapi.Mutation) (map[string]string, map[string][]string, error)
+	mutations []*dgoapi.Mutation) (map[string]string, map[string][]string, *schema.Extensions,
+	error)
 
 // Resolve calls mr(ctx, mutation)
 func (mr MutationResolverFunc) Resolve(
@@ -128,7 +130,8 @@ func (mr MutationResolverFunc) Resolve(
 func (me MutationExecutionFunc) Mutate(
 	ctx context.Context,
 	query *gql.GraphQuery,
-	mutations []*dgoapi.Mutation) (map[string]string, map[string][]string, error) {
+	mutations []*dgoapi.Mutation) (map[string]string, map[string][]string, *schema.Extensions,
+	error) {
 	return me(ctx, query, mutations)
 }
 
@@ -204,7 +207,7 @@ func (mr *mutationResolver) rewriteAndExecute(
 			resolverFailed
 	}
 
-	assigned, result, err := mr.mutationExecutor.Mutate(ctx, query, mutations)
+	assigned, result, extM, err := mr.mutationExecutor.Mutate(ctx, query, mutations)
 	if err != nil {
 		gqlErr := schema.GQLWrapLocationf(
 			err, mutation.Location(), "mutation %s failed", mutation.Name())
@@ -220,10 +223,17 @@ func (mr *mutationResolver) rewriteAndExecute(
 		return emptyResult(errs), resolverFailed
 	}
 
-	resp, err := mr.queryExecutor.Query(ctx, dgQuery)
+	resp, extQ, err := mr.queryExecutor.Query(ctx, dgQuery)
 	errs = schema.AppendGQLErrs(errs, schema.GQLWrapf(err,
 		"couldn't rewrite query for mutation %s", mutation.Name()))
 
+	// merge the extensions we got from .Mutate() and .Query() into extM
+	if extM == nil {
+		extM = extQ
+	} else {
+		extM.Merge(extQ)
+	}
+
 	resolved := completeDgraphResult(ctx, mutation.QueryField(), resp, errs)
 	if resolved.Data == nil && resolved.Err != nil {
 		return &Resolved{
@@ -233,8 +243,9 @@ func (mr *mutationResolver) rewriteAndExecute(
 					schema.Typename:                      mutation.TypeName,
 					mutation.QueryField().ResponseName(): nil,
 				}},
-			Field: mutation,
-			Err:   err,
+			Field:      mutation,
+			Err:        err,
+			Extensions: extM,
 		}, resolverSucceeded
 	}
 
@@ -245,9 +256,11 @@ func (mr *mutationResolver) rewriteAndExecute(
 	dgRes := resolved.Data.(map[string]interface{})
 	dgRes[schema.NumUid] = numUids
 	dgRes[schema.Typename] = mutation.Type().Name()
-	resolved.Data = map[string]interface{}{mutation.ResponseName(): dgRes}
 
+	resolved.Data = map[string]interface{}{mutation.ResponseName(): dgRes}
 	resolved.Field = mutation
+	resolved.Extensions = extM
+
 	return resolved, resolverSucceeded
 }