github.com/fortexxx/gqlgen@v0.10.3-0.20191216030626-ca5ea8b21ead/graphql/handler/extension/complexity.go (about)

     1  package extension
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  
     7  	"github.com/99designs/gqlgen/complexity"
     8  	"github.com/99designs/gqlgen/graphql"
     9  	"github.com/vektah/gqlparser/gqlerror"
    10  )
    11  
    12  // ComplexityLimit allows you to define a limit on query complexity
    13  //
    14  // If a query is submitted that exceeds the limit, a 422 status code will be returned.
    15  type ComplexityLimit struct {
    16  	Func func(ctx context.Context, rc *graphql.OperationContext) int
    17  
    18  	es graphql.ExecutableSchema
    19  }
    20  
    21  var _ interface {
    22  	graphql.OperationContextMutator
    23  	graphql.HandlerExtension
    24  } = &ComplexityLimit{}
    25  
    26  const complexityExtension = "ComplexityLimit"
    27  
    28  type ComplexityStats struct {
    29  	// The calculated complexity for this request
    30  	Complexity int
    31  
    32  	// The complexity limit for this request returned by the extension func
    33  	ComplexityLimit int
    34  }
    35  
    36  // FixedComplexityLimit sets a complexity limit that does not change
    37  func FixedComplexityLimit(limit int) *ComplexityLimit {
    38  	return &ComplexityLimit{
    39  		Func: func(ctx context.Context, rc *graphql.OperationContext) int {
    40  			return limit
    41  		},
    42  	}
    43  }
    44  
    45  func (c ComplexityLimit) ExtensionName() string {
    46  	return complexityExtension
    47  }
    48  
    49  func (c *ComplexityLimit) Validate(schema graphql.ExecutableSchema) error {
    50  	if c.Func == nil {
    51  		return fmt.Errorf("ComplexityLimit func can not be nil")
    52  	}
    53  	c.es = schema
    54  	return nil
    55  }
    56  
    57  func (c ComplexityLimit) MutateOperationContext(ctx context.Context, rc *graphql.OperationContext) *gqlerror.Error {
    58  	op := rc.Doc.Operations.ForName(rc.OperationName)
    59  	complexity := complexity.Calculate(c.es, op, rc.Variables)
    60  
    61  	limit := c.Func(ctx, rc)
    62  
    63  	rc.Stats.SetExtension(complexityExtension, &ComplexityStats{
    64  		Complexity:      complexity,
    65  		ComplexityLimit: limit,
    66  	})
    67  
    68  	if complexity > limit {
    69  		return gqlerror.Errorf("operation has complexity %d, which exceeds the limit of %d", complexity, limit)
    70  	}
    71  
    72  	return nil
    73  }
    74  
    75  func GetComplexityStats(ctx context.Context) *ComplexityStats {
    76  	rc := graphql.GetOperationContext(ctx)
    77  	if rc == nil {
    78  		return nil
    79  	}
    80  
    81  	s, _ := rc.Stats.GetExtension(complexityExtension).(*ComplexityStats)
    82  	return s
    83  }