Merge pull request #315 from edsrzf/query-complexity

Query complexity calculation and limits
99designs · Aug 31, 2018 · 7d44dd6 · 7d44dd6
2 parents d08b9c4 + 2ab857e
commit 7d44dd6
Show file tree

Hide file tree

Showing 19 changed files with 2,610 additions and 11 deletions.
diff --git a/codegen/config.go b/codegen/config.go
@@ -84,8 +84,8 @@ type TypeMapEntry struct {
 }
 
 type TypeMapField struct {
-	Resolver  bool   `yaml:"resolver"`
-	FieldName string `yaml:"fieldName"`
+	Resolver   bool   `yaml:"resolver"`
+	FieldName  string `yaml:"fieldName"`
 }
 
 func (c *PackageConfig) normalize() error {

diff --git a/codegen/object.go b/codegen/object.go
@@ -81,10 +81,18 @@ func (o *Object) IsConcurrent() bool {
 	return false
 }
 
+func (o *Object) IsReserved() bool {
+	return strings.HasPrefix(o.GQLType, "__")
+}
+
 func (f *Field) IsResolver() bool {
 	return f.GoFieldName == ""
 }
 
+func (f *Field) IsReserved() bool {
+	return strings.HasPrefix(f.GQLName, "__")
+}
+
 func (f *Field) IsMethod() bool {
 	return f.GoFieldType == GoFieldMethod
 }
@@ -165,6 +173,24 @@ func (f *Field) ResolverDeclaration() string {
 	return res
 }
 
+func (f *Field) ComplexitySignature() string {
+	res := fmt.Sprintf("func(childComplexity int")
+	for _, arg := range f.Args {
+		res += fmt.Sprintf(", %s %s", arg.GoVarName, arg.Signature())
+	}
+	res += ") int"
+	return res
+}
+
+func (f *Field) ComplexityArgs() string {
+	var args []string
+	for _, arg := range f.Args {
+		args = append(args, "args["+strconv.Quote(arg.GQLName)+"].("+arg.Signature()+")")
+	}
+
+	return strings.Join(args, ", ")
+}
+
 func (f *Field) CallArgs() string {
 	var args []string
 
@@ -198,7 +224,7 @@ func (f *Field) doWriteJson(val string, remainingMods []string, astType *ast.Typ
 						ec.Errorf(ctx, "must not be null")
 					}
 				{{- end }}
-				return graphql.Null 
+				return graphql.Null
 			}
 			{{.next }}`, map[string]interface{}{
 			"val":     val,

diff --git a/codegen/templates/data.go b/codegen/templates/data.go
diff --git a/codegen/templates/generated.gotpl b/codegen/templates/generated.gotpl
@@ -13,12 +13,14 @@ func NewExecutableSchema(cfg Config) graphql.ExecutableSchema {
 	return &executableSchema{
 		resolvers: cfg.Resolvers,
 		directives: cfg.Directives,
+		complexity: cfg.Complexity,
 	}
 }
 
 type Config struct {
 	Resolvers  ResolverRoot
 	Directives DirectiveRoot
+	Complexity ComplexityRoot
 }
 
 type ResolverRoot interface {
@@ -35,7 +37,21 @@ type DirectiveRoot struct {
 {{ end }}
 }
 
-{{- range $object := .Objects -}}
+type ComplexityRoot struct {
+{{ range $object := .Objects }}
+	{{ if $object.IsReserved }}{{ else -}}
+		{{ $object.GQLType|toCamel }} struct {
+		{{ range $field := $object.Fields -}}
+			{{ if $field.IsReserved }}{{ else -}}
+				{{ $field.GQLName|toCamel }} {{ $field.ComplexitySignature }}
+			{{ end }}
+		{{- end }}
+		}
+	{{- end }}
+{{ end }}
+}
+
+{{ range $object := .Objects -}}
 	{{ if $object.HasResolvers }}
 		type {{$object.GQLType}}Resolver interface {
 		{{ range $field := $object.Fields -}}
@@ -48,12 +64,44 @@ type DirectiveRoot struct {
 type executableSchema struct {
 	resolvers  ResolverRoot
 	directives DirectiveRoot
+	complexity ComplexityRoot
 }
 
 func (e *executableSchema) Schema() *ast.Schema {
 	return parsedSchema
 }
 
+func (e *executableSchema) Complexity(typeName, field string, childComplexity int, rawArgs map[string]interface{}) (int, bool) {
+	switch typeName + "." + field {
+	{{ range $object := .Objects }}
+		{{ if $object.IsReserved }}{{ else }}
+			{{ range $field := $object.Fields }}
+				{{ if $field.IsReserved }}{{ else }}
+					case "{{$object.GQLType}}.{{$field.GQLName}}":
+						if e.complexity.{{$object.GQLType|toCamel}}.{{$field.GQLName|toCamel}} == nil {
+							break
+						}
+						{{ if $field.Args }}args := map[string]interface{}{} {{end}}
+						{{ range $i, $arg := $field.Args }}
+							var arg{{$i}} {{$arg.Signature }}
+							if tmp, ok := rawArgs[{{$arg.GQLName|quote}}]; ok {
+								var err error
+								{{$arg.Unmarshal (print "arg" $i) "tmp" }}
+								if err != nil {
+									return 0, false
+								}
+							}
+							args[{{$arg.GQLName|quote}}] = arg{{$i}}
+						{{ end }}
+						return e.complexity.{{$object.GQLType|toCamel}}.{{$field.GQLName|toCamel}}(childComplexity{{if $field.Args}}, {{$field.ComplexityArgs}} {{end}}), true
+				{{ end }}
+			{{ end }}
+		{{ end }}
+	{{ end }}
+	}
+	return 0, false
+}
+
 func (e *executableSchema) Query(ctx context.Context, op *ast.OperationDefinition) *graphql.Response {
 	{{- if .QueryRoot }}
 		ec := executionContext{graphql.GetRequestContext(ctx), e}