Fix N+1 Query Issue with Nested Eager Loading on Union Types

src/Support/SelectFields.php

@@ -207,14 +207,29 @@
 
                         static::addAlwaysFields($fieldObject, $field, $parentTable, true);
 
-                        $with[$relationsKey] = static::getSelectableFieldsAndRelations(
-                            $queryArgs,
-                            $field,
-                            $newParentType,
-                            $customQuery,
-                            false,
-                            $ctx
-                        );
+                        // Check if this is a MorphTo relation
+                        if (is_a($relation, MorphTo::class)) {
+                            // For MorphTo relations, we need to handle them specially
+                            // because they can have different models, and we need to eager load based on the query
+                            static::handleMorphToRelation(
+                                $queryArgs,
+                                $field,
+                                $with,
+                                $ctx,
+                                $fieldObject,
+                                $key,
+                                $customQuery
+                            );
+                        } else {
+                            $with[$relationsKey] = static::getSelectableFieldsAndRelations(
+                                $queryArgs,
+                                $field,
+                                $newParentType,
+                                $customQuery,
+                                false,
+                                $ctx
+                            );
+                        }
                     } elseif (is_a($parentTypeUnwrapped, \GraphQL\Type\Definition\InterfaceType::class)) {
                         static::handleInterfaceFields(
                             $queryArgs,
@@ -307,7 +322,7 @@
         $selectable = true;
 
         // If not a selectable field
         if (isset($fieldObject->config['selectable']) && false === $fieldObject->config['selectable']) {
             $selectable = false;
         }
 
@@ -518,6 +533,65 @@
         };
     }
 
+    /**
+     * Handle MorphTo relations
+     * @param mixed $ctx
+     */
+    protected static function handleMorphToRelation(
+        array $queryArgs,
+        array $field,
+        array &$with,
+        $ctx,
+        FieldDefinition $fieldObject,
+        string $key,
+        ?Closure $customQuery
+    ): void {
+        $relationsKey = Arr::get($fieldObject->config, 'alias', $key);
+
+        /* @var GraphqlType $fieldType */
+        $fieldType = $fieldObject->config['type'];
+
+        if ($fieldType instanceof WrappingType) {
+            $fieldType = $fieldType->getInnermostType();
+        }
+
+        /** @var UnionType $union */
+        $union = $fieldType;
+
+        $relationNames = (isset($union->config['relationName']) && \is_callable($union->config['relationName']))
+            ? $union->config['relationName']()
+            : null;
+        $with[$relationsKey] = function (MorphTo $relation) use ($queryArgs, $field, $union, $relationNames, $customQuery, $ctx): void {
+            $morphRelation = [];
+
+            foreach ($union->getTypes() as $unionType) {
+                // Get the model class name for the morph type
+                if (isset($unionType->config['model'])) {
+                    $modelClass = $unionType->config['model'];
+                } else {
+                    // Fallback to type name if no model is configured
+                    $modelClass = $relationNames[$unionType->name()] ?? $unionType->name();
+                }
+
+                /** @var callable $callable */
+                $callable = static::getSelectableFieldsAndRelations(
+                    $queryArgs,
+                    $field,
+                    $unionType,
+                    $customQuery,
+                    false,
+                    $ctx
+                );
+
+                $morphRelation[$modelClass] = $callable;
+            }
+
+            if (!empty($morphRelation)) {
+                $relation->constrain($morphRelation);
+            }
+        };
+    }
+
     public function getSelect(): array
     {
         return $this->select;

src/Support/UnionType.php

@@ -30,6 +30,10 @@ public function getAttributes(): array
             $attributes['resolveType'] = [$this, 'resolveType'];
         }
 
+        if (method_exists($this, 'relationName')) {
+            $attributes['relationName'] = [$this, 'relationName'];
+        }
+
         return $attributes;
     }
 

tests/Database/SelectFields/UnionMorphTests/CommentType.php

@@ -0,0 +1,38 @@
+<?php
+
+declare(strict_types = 1);
+namespace Rebing\GraphQL\Tests\Database\SelectFields\UnionMorphTests;
+
+use GraphQL\Type\Definition\Type;
+use Rebing\GraphQL\Support\Facades\GraphQL;
+use Rebing\GraphQL\Support\Type as GraphQLType;
+use Rebing\GraphQL\Tests\Support\Models\Comment;
+
+class CommentType extends GraphQLType
+{
+    protected $attributes = [
+        'name' => 'Comment',
+        'model' => Comment::class,
+    ];
+
+    public function fields(): array
+    {
+        return [
+            'id' => [
+                'type' => Type::nonNull(Type::id()),
+            ],
+            'title' => [
+                'type' => Type::nonNull(Type::string()),
+            ],
+            'body' => [
+                'type' => Type::string(),
+            ],
+            'file' => [
+                'type' => GraphQL::type('File'),
+            ],
+            'commentable' => [
+                'type' => GraphQL::type('CommentableUnion'),
+            ],
+        ];
+    }
+}

tests/Database/SelectFields/UnionMorphTests/CommentableUnionType.php

@@ -0,0 +1,49 @@
+<?php
+
+declare(strict_types = 1);
+namespace Rebing\GraphQL\Tests\Database\SelectFields\UnionMorphTests;
+
+use GraphQL\Type\Definition\Type as GraphqlType;
+use Rebing\GraphQL\Support\Facades\GraphQL;
+use Rebing\GraphQL\Support\UnionType;
+use Rebing\GraphQL\Tests\Support\Models\Post;
+use Rebing\GraphQL\Tests\Support\Models\Product;
+
+class CommentableUnionType extends UnionType
+{
+    protected $attributes = [
+        'name' => 'CommentableUnion',
+    ];
+
+    public function types(): array
+    {
+        return [
+            GraphQL::type('Post'),
+            GraphQL::type('Product'),
+        ];
+    }
+
+    public function relationName(): array
+    {
+        return [
+            Post::class => 'post',
+            Product::class => 'product',
+        ];
+    }
+
+    /**
+     * @param object $value
+     */
+    public function resolveType($value): ?GraphqlType
+    {
+        if ($value instanceof Post) {
+            return GraphQL::type('Post');
+        }
+
+        if ($value instanceof Product) {
+            return GraphQL::type('Product');
+        }
+
+        return null;
+    }
+}

tests/Database/SelectFields/UnionMorphTests/CommentsQuery.php

@@ -0,0 +1,38 @@
+<?php
+
+declare(strict_types = 1);
+namespace Rebing\GraphQL\Tests\Database\SelectFields\UnionMorphTests;
+
+use Closure;
+use GraphQL\Type\Definition\ResolveInfo;
+use GraphQL\Type\Definition\Type;
+use Rebing\GraphQL\Support\Facades\GraphQL;
+use Rebing\GraphQL\Support\Query;
+use Rebing\GraphQL\Support\SelectFields;
+use Rebing\GraphQL\Tests\Support\Models\Comment;
+
+class CommentsQuery extends Query
+{
+    protected $attributes = [
+        'name' => 'comments',
+    ];
+
+    public function type(): Type
+    {
+        return Type::listOf(GraphQL::type('Comment'));
+    }
+
+    public function resolve($root, $args, $context, ResolveInfo $info, Closure $getSelectFields)
+    {
+        /** @var SelectFields $selectFields */
+        $selectFields = $getSelectFields();
+
+        /** @var Comment[] $comments */
+        $comments = Comment::query()
+            ->select($selectFields->getSelect())
+            ->with($selectFields->getRelations())
+            ->get();
+
+        return $comments;
+    }
+}

tests/Database/SelectFields/UnionMorphTests/FileType.php

@@ -0,0 +1,35 @@
+<?php
+
+declare(strict_types = 1);
+namespace Rebing\GraphQL\Tests\Database\SelectFields\UnionMorphTests;
+
+use GraphQL\Type\Definition\Type;
+use Rebing\GraphQL\Support\Facades\GraphQL;
+use Rebing\GraphQL\Support\Type as GraphQLType;
+use Rebing\GraphQL\Tests\Support\Models\File;
+
+class FileType extends GraphQLType
+{
+    protected $attributes = [
+        'name' => 'File',
+        'model' => File::class,
+    ];
+
+    public function fields(): array
+    {
+        return [
+            'id' => [
+                'type' => Type::nonNull(Type::id()),
+            ],
+            'name' => [
+                'type' => Type::string(),
+            ],
+            'path' => [
+                'type' => Type::string(),
+            ],
+            'folder' => [
+                'type' => GraphQL::type('Folder'),
+            ],
+        ];
+    }
+}

tests/Database/SelectFields/UnionMorphTests/FolderType.php

@@ -0,0 +1,28 @@
+<?php
+
+declare(strict_types = 1);
+namespace Rebing\GraphQL\Tests\Database\SelectFields\UnionMorphTests;
+
+use GraphQL\Type\Definition\Type;
+use Rebing\GraphQL\Support\Type as GraphQLType;
+use Rebing\GraphQL\Tests\Support\Models\Folder;
+
+class FolderType extends GraphQLType
+{
+    protected $attributes = [
+        'name' => 'Folder',
+        'model' => Folder::class,
+    ];
+
+    public function fields(): array
+    {
+        return [
+            'id' => [
+                'type' => Type::nonNull(Type::id()),
+            ],
+            'name' => [
+                'type' => Type::string(),
+            ],
+        ];
+    }
+}

tests/Database/SelectFields/UnionMorphTests/PostType.php

@@ -0,0 +1,32 @@
+<?php
+
+declare(strict_types = 1);
+namespace Rebing\GraphQL\Tests\Database\SelectFields\UnionMorphTests;
+
+use GraphQL\Type\Definition\Type;
+use Rebing\GraphQL\Support\Facades\GraphQL;
+use Rebing\GraphQL\Support\Type as GraphQLType;
+use Rebing\GraphQL\Tests\Support\Models\Post;
+
+class PostType extends GraphQLType
+{
+    protected $attributes = [
+        'name' => 'Post',
+        'model' => Post::class,
+    ];
+
+    public function fields(): array
+    {
+        return [
+            'id' => [
+                'type' => Type::nonNull(Type::id()),
+            ],
+            'title' => [
+                'type' => Type::string(),
+            ],
+            'file' => [
+                'type' => GraphQL::type('File'),
+            ],
+        ];
+    }
+}

tests/Database/SelectFields/UnionMorphTests/ProductType.php

@@ -0,0 +1,35 @@
+<?php
+
+declare(strict_types = 1);
+namespace Rebing\GraphQL\Tests\Database\SelectFields\UnionMorphTests;
+
+use GraphQL\Type\Definition\Type;
+use Rebing\GraphQL\Support\Facades\GraphQL;
+use Rebing\GraphQL\Support\Type as GraphQLType;
+use Rebing\GraphQL\Tests\Support\Models\Product;
+
+class ProductType extends GraphQLType
+{
+    protected $attributes = [
+        'name' => 'Product',
+        'model' => Product::class,
+    ];
+
+    public function fields(): array
+    {
+        return [
+            'id' => [
+                'type' => Type::nonNull(Type::id()),
+            ],
+            'name' => [
+                'type' => Type::string(),
+            ],
+            'price' => [
+                'type' => Type::float(),
+            ],
+            'file' => [
+                'type' => GraphQL::type('File'),
+            ],
+        ];
+    }
+}