From 32705737b4770c7ee77f609ec3fb5f31b4060276 Mon Sep 17 00:00:00 2001
From: Michael Hayes <michael@hayes.io>
Date: Thu, 25 Aug 2022 16:24:20 -0700
Subject: [PATCH 1/2] Add a context option for passing the request context to
 estimators

---
 README.md                             |  6 +++++
 src/QueryComplexity.ts                |  9 +++++++
 src/__tests__/QueryComplexity-test.ts | 34 ++++++++++++++++++++++++---
 3 files changed, 46 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index dce7d3e..e94943a 100644
--- a/README.md
+++ b/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;
diff --git a/src/QueryComplexity.ts b/src/QueryComplexity.ts
index 348c4f5..091f45c 100644
--- a/src/QueryComplexity.ts
+++ b/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);
diff --git a/src/__tests__/QueryComplexity-test.ts b/src/__tests__/QueryComplexity-test.ts
index a0979db..6c22ae3 100644
--- a/src/__tests__/QueryComplexity-test.ts
+++ b/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);
+  });
 });

From 7a2aef4ced947fb5c0ca91cca08232b7fddca6f9 Mon Sep 17 00:00:00 2001
From: Michael Hayes <michael@hayes.io>
Date: Thu, 25 Aug 2022 21:07:55 -0700
Subject: [PATCH 2/2] Fix tests

---
 src/__tests__/utils/compatResolveType.ts | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/src/__tests__/utils/compatResolveType.ts b/src/__tests__/utils/compatResolveType.ts
index 26d229c..c96c813 100644
--- a/src/__tests__/utils/compatResolveType.ts
+++ b/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;