Add support for extend union syntax

Co-authored-by: Jaap van der Vis <[email protected]>
nuwave · Nov 30, 2023 · 7e4b3be · 7e4b3be
1 parent afc87a5
commit 7e4b3be
Show file tree

Hide file tree

Showing 3 changed files with 74 additions and 2 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -9,6 +9,10 @@ You can find and compare releases at the [GitHub release page](https://github.co
 
 ## Unreleased
 
+### Added
+
+- Add support for `extend union` syntax https://github.com/nuwave/lighthouse/pull/2468
+
 ## v6.23.2
 
 ### Fixed

diff --git a/src/Schema/AST/ASTBuilder.php b/src/Schema/AST/ASTBuilder.php
@@ -10,6 +10,8 @@
 use GraphQL\Language\AST\InterfaceTypeExtensionNode;
 use GraphQL\Language\AST\ObjectTypeDefinitionNode;
 use GraphQL\Language\AST\ObjectTypeExtensionNode;
+use GraphQL\Language\AST\UnionTypeDefinitionNode;
+use GraphQL\Language\AST\UnionTypeExtensionNode;
 use GraphQL\Language\Parser;
 use Illuminate\Contracts\Events\Dispatcher as EventsDispatcher;
 use Illuminate\Support\Arr;
@@ -31,6 +33,7 @@ class ASTBuilder
         InputObjectTypeExtensionNode::class => InputObjectTypeDefinitionNode::class,
         InterfaceTypeExtensionNode::class => InterfaceTypeDefinitionNode::class,
         EnumTypeExtensionNode::class => EnumTypeDefinitionNode::class,
+        UnionTypeExtensionNode::class => UnionTypeDefinitionNode::class,
     ];
 
     /** Initialized lazily in $this->documentAST(). */
@@ -119,6 +122,8 @@ protected function applyTypeExtensionManipulators(): void
                     $this->extendObjectLikeType($typeName, $typeExtension);
                 } elseif ($typeExtension instanceof EnumTypeExtensionNode) {
                     $this->extendEnumType($typeName, $typeExtension);
+                } elseif ($typeExtension instanceof UnionTypeExtensionNode) {
+                    $this->extendUnionType($typeName, $typeExtension);
                 }
             }
         }
@@ -173,12 +178,26 @@ protected function extendEnumType(string $typeName, EnumTypeExtensionNode $typeE
         );
     }
 
-    protected function missingBaseDefinition(string $typeName, ObjectTypeExtensionNode|InputObjectTypeExtensionNode|InterfaceTypeExtensionNode|EnumTypeExtensionNode $typeExtension): string
+    protected function extendUnionType(string $typeName, UnionTypeExtensionNode $typeExtension): void
+    {
+        $extendedUnion = $this->documentAST->types[$typeName]
+            ?? throw new DefinitionException($this->missingBaseDefinition($typeName, $typeExtension));
+        assert($extendedUnion instanceof UnionTypeDefinitionNode);
+
+        $this->assertExtensionMatchesDefinition($typeExtension, $extendedUnion);
+
+        $extendedUnion->types = ASTHelper::mergeUniqueNodeList(
+            $extendedUnion->types,
+            $typeExtension->types,
+        );
+    }
+
+    protected function missingBaseDefinition(string $typeName, ObjectTypeExtensionNode|InputObjectTypeExtensionNode|InterfaceTypeExtensionNode|EnumTypeExtensionNode|UnionTypeExtensionNode $typeExtension): string
     {
         return "Could not find a base definition {$typeName} of kind {$typeExtension->kind} to extend.";
     }
 
-    protected function assertExtensionMatchesDefinition(ObjectTypeExtensionNode|InputObjectTypeExtensionNode|InterfaceTypeExtensionNode|EnumTypeExtensionNode $extension, ObjectTypeDefinitionNode|InputObjectTypeDefinitionNode|InterfaceTypeDefinitionNode|EnumTypeDefinitionNode $definition): void
+    protected function assertExtensionMatchesDefinition(ObjectTypeExtensionNode|InputObjectTypeExtensionNode|InterfaceTypeExtensionNode|EnumTypeExtensionNode|UnionTypeExtensionNode $extension, ObjectTypeDefinitionNode|InputObjectTypeDefinitionNode|InterfaceTypeDefinitionNode|EnumTypeDefinitionNode|UnionTypeDefinitionNode $definition): void
     {
         if (static::EXTENSION_TO_DEFINITION_CLASS[$extension::class] !== $definition::class) {
             throw new DefinitionException("The type extension {$extension->name->value} of kind {$extension->kind} can not extend a definition of kind {$definition->kind}.");

diff --git a/tests/Unit/Schema/AST/ASTBuilderTest.php b/tests/Unit/Schema/AST/ASTBuilderTest.php
@@ -7,6 +7,7 @@
 use GraphQL\Language\AST\InterfaceTypeDefinitionNode;
 use GraphQL\Language\AST\NodeKind;
 use GraphQL\Language\AST\ObjectTypeDefinitionNode;
+use GraphQL\Language\AST\UnionTypeDefinitionNode;
 use Illuminate\Support\Collection;
 use Nuwave\Lighthouse\Exceptions\DefinitionException;
 use Nuwave\Lighthouse\Schema\AST\ASTBuilder;
@@ -152,6 +153,27 @@ enum MyEnum {
         $this->assertCount(4, $myEnum->values);
     }
 
+    public function testMergeUnionExtensionFields(): void
+    {
+        $this->schema = /** @lang GraphQL */ '
+            type Foo
+            type Bar
+            type Baz
+
+            union MyUnion = Foo
+
+            extend union MyUnion = Bar
+
+            extend union MyUnion = Baz
+        ';
+        $documentAST = $this->astBuilder->documentAST();
+
+        $myUnion = $documentAST->types['MyUnion'];
+        assert($myUnion instanceof UnionTypeDefinitionNode);
+
+        $this->assertCount(3, $myUnion->types);
+    }
+
     public function testDoesNotAllowExtendingUndefinedTypes(): void
     {
         $this->schema = /** @lang GraphQL */ '
@@ -168,6 +190,18 @@ public function testDoesNotAllowExtendingUndefinedTypes(): void
         $this->astBuilder->documentAST();
     }
 
+    public function testDoesNotAllowExtendingUndefinedUnions(): void
+    {
+        $this->schema = /** @lang GraphQL */ '
+        union MyFirstEnum = String
+
+        extend union MySecondUnion = Int
+        ';
+
+        $this->expectExceptionObject(new DefinitionException('Could not find a base definition MySecondUnion of kind ' . NodeKind::UNION_TYPE_EXTENSION . ' to extend.'));
+        $this->astBuilder->documentAST();
+    }
+
     public function testDoesNotAllowDuplicateFieldsOnTypeExtensions(): void
     {
         $this->schema = /** @lang GraphQL */ '
@@ -234,6 +268,21 @@ enum MyEnum {
         $this->astBuilder->documentAST();
     }
 
+    public function testDoesNotAllowDuplicateTypesOnUnionExtensions(): void
+    {
+        $this->schema = /** @lang GraphQL */ '
+        type Foo
+        type Bar
+
+        union MyUnion = Foo | Bar
+
+        extend union MyUnion = Bar
+        ';
+
+        $this->expectExceptionObject(new DefinitionException(ASTHelper::duplicateDefinition('Bar')));
+        $this->astBuilder->documentAST();
+    }
+
     public function testDoesNotAllowMergingNonMatchingTypes(): void
     {
         $this->schema = /** @lang GraphQL */ '