Add support for Interface and Union types max complexity #26

ivome · ivome · commit c6bf0f5b305c · 2020-08-12T17:49:45.000-05:00
diff --git a/src/QueryComplexity.ts b/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,
@@ -185,9 +190,19 @@ export default class QueryComplexity {
       if (typeDef instanceof GraphQLObjectType || typeDef instanceof GraphQLInterfaceType) {
         fields = typeDef.getFields();
       }
-      return 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;
@@ -209,7 +224,7 @@ export default class QueryComplexity {
           });
 
           if (!includeNode || skipNode) {
-            return total;
+            return complexities;
           }
 
           switch (childNode.kind) {
@@ -247,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;
                 }
 
@@ -268,7 +287,22 @@ 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: {
@@ -279,16 +313,41 @@ export default class QueryComplexity {
                 );
               }
 
-              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;
-        }, 0);
+
+          return complexities;
+        }, {});
+      // Only return max complexity of all possible types
+      if (!selectionSetComplexities) {
+        return NaN;
+      }
+      return Math.max(...Object.values(selectionSetComplexities), 0);
     }
     return 0;
   }
@@ -306,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;
+}
diff --git a/src/__tests__/QueryComplexity-test.ts b/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 {
@@ -510,4 +510,151 @@ describe('QueryComplexity analysis', () => {
     });
     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
+  });
 });
diff --git a/src/__tests__/fixtures/schema.ts b/src/__tests__/fixtures/schema.ts
@@ -63,6 +63,14 @@ const NameInterface = new GraphQLInterfaceType({
   resolveType: () => Item
 });
 
+const UnionInterface = new GraphQLInterfaceType({
+  name: 'UnionInterface',
+  fields: () => ({
+    union: { type: Union }
+  }),
+  resolveType: () => Item
+});
+
 const SecondItem = new GraphQLObjectType({
   name: 'SecondItem',
   fields: () => ({
@@ -92,7 +100,15 @@ const SDLInterface = new GraphQLInterfaceType({
   fields: {
     sdl: { type: GraphQLString }
   },
-  resolveType: () => '"SDL"'
+  resolveType: () => 'SDL'
+});
+
+const SDL = new GraphQLObjectType({
+  name: 'SDL',
+  fields: {
+    sdl: { type: GraphQLString }
+  },
+  interfaces: () => [SDLInterface],
 });
 
 const Query = new GraphQLObjectType({
@@ -145,6 +161,12 @@ const Query = new GraphQLObjectType({
     },
     _service: {type: SDLInterface},
   }),
+  interfaces: () => [NameInterface, UnionInterface,]
 });
 
-export default new GraphQLSchema({ query: Query });
+export default new GraphQLSchema({
+  query: Query,
+  types: [
+    SDL,
+  ],
+});
diff --git a/tsconfig.json b/tsconfig.json
@@ -17,6 +17,7 @@
     },
     "lib": [
       "es2015",
+      "es2018",
       "esnext.asynciterable",
       "dom"
     ]

Original file line number	Diff line number	Diff line change
`@@ -17,6 +17,7 @@`
`17`	`17`	`},`
`18`	`18`	`"lib": [`
`19`	`19`	`"es2015",`
	`20`	`+ "es2018",`
`20`	`21`	`"esnext.asynciterable",`
`21`	`22`	`"dom"`
`22`	`23`	`]`