Parse schema extensions

3 files changed, 20 insertions, 8 deletions

diff --git a/src/Language/GraphQL/AST/Document.hs b/src/Language/GraphQL/AST/Document.hs
index e3fe78c..6ed5f50 100644
--- a/src/Language/GraphQL/AST/Document.hs
+++ b/src/Language/GraphQL/AST/Document.hs
@@ -25,7 +25,6 @@ module Language.GraphQL.AST.Document
     , OperationDefinition(..)
     , OperationType(..)
     , OperationTypeDefinition(..)
-    , OperationTypeDefinitions
     , SchemaExtension(..)
     , Selection(..)
     , SelectionSet
@@ -247,7 +246,7 @@ data Directive = Directive Name [Argument] deriving (Eq, Show)
 -- * Type System
 
 data TypeSystemDefinition
-    = SchemaDefinition [Directive] OperationTypeDefinitions
+    = SchemaDefinition [Directive] (NonEmpty OperationTypeDefinition)
     | TypeDefinition TypeDefinition
     | DirectiveDefinition
         Description Name ArgumentsDefinition (NonEmpty DirectiveLocation)
@@ -262,14 +261,12 @@ data TypeSystemExtension
 
 -- ** Schema
 
-type OperationTypeDefinitions = NonEmpty OperationTypeDefinition
-
 data OperationTypeDefinition
     = OperationTypeDefinition OperationType NamedType
     deriving (Eq, Show)
 
 data SchemaExtension
-    = SchemaOperationExtension [Directive] OperationTypeDefinitions
+    = SchemaOperationExtension [Directive] (NonEmpty OperationTypeDefinition)
     | SchemaDirectiveExtension (NonEmpty Directive)
     deriving (Eq, Show)
 
diff --git a/src/Language/GraphQL/AST/Lexer.hs b/src/Language/GraphQL/AST/Lexer.hs
index 7303fdf..e119303 100644
--- a/src/Language/GraphQL/AST/Lexer.hs
+++ b/src/Language/GraphQL/AST/Lexer.hs
@@ -32,7 +32,6 @@ module Language.GraphQL.AST.Lexer
 import Control.Applicative (Alternative(..), liftA2)
 import Data.Char (chr, digitToInt, isAsciiLower, isAsciiUpper, ord)
 import Data.Foldable (foldl')
-import Data.Functor (($>))
 import Data.List (dropWhileEnd)
 import Data.Proxy (Proxy(..))
 import Data.Void (Void)
@@ -222,4 +221,4 @@ unicodeBOM = optional (char '\xfeff') >> pure ()
 
 -- | Parses "extend" followed by a 'symbol'. It is used by schema extensions.
 extend :: Text -> Parser ()
-extend token = symbol "extend" $> extend token >> pure ()
+extend token = symbol "extend" *> symbol token >> pure ()
diff --git a/src/Language/GraphQL/AST/Parser.hs b/src/Language/GraphQL/AST/Parser.hs
index 33deb15..a750651 100644
--- a/src/Language/GraphQL/AST/Parser.hs
+++ b/src/Language/GraphQL/AST/Parser.hs
@@ -30,6 +30,7 @@ document = unicodeBOM
 definition :: Parser Definition
 definition = ExecutableDefinition <$> executableDefinition
     <|> TypeSystemDefinition <$> typeSystemDefinition
+    <|> TypeSystemExtension <$> typeSystemExtension
     <?> "Definition"
 
 executableDefinition :: Parser ExecutableDefinition
@@ -43,6 +44,10 @@ typeSystemDefinition = schemaDefinition
     <|> directiveDefinition
     <?> "TypeSystemDefinition"
 
+typeSystemExtension :: Parser TypeSystemExtension
+typeSystemExtension = SchemaExtension <$> schemaExtension
+    <?> "TypeSystemExtension"
+
 directiveDefinition :: Parser TypeSystemDefinition
 directiveDefinition = DirectiveDefinition
     <$> description
@@ -214,8 +219,19 @@ schemaDefinition = SchemaDefinition
     <*> directives
     <*> operationTypeDefinitions
     <?> "SchemaDefinition"
+
+operationTypeDefinitions :: Parser (NonEmpty OperationTypeDefinition)
+operationTypeDefinitions = braces $ NonEmpty.some operationTypeDefinition
+
+schemaExtension :: Parser SchemaExtension
+schemaExtension = extend "schema"
+    >> try schemaOperationExtension
+    <|> SchemaDirectiveExtension <$> NonEmpty.some directive
+    <?> "SchemaExtension"
   where
-    operationTypeDefinitions = braces $ NonEmpty.some operationTypeDefinition
+    schemaOperationExtension = SchemaOperationExtension
+        <$> directives
+        <*> operationTypeDefinitions
 
 operationTypeDefinition :: Parser OperationTypeDefinition
 operationTypeDefinition = OperationTypeDefinition