Add a context option for passing the request context to estimators (#77)

* Add a context option for passing the request context to estimators

* Fix tests

Author: hayes

README.md

@@ -38,6 +38,9 @@ const rule = createComplexityRule({
   // in the visitor of the graphql-js library
   variables: {},
 
+  // The context object for the request (optional)
+  context: {}
+
   // specify operation name only when pass multi-operation documents
   operationName?: string,
 
@@ -110,6 +113,9 @@ type ComplexityEstimatorArgs = {
 
   // The complexity of all child selections for that field
   childComplexity: number;
+
+  // The context object for the request if it was provided
+  context?: Record<string, any>;
 };
 
 type ComplexityEstimator = (options: ComplexityEstimatorArgs) => number | void;

src/QueryComplexity.ts

@@ -43,6 +43,7 @@ export type ComplexityEstimatorArgs = {
   node: FieldNode;
   args: { [key: string]: any };
   childComplexity: number;
+  context?: Record<string, any>;
 };
 
 export type ComplexityEstimator = (
@@ -78,6 +79,9 @@ export interface QueryComplexityOptions {
 
   // An array of complexity estimators to use for estimating the complexity
   estimators: Array<ComplexityEstimator>;
+
+  // Pass request context to the estimators via estimationContext
+  context?: Record<string, any>;
 }
 
 function queryComplexityMessage(max: number, actual: number): string {
@@ -93,6 +97,7 @@ export function getComplexity(options: {
   query: DocumentNode;
   variables?: Record<string, any>;
   operationName?: string;
+  context?: Record<string, any>;
 }): number {
   const typeInfo = new TypeInfo(options.schema);
 
@@ -109,6 +114,7 @@ export function getComplexity(options: {
     estimators: options.estimators,
     variables: options.variables,
     operationName: options.operationName,
+    context: options.context,
   });
 
   visit(options.query, visitWithTypeInfo(typeInfo, visitor));
@@ -130,6 +136,7 @@ export default class QueryComplexity {
   includeDirectiveDef: GraphQLDirective;
   skipDirectiveDef: GraphQLDirective;
   variableValues: Record<string, any>;
+  requestContext?: Record<string, any>;
 
   constructor(context: ValidationContext, options: QueryComplexityOptions) {
     if (
@@ -149,6 +156,7 @@ export default class QueryComplexity {
     this.skipDirectiveDef = this.context.getSchema().getDirective('skip');
     this.estimators = options.estimators;
     this.variableValues = {};
+    this.requestContext = options.context;
 
     this.OperationDefinition = {
       enter: this.onOperationDefinitionEnter,
@@ -327,6 +335,7 @@ export default class QueryComplexity {
                   field,
                   node: childNode,
                   type: typeDef,
+                  context: this.requestContext,
                 };
                 const validScore = this.estimators.find((estimator) => {
                   const tmpComplexity = estimator(estimatorArgs);

src/__tests__/QueryComplexity-test.ts

@@ -639,7 +639,7 @@ describe('QueryComplexity analysis', () => {
           ...on Union {
             ...on Item {
               complexScalar1: complexScalar
-            } 
+            }
           }
           ...on SecondItem {
             scalar
@@ -678,7 +678,7 @@ describe('QueryComplexity analysis', () => {
       fragment F on Union {
         ...on Item {
           complexScalar1: complexScalar
-        } 
+        }
       }
     `);
 
@@ -759,7 +759,7 @@ describe('QueryComplexity analysis', () => {
           }
         }
       }
-      
+
       fragment F on Query {
         complexScalar
         ...on Query {
@@ -832,4 +832,32 @@ describe('QueryComplexity analysis', () => {
       }),
     ]);
   });
+
+  it('passed context to estimators', () => {
+    const ast = parse(`
+      query {
+        scalar
+        requiredArgs(count: 10) {
+          scalar
+        }
+      }
+    `);
+
+    const contextEstimator: ComplexityEstimator = ({
+      context,
+      childComplexity,
+    }) => {
+      return context['complexityMultiplier'] * (childComplexity || 1);
+    };
+
+    const complexity = getComplexity({
+      estimators: [contextEstimator],
+      schema,
+      query: ast,
+      context: { complexityMultiplier: 5 },
+    });
+
+    // query.scalar(5) + query.requiredArgs(5) * requiredArgs.scalar(5)
+    expect(complexity).to.equal(30);
+  });
 });

src/__tests__/utils/compatResolveType.ts

@@ -1,6 +1,4 @@
-import { GraphQLType } from 'graphql';
-
-import graphqlPackage from 'graphql/package.json';
+import * as graphql from 'graphql';
 import semver from 'semver';
 
 /**
@@ -10,8 +8,8 @@ import semver from 'semver';
  * @param type
  * @returns
  */
-export function compatResolveType(type: GraphQLType): any {
-  if (semver.gte(graphqlPackage.version, '16.0.0')) {
+export function compatResolveType(type: graphql.GraphQLType): any {
+  if (graphql.version && semver.gte(graphql.version, '16.0.0')) {
     return () => type.toString();
   } else {
     return () => type;