From d8def8233c1eb4228f00aad7adc3c72f88ae6699 Mon Sep 17 00:00:00 2001
From: Wei Xue <wxue@yelp.com>
Date: Wed, 7 Apr 2021 10:26:36 -0700
Subject: [PATCH] Added support for client defined required field

---
 src/execution/execute.js |  4 +++-
 src/language/ast.d.ts    |  1 +
 src/language/parser.js   | 34 ++++++++++++++++++++++++----------
 src/language/printer.js  |  6 ++++--
 4 files changed, 32 insertions(+), 13 deletions(-)

diff --git a/src/execution/execute.js b/src/execution/execute.js
index 0d7f0f7a12..b3732f9860 100644
--- a/src/execution/execute.js
+++ b/src/execution/execute.js
@@ -38,6 +38,7 @@ import type {
   GraphQLResolveInfo,
   GraphQLTypeResolver,
   GraphQLList,
+  GraphQLNonNull,
 } from '../type/definition';
 import { assertValidSchema } from '../type/validate';
 import {
@@ -594,7 +595,8 @@ function resolveField(
     return;
   }
 
-  const returnType = fieldDef.type;
+  const returnType = fieldNode.required ? new GraphQLNonNull(fieldDef.type) : fieldDef.type
+
   const resolveFn = fieldDef.resolve ?? exeContext.fieldResolver;
 
   const info = buildResolveInfo(
diff --git a/src/language/ast.d.ts b/src/language/ast.d.ts
index 2e2850a423..331bf1f4fe 100644
--- a/src/language/ast.d.ts
+++ b/src/language/ast.d.ts
@@ -267,6 +267,7 @@ export interface FieldNode {
   readonly arguments?: ReadonlyArray<ArgumentNode>;
   readonly directives?: ReadonlyArray<DirectiveNode>;
   readonly selectionSet?: SelectionSetNode;
+  readonly required?: Boolean;
 }
 
 export interface ArgumentNode {
diff --git a/src/language/parser.js b/src/language/parser.js
index 9205382228..84b273393c 100644
--- a/src/language/parser.js
+++ b/src/language/parser.js
@@ -219,9 +219,10 @@ export class Parser {
     if (this.peek(TokenKind.NAME)) {
       switch (this._lexer.token.value) {
         case 'query':
+          return this.parseOperationDefinition(true);
         case 'mutation':
         case 'subscription':
-          return this.parseOperationDefinition();
+          return this.parseOperationDefinition(false);
         case 'fragment':
           return this.parseFragmentDefinition();
         case 'schema':
@@ -252,7 +253,7 @@ export class Parser {
    *  - SelectionSet
    *  - OperationType Name? VariableDefinitions? Directives? SelectionSet
    */
-  parseOperationDefinition(): OperationDefinitionNode {
+  parseOperationDefinition(allowRequiredField: boolean): OperationDefinitionNode {
     const start = this._lexer.token;
     if (this.peek(TokenKind.BRACE_L)) {
       return this.node(start, {
@@ -261,7 +262,7 @@ export class Parser {
         name: undefined,
         variableDefinitions: [],
         directives: [],
-        selectionSet: this.parseSelectionSet(),
+        selectionSet: this.parseSelectionSet(allowRequiredField),
       });
     }
     const operation = this.parseOperationType();
@@ -275,7 +276,7 @@ export class Parser {
       name,
       variableDefinitions: this.parseVariableDefinitions(),
       directives: this.parseDirectives(false),
-      selectionSet: this.parseSelectionSet(),
+      selectionSet: this.parseSelectionSet(allowRequiredField),
     });
   }
 
@@ -337,12 +338,17 @@ export class Parser {
   /**
    * SelectionSet : { Selection+ }
    */
-  parseSelectionSet(): SelectionSetNode {
+  parseSelectionSet(allowRequiredField: boolean): SelectionSetNode {
+    const _allowRequiredField = allowRequiredField
+    let parseSelectionFn: () => SelectionNode = function (): SelectionNode {
+      return this.parseSelection(_allowRequiredField)
+    }
+
     return this.node(this._lexer.token, {
       kind: Kind.SELECTION_SET,
       selections: this.many(
         TokenKind.BRACE_L,
-        this.parseSelection,
+        parseSelectionFn,
         TokenKind.BRACE_R,
       ),
     });
@@ -354,10 +360,10 @@ export class Parser {
    *   - FragmentSpread
    *   - InlineFragment
    */
-  parseSelection(): SelectionNode {
+  parseSelection(allowRequiredField: boolean): SelectionNode {
     return this.peek(TokenKind.SPREAD)
       ? this.parseFragment()
-      : this.parseField();
+      : this.parseField(allowRequiredField);
   }
 
   /**
@@ -365,12 +371,19 @@ export class Parser {
    *
    * Alias : Name :
    */
-  parseField(): FieldNode {
+  parseField(allowRequiredField: boolean): FieldNode {
     const start = this._lexer.token;
 
     const nameOrAlias = this.parseName();
     let alias;
     let name;
+    let required;
+    if (allowRequiredField && this.expectOptionalToken(TokenKind.BANG)) {
+      required = true
+    } else {
+      required = false
+    }
+
     if (this.expectOptionalToken(TokenKind.COLON)) {
       alias = nameOrAlias;
       name = this.parseName();
@@ -385,8 +398,9 @@ export class Parser {
       arguments: this.parseArguments(false),
       directives: this.parseDirectives(false),
       selectionSet: this.peek(TokenKind.BRACE_L)
-        ? this.parseSelectionSet()
+        ? this.parseSelectionSet(allowRequiredField)
         : undefined,
+      required: required,
     });
   }
 
diff --git a/src/language/printer.js b/src/language/printer.js
index 43bd627195..6476a787dd 100644
--- a/src/language/printer.js
+++ b/src/language/printer.js
@@ -52,9 +52,10 @@ const printDocASTReducer: any = {
   },
   SelectionSet: { leave: ({ selections }) => block(selections) },
 
+
   Field: {
-    leave({ alias, name, arguments: args, directives, selectionSet }) {
-      const prefix = wrap('', alias, ': ') + name;
+    leave({ alias, name, arguments: args, directives, selectionSet, required }) {
+      const prefix = wrap('', alias, ': ') + name + (required ? '!' : '');
       let argsLine = prefix + wrap('(', join(args, ', '), ')');
 
       if (argsLine.length > MAX_LINE_LENGTH) {
@@ -67,6 +68,7 @@ const printDocASTReducer: any = {
 
   Argument: { leave: ({ name, value }) => name + ': ' + value },
 
+
   // Fragments
 
   FragmentSpread: {