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 }