feat: entity type customization, type loader support (#33)

* feat: entity type customization, type loader support

* feat: lazy union type loading

* Update README.md

Co-authored-by: Tom Carrio <tom@carrio.dev>

* fix: type docs updating

---------

Co-authored-by: Tom Carrio <tom@carrio.dev>

Author: jpasquers

README.md

@@ -105,6 +105,12 @@ $query = 'query GetServiceSDL { _service { sdl } }';
 $result = GraphQL::executeQuery($schema, $query);
 ```
 
+#### Config
+
+The config parameter for the `FederatedSchema` object is entirely compatible with the `Schema` config argument. On top of this, we support the following optional parameters:
+
+1. `entityTypes` - the entity types (which extend `EntityObjectType`) which will form the `_Entity` type on the federated schema. If not provided, the Schema will scan the `Query` type tree for all types extending `EntityObjectType`.
+
 ## Disclaimer
 
 Documentation in this project include content quoted directly from the [Apollo official documentation](https://www.apollographql.com/docs) to reduce redundancy.

src/FederatedSchema.php

@@ -4,17 +4,17 @@
 
 namespace Apollo\Federation;
 
+use Apollo\Federation\Types\AnyType;
 use GraphQL\Type\Schema;
-use GraphQL\Type\Definition\CustomScalarType;
 use GraphQL\Type\Definition\Directive;
 use GraphQL\Type\Definition\ObjectType;
-use GraphQL\Type\Definition\UnionType;
 use GraphQL\Type\Definition\Type;
 use GraphQL\Utils\TypeInfo;
 use GraphQL\Utils\Utils;
 
 use Apollo\Federation\Types\EntityObjectType;
-use Apollo\Federation\Utils\FederatedSchemaPrinter;
+use Apollo\Federation\Types\EntityUnionType;
+use Apollo\Federation\Types\ServiceDefinitionType;
 
 /**
  * A federated GraphQL schema definition (see [related docs](https://www.apollographql.com/docs/apollo-server/federation/introduction))
@@ -55,18 +55,36 @@
  */
 class FederatedSchema extends Schema
 {
-    /** @var EntityObjectType[] */
+    /** @var EntityObjectType[]|callable: EntityObjectType[] */
     protected $entityTypes;
 
     /** @var Directive[] */
     protected $entityDirectives;
 
+    protected ServiceDefinitionType $serviceDefinitionType;
+    protected EntityUnionType $entityUnionType;
+    protected AnyType $anyType;
+
+    /**
+     * 
+     * We will provide the parts that we need to operate against.
+     * 
+     * @param array{?entityTypes: array<EntityObjectType>, ?typeLoader: callable, query: array} $config
+     */
     public function __construct($config)
     {
-        $this->entityTypes = $this->extractEntityTypes($config);
+        $this->entityTypes = $config['entityTypes'] ?? $this->lazyEntityTypeExtractor($config);
         $this->entityDirectives = array_merge(Directives::getDirectives(), Directive::getInternalDirectives());
-
-        $config = array_merge($config, $this->getEntityDirectivesConfig($config), $this->getQueryTypeConfig($config));
+        
+        $this->serviceDefinitionType = new ServiceDefinitionType($this);
+        $this->entityUnionType = new EntityUnionType($this->entityTypes);
+        $this->anyType = new AnyType();
+
+        $config = array_merge($config, 
+            $this->getEntityDirectivesConfig($config), 
+            $this->getQueryTypeConfig($config),
+            $this->supplementTypeLoader($config)
+        );
 
         parent::__construct($config);
     }
@@ -78,7 +96,9 @@ public function __construct($config)
      */
     public function getEntityTypes(): array
     {
-        return $this->entityTypes;
+        return is_callable($this->entityTypes)
+            ? ($this->entityTypes)()
+            : $this->entityTypes;
     }
 
     /**
@@ -121,24 +141,42 @@ private function getQueryTypeConfig(array $config): array
         ];
     }
 
+    /**
+     * Add type loading functionality for the types required for the federated schema to function.
+     */
+    private function supplementTypeLoader(array $config): array
+    {
+        if (!array_key_exists('typeLoader', $config) || !is_callable($config['typeLoader'])) {
+            return [];
+        }
+
+        return [
+            'typeLoader' => function ($typeName) use ($config) {
+                $map = $this->builtInTypeMap();
+                if (array_key_exists($typeName, $map)) {
+                    return $map[$typeName];
+                }
+
+                return $config['typeLoader']($typeName);
+            }
+        ];
+    }
+
+    private function builtInTypeMap(): array 
+    {
+        return [
+            EntityUnionType::getTypeName() => $this->entityUnionType,
+            ServiceDefinitionType::getTypeName() => $this->serviceDefinitionType,
+            AnyType::getTypeName() => $this->anyType
+        ];
+    }
+
     /** @var array */
     private function getQueryTypeServiceFieldConfig(): array
     {
-        $serviceType = new ObjectType([
-            'name' => '_Service',
-            'fields' => [
-                'sdl' => [
-                    'type' => Type::string(),
-                    'resolve' => function () {
-                        return FederatedSchemaPrinter::doPrint($this);
-                    }
-                ]
-            ]
-        ]);
-
         return [
             '_service' => [
-                'type' => Type::nonNull($serviceType),
+                'type' => Type::nonNull($this->serviceDefinitionType),
                 'resolve' => function () {
                     return [];
                 }
@@ -149,28 +187,12 @@ private function getQueryTypeServiceFieldConfig(): array
     /** @var array */
     private function getQueryTypeEntitiesFieldConfig(?array $config): array
     {
-        if (!$this->hasEntityTypes()) {
-            return [];
-        }
-
-        $entityType = new UnionType([
-            'name' => '_Entity',
-            'types' => array_values($this->getEntityTypes())
-        ]);
-
-        $anyType = new CustomScalarType([
-            'name' => '_Any',
-            'serialize' => function ($value) {
-                return $value;
-            }
-        ]);
-
         return [
             '_entities' => [
-                'type' => Type::listOf($entityType),
+                'type' => Type::listOf($this->entityUnionType),
                 'args' => [
                     'representations' => [
-                        'type' => Type::nonNull(Type::listOf(Type::nonNull($anyType)))
+                        'type' => Type::nonNull(Type::listOf(Type::nonNull($this->anyType)))
                     ]
                 ],
                 'resolve' => function ($root, $args, $context, $info) use ($config) {
@@ -208,22 +230,25 @@ private function resolve($root, $args, $context, $info)
             return $r;
         }, $args['representations']);
     }
+
     /**
      * @param array $config
      *
-     * @return EntityObjectType[]
+     * @return callable: EntityObjectType[]
      */
-    private function extractEntityTypes(array $config): array
+    private function lazyEntityTypeExtractor(array $config): callable
     {
-        $resolvedTypes = TypeInfo::extractTypes($config['query']);
-        $entityTypes = [];
-
-        foreach ($resolvedTypes as $type) {
-            if ($type instanceof EntityObjectType) {
-                $entityTypes[$type->name] = $type;
+        return function () use ($config) {
+            $resolvedTypes = TypeInfo::extractTypes($config['query']);
+            $entityTypes = [];
+    
+            foreach ($resolvedTypes as $type) {
+                if ($type instanceof EntityObjectType) {
+                    $entityTypes[$type->name] = $type;
+                }
             }
-        }
 
-        return $entityTypes;
+            return $entityTypes;
+        };
     }
 }

src/Types/AnyType.php

@@ -0,0 +1,29 @@
+<?php
+
+declare(strict_types=1);
+
+namespace Apollo\Federation\Types;
+
+use GraphQL\Type\Definition\CustomScalarType;
+
+/**
+ * Simple representation of an agnostic scalar value.
+ */
+class AnyType extends CustomScalarType
+{
+    public function __construct()
+    {
+        $config = [
+            'name' => self::getTypeName(),
+            'serialize' => function ($value) {
+                return $value;
+            }
+        ];
+        parent::__construct($config);
+    }
+
+    public static function getTypeName(): string
+    {
+        return '_Any';
+    }
+}

src/Types/EntityUnionType.php

@@ -0,0 +1,34 @@
+<?php
+
+declare(strict_types=1);
+
+namespace Apollo\Federation\Types;
+
+use GraphQL\Type\Definition\UnionType;
+
+/**
+ * The union of all entities defined within this schema.
+ */
+class EntityUnionType extends UnionType
+{
+
+    /**
+     * @param array|callable $entityTypes all entity types or a callable to retrieve them
+     */
+    public function __construct($entityTypes)
+    {
+        $config = [
+            'name' => self::getTypeName(),
+            'types' => is_callable($entityTypes) 
+                ? fn () => array_values($entityTypes())
+                : array_values($entityTypes)
+                
+        ];
+        parent::__construct($config);
+    }
+
+    public static function getTypeName(): string
+    {
+        return '_Entity';
+    }
+}

src/Types/ServiceDefinitionType.php

@@ -0,0 +1,39 @@
+<?php
+
+declare(strict_types=1);
+
+namespace Apollo\Federation\Types;
+
+use Apollo\Federation\Utils\FederatedSchemaPrinter;
+use GraphQL\Type\Definition\ObjectType;
+use GraphQL\Type\Definition\Type;
+use GraphQL\Type\Schema;
+
+/**
+ * The type of the service definition required for federated schemas.
+ */
+class ServiceDefinitionType extends ObjectType
+{
+
+    /**
+     * @param Schema $schema - the schemas whose SDL should be printed.
+     */
+    public function __construct(Schema $schema)
+    {
+        $config = [
+            'name' => self::getTypeName(),
+            'fields' => [
+                'sdl' => [
+                    'type' => Type::string(),
+                    'resolve' => fn () => FederatedSchemaPrinter::doPrint($schema)
+                ]     
+            ]
+        ];
+        parent::__construct($config);
+    }
+
+    public static function getTypeName(): string
+    {
+        return '_Service';
+    }
+}

test/SchemaTest.php

@@ -8,11 +8,9 @@
 use Spatie\Snapshots\MatchesSnapshots;
 
 use GraphQL\GraphQL;
-use GraphQL\Type\Definition\Type;
 use GraphQL\Utils\SchemaPrinter;
 
 use Apollo\Federation\Tests\StarWarsSchema;
-use Apollo\Federation\Tests\DungeonsAndDragonsSchema;
 
 class SchemaTest extends TestCase
 {
@@ -88,6 +86,14 @@ public function testSchemaSdl()
         $this->assertMatchesSnapshot($schemaSdl);
     }
 
+    public function testSchemaSdlForProvidedEntities()
+    {
+        $schema = StarWarsSchema::getEpisodesSchemaWithProvidedEntities();
+        $schemaSdl = SchemaPrinter::doPrint($schema);
+
+        $this->assertMatchesSnapshot($schemaSdl);
+    }
+
     public function testResolvingEntityReferences()
     {
         $schema = StarWarsSchema::getEpisodesSchema();