Merge branch 'master' into feat-opencensus

99designs · Nov 2, 2018 · 4e7e6a1 · 4e7e6a1
2 parents 59ef91a + 926ad17
commit 4e7e6a1
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 9 deletions.
diff --git a/graphql/context.go b/graphql/context.go
@@ -17,6 +17,10 @@ type RequestContext struct {
 	RawQuery  string
 	Variables map[string]interface{}
 	Doc       *ast.QueryDocument
+
+	ComplexityLimit     int
+	OperationComplexity int
+
 	// ErrorPresenter will be used to generate the error
 	// message from errors given to Error().
 	ErrorPresenter      ErrorPresenterFunc

diff --git a/handler/graphql.go b/handler/graphql.go
@@ -35,7 +35,7 @@ type Config struct {
 	complexityLimit int
 }
 
-func (c *Config) newRequestContext(doc *ast.QueryDocument, query string, variables map[string]interface{}) *graphql.RequestContext {
+func (c *Config) newRequestContext(es graphql.ExecutableSchema, doc *ast.QueryDocument, op *ast.OperationDefinition, query string, variables map[string]interface{}) *graphql.RequestContext {
 	reqCtx := graphql.NewRequestContext(doc, query, variables)
 	if hook := c.recover; hook != nil {
 		reqCtx.Recover = hook
@@ -59,6 +59,12 @@ func (c *Config) newRequestContext(doc *ast.QueryDocument, query string, variabl
 		reqCtx.Tracer = &graphql.NopTracer{}
 	}
 
+	if c.complexityLimit > 0 {
+		reqCtx.ComplexityLimit = c.complexityLimit
+		operationComplexity := complexity.Calculate(es, op, variables)
+		reqCtx.OperationComplexity = operationComplexity
+	}
+
 	return reqCtx
 }
 
@@ -342,7 +348,7 @@ func (gh *graphqlHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		gh.cache.Add(reqParams.Query, doc)
 	}
 
-	reqCtx := gh.cfg.newRequestContext(doc, reqParams.Query, vars)
+	reqCtx := gh.cfg.newRequestContext(gh.exec, doc, op, reqParams.Query, vars)
 	ctx = graphql.WithRequestContext(ctx, reqCtx)
 
 	defer func() {
@@ -352,12 +358,9 @@ func (gh *graphqlHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		}
 	}()
 
-	if gh.cfg.complexityLimit > 0 {
-		queryComplexity := complexity.Calculate(gh.exec, op, vars)
-		if queryComplexity > gh.cfg.complexityLimit {
-			sendErrorf(w, http.StatusUnprocessableEntity, "query has complexity %d, which exceeds the limit of %d", queryComplexity, gh.cfg.complexityLimit)
-			return
-		}
+	if reqCtx.ComplexityLimit > 0 && reqCtx.OperationComplexity > reqCtx.ComplexityLimit {
+		sendErrorf(w, http.StatusUnprocessableEntity, "operation has complexity %d, which exceeds the limit of %d", reqCtx.OperationComplexity, reqCtx.ComplexityLimit)
+		return
 	}
 
 	switch op.Operation {

diff --git a/handler/websocket.go b/handler/websocket.go
@@ -165,7 +165,7 @@ func (c *wsConnection) subscribe(message *operationMessage) bool {
 		c.sendError(message.ID, err)
 		return true
 	}
-	reqCtx := c.cfg.newRequestContext(doc, reqParams.Query, vars)
+	reqCtx := c.cfg.newRequestContext(c.exec, doc, op, reqParams.Query, vars)
 	ctx := graphql.WithRequestContext(c.ctx, reqCtx)
 
 	if c.initPayload != nil {