Parse schema extensions

6 files changed, 44 insertions, 10 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 22da347..7d5e230 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -11,6 +11,7 @@ and this project adheres to
 - AST for the GraphQL schema.
 - Parser for the TypeSystemDefinition.
 - `Trans.argument`.
+- Schema extension parser.
 
 ### Changed
 - Rename `AST.Definition` into `AST.Document.ExecutableDefinition`.
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
diff --git a/tests/Language/GraphQL/AST/LexerSpec.hs b/tests/Language/GraphQL/AST/LexerSpec.hs
index 9b5d6aa..4938b0f 100644
--- a/tests/Language/GraphQL/AST/LexerSpec.hs
+++ b/tests/Language/GraphQL/AST/LexerSpec.hs
@@ -8,7 +8,7 @@ import Data.Text (Text)
 import Data.Void (Void)
 import Language.GraphQL.AST.Lexer
 import Test.Hspec (Spec, context, describe, it)
-import Test.Hspec.Megaparsec (shouldParse, shouldSucceedOn)
+import Test.Hspec.Megaparsec (shouldParse, shouldFailOn, shouldSucceedOn)
 import Text.Megaparsec (ParseErrorBundle, parse)
 import Text.RawString.QQ (r)
 
@@ -89,6 +89,8 @@ spec = describe "Lexer" $ do
             parse amp "" "&" `shouldParse` "&"
         it "lexes schema extensions" $
             parse (extend "schema") "" `shouldSucceedOn` "extend schema"
+        it "fails if the given token doesn't match" $
+            parse (extend "schema") "" `shouldFailOn` "extend shema"
 
 runBetween :: (Parser () -> Parser ()) -> Text -> Either (ParseErrorBundle Text Void) ()
 runBetween parser = parse (parser $ pure ()) ""
diff --git a/tests/Language/GraphQL/AST/ParserSpec.hs b/tests/Language/GraphQL/AST/ParserSpec.hs
index 2e55389..f06f6c1 100644
--- a/tests/Language/GraphQL/AST/ParserSpec.hs
+++ b/tests/Language/GraphQL/AST/ParserSpec.hs
@@ -4,9 +4,11 @@ module Language.GraphQL.AST.ParserSpec
     ( spec
     ) where
 
+import Data.List.NonEmpty (NonEmpty(..))
+import Language.GraphQL.AST.Document
 import Language.GraphQL.AST.Parser
 import Test.Hspec (Spec, describe, it)
-import Test.Hspec.Megaparsec (shouldSucceedOn)
+import Test.Hspec.Megaparsec (shouldParse, shouldSucceedOn)
 import Text.Megaparsec (parse)
 import Text.RawString.QQ (r)
 
@@ -116,3 +118,20 @@ spec = describe "Parser" $ do
               | FIELD
               | FRAGMENT_SPREAD
         |]
+
+    it "parses schema extension with a new directive" $
+        parse document "" `shouldSucceedOn`[r|
+            extend schema @newDirective
+        |]
+
+    it "parses schema extension with an operation type definition" $
+        parse document "" `shouldSucceedOn` [r|extend schema { query: Query }|]
+
+    it "parses schema extension with an operation type and directive" $
+        let newDirective = Directive "newDirective" []
+            testSchemaExtension = TypeSystemExtension
+                $ SchemaExtension
+                $ SchemaOperationExtension [newDirective]
+                $ OperationTypeDefinition Query "Query" :| []
+            query = [r|extend schema @newDirective { query: Query }|]
+         in parse document "" query `shouldParse` (testSchemaExtension :| [])