Return max complexity for UnionType / Interfaces

src/QueryComplexity.ts

@@ -19,7 +19,7 @@ import {
   GraphQLField, isCompositeType, GraphQLCompositeType, GraphQLFieldMap,
   GraphQLSchema, DocumentNode, TypeInfo,
   visit, visitWithTypeInfo,
-  GraphQLDirective,
+  GraphQLDirective, isAbstractType,
 } from 'graphql';
 import {
   GraphQLUnionType,
@@ -48,6 +48,11 @@ export type ComplexityEstimator = (options: ComplexityEstimatorArgs) => number |
 // Complexity can be anything that is supported by the configured estimators
 export type Complexity = any;
 
+// Map of complexities for possible types (of Union, Interface types)
+type ComplexityMap = {
+  [typeName: string]: number,
+}
+
 export interface QueryComplexityOptions {
   // The maximum allowed query complexity, queries above this threshold will be rejected
   maximumComplexity: number,
@@ -179,16 +184,25 @@ export default class QueryComplexity {
   nodeComplexity(
     node: FieldNode | FragmentDefinitionNode | InlineFragmentNode | OperationDefinitionNode,
     typeDef: GraphQLObjectType | GraphQLInterfaceType | GraphQLUnionType,
-    complexity: number = 0
   ): number {
     if (node.selectionSet) {
       let fields:GraphQLFieldMap<any, any> = {};
       if (typeDef instanceof GraphQLObjectType || typeDef instanceof GraphQLInterfaceType) {
         fields = typeDef.getFields();
       }
-      return complexity + node.selectionSet.selections.reduce(
-        (total: number, childNode: FieldNode | FragmentSpreadNode | InlineFragmentNode) => {
-          let nodeComplexity = 0;
+
+      // Determine all possible types of the current node
+      let possibleTypeNames: string[];
+      if (isAbstractType(typeDef)) {
+        possibleTypeNames = this.context.getSchema().getPossibleTypes(typeDef).map(t => t.name);
+      } else {
+        possibleTypeNames = [typeDef.name];
+      }
+
+      // Collect complexities for all possible types individually
+      const selectionSetComplexities: ComplexityMap = node.selectionSet.selections.reduce(
+        (complexities: ComplexityMap, childNode: FieldNode | FragmentSpreadNode | InlineFragmentNode) => {
+          // let nodeComplexity = 0;
 
           let includeNode = true;
           let skipNode = false;
@@ -210,7 +224,7 @@ export default class QueryComplexity {
           });
 
           if (!includeNode || skipNode) {
-            return total;
+            return complexities;
           }
 
           switch (childNode.kind) {
@@ -248,7 +262,11 @@ export default class QueryComplexity {
                 const tmpComplexity = estimator(estimatorArgs);
 
                 if (typeof tmpComplexity === 'number' && !isNaN(tmpComplexity)) {
-                  nodeComplexity = tmpComplexity;
+                  complexities = addComplexities(
+                    tmpComplexity,
+                    complexities,
+                    possibleTypeNames,
+                  );
                   return true;
                 }
 
@@ -269,30 +287,69 @@ export default class QueryComplexity {
               const fragmentType = assertCompositeType(
                 this.context.getSchema().getType(fragment.typeCondition.name.value)
               );
-              nodeComplexity = this.nodeComplexity(fragment, fragmentType);
+              const nodeComplexity = this.nodeComplexity(fragment, fragmentType);
+              if (isAbstractType(fragmentType)) {
+                // Add fragment complexity for all possible types
+                complexities = addComplexities(
+                  nodeComplexity,
+                  complexities,
+                  this.context.getSchema().getPossibleTypes(fragmentType).map(t => t.name),
+                );
+              } else {
+                // Add complexity for object type
+                complexities = addComplexities(
+                  nodeComplexity,
+                  complexities,
+                  [fragmentType.name],
+                );
+              }
               break;
             }
             case Kind.INLINE_FRAGMENT: {
               let inlineFragmentType = typeDef;
               if (childNode.typeCondition && childNode.typeCondition.name) {
-                // $FlowFixMe: Not sure why flow thinks this can still be NULL
                 inlineFragmentType = assertCompositeType(
                   this.context.getSchema().getType(childNode.typeCondition.name.value)
                 );
               }
 
-              nodeComplexity = this.nodeComplexity(childNode, inlineFragmentType);
+              const nodeComplexity = this.nodeComplexity(childNode, inlineFragmentType);
+              if (isAbstractType(inlineFragmentType)) {
+                // Add fragment complexity for all possible types
+                complexities = addComplexities(
+                  nodeComplexity,
+                  complexities,
+                  this.context.getSchema().getPossibleTypes(inlineFragmentType).map(t => t.name),
+                );
+              } else {
+                // Add complexity for object type
+                complexities = addComplexities(
+                  nodeComplexity,
+                  complexities,
+                  [inlineFragmentType.name],
+                );
+              }
               break;
             }
             default: {
-              nodeComplexity = this.nodeComplexity(childNode, typeDef);
+              complexities = addComplexities(
+                this.nodeComplexity(childNode, typeDef),
+                complexities,
+                possibleTypeNames,
+              );
               break;
             }
           }
-          return Math.max(nodeComplexity, 0) + total;
-        }, complexity);
+
+          return complexities;
+        }, {});
+      // Only return max complexity of all possible types
+      if (!selectionSetComplexities) {
+        return NaN;
+      }
+      return Math.max(...Object.values(selectionSetComplexities), 0);
     }
-    return complexity;
+    return 0;
   }
 
   createError(): GraphQLError {
@@ -308,3 +365,24 @@ export default class QueryComplexity {
     ));
   }
 }
+
+/**
+ * Adds a complexity to the complexity map for all possible types
+ * @param complexity
+ * @param complexityMap
+ * @param possibleTypes
+ */
+function addComplexities(
+  complexity: number,
+  complexityMap: ComplexityMap,
+  possibleTypes: string[],
+): ComplexityMap {
+  for (const type of possibleTypes) {
+    if (complexityMap.hasOwnProperty(type)) {
+      complexityMap[type] = complexityMap[type] + complexity;
+    } else {
+      complexityMap[type] = complexity;
+    }
+  }
+  return complexityMap;
+}

src/__tests__/QueryComplexity-test.ts

@@ -408,7 +408,7 @@ describe('QueryComplexity analysis', () => {
     );
   });
 
-  it('should return NaN when no astNode available on field when use directiveEstimator', () => {
+  it('should return NaN when no astNode available on field when using directiveEstimator', () => {
     const ast = parse(`
       query {
         _service {
@@ -482,4 +482,179 @@ describe('QueryComplexity analysis', () => {
     });
     expect(complexity2).to.equal(20);
   });
+
+  it('should calculate max complexity for fragment on union type', () => {
+    const query = parse(`
+      query Primary {
+        union {
+          ...on Item {
+            scalar
+          }
+          ...on SecondItem {
+            scalar
+          }
+          ...on SecondItem {
+            scalar
+          }
+        }
+      }
+    `);
+
+    const complexity = getComplexity({
+      estimators: [
+        fieldExtensionsEstimator(),
+        simpleEstimator({defaultComplexity: 1})
+      ],
+      schema,
+      query,
+    });
+    expect(complexity).to.equal(3);
+  });
+
+  it('should calculate max complexity for nested fragment on union type', () => {
+    const query = parse(`
+      query Primary {
+        union {
+          ...on Union {
+            ...on Item {
+              complexScalar1: complexScalar
+            } 
+          }
+          ...on SecondItem {
+            scalar
+          }
+          ...on Item {
+            complexScalar2: complexScalar
+          }
+        }
+      }
+    `);
+
+    const complexity = getComplexity({
+      estimators: [
+        fieldExtensionsEstimator(),
+        simpleEstimator({defaultComplexity: 0})
+      ],
+      schema,
+      query,
+    });
+    expect(complexity).to.equal(40);
+  });
+
+  it('should calculate max complexity for nested fragment on union type + named fragment', () => {
+    const query = parse(`
+      query Primary {
+        union {
+          ...F
+          ...on SecondItem {
+            scalar
+          }
+          ...on Item {
+            complexScalar2: complexScalar
+          }
+        }
+      }
+      fragment F on Union {
+        ...on Item {
+          complexScalar1: complexScalar
+        } 
+      }
+    `);
+
+    const complexity = getComplexity({
+      estimators: [
+        fieldExtensionsEstimator(),
+        simpleEstimator({defaultComplexity: 0})
+      ],
+      schema,
+      query,
+    });
+    expect(complexity).to.equal(40);
+  });
+
+  it('should calculate max complexity for multiple interfaces', () => {
+    const query = parse(`
+      query Primary {
+        interface {
+          ...on Query {
+            complexScalar
+          }
+          ...on SecondItem {
+            name
+            name2: name
+          }
+        }
+      }
+    `);
+
+    const complexity = getComplexity({
+      estimators: [
+        fieldExtensionsEstimator(),
+        simpleEstimator({defaultComplexity: 1})
+      ],
+      schema,
+      query,
+    });
+    expect(complexity).to.equal(21);
+  });
+
+  it('should calculate max complexity for multiple interfaces with nesting', () => {
+    const query = parse(`
+      query Primary {
+        interface {
+          ...on Query {
+            complexScalar
+            ...on Query {
+              a: complexScalar
+            }
+          }
+          ...on SecondItem {
+            name
+            name2: name
+          }
+        }
+      }
+    `);
+
+    const complexity = getComplexity({
+      estimators: [
+        fieldExtensionsEstimator(),
+        simpleEstimator({defaultComplexity: 1})
+      ],
+      schema,
+      query,
+    });
+    expect(complexity).to.equal(41); // 1 for interface, 20 * 2 for complexScalar
+  });
+
+  it('should calculate max complexity for multiple interfaces with nesting + named fragment', () => {
+    const query = parse(`
+      query Primary {
+        interface {
+          ...F
+          ...on SecondItem {
+            name
+            name2: name
+          }
+        }
+      }
+
+      fragment F on Query {
+        complexScalar
+        ...on Query {
+          a: complexScalar
+        }
+      }
+    `);
+
+    const complexity = getComplexity({
+      estimators: [
+        fieldExtensionsEstimator(),
+        simpleEstimator({defaultComplexity: 1})
+      ],
+      schema,
+      query,
+    });
+    expect(complexity).to.equal(41); // 1 for interface, 20 * 2 for complexScalar
+  });
 });