From eb467d63f552f68529bb6476a0c1efa18d34f28f Mon Sep 17 00:00:00 2001
From: Klaus Purer <klaus.purer@protonmail.ch>
Date: Wed, 20 Dec 2023 13:57:44 +0100
Subject: [PATCH] fix(schema): Performance improvements for better AST caching
 (#1379)

---
 .../Schema/AlterableComposableSchema.php      | 10 +--
 .../GraphQL/Schema/SdlSchemaPluginBase.php    | 75 ++++++++++++++-----
 2 files changed, 56 insertions(+), 29 deletions(-)

diff --git a/src/Plugin/GraphQL/Schema/AlterableComposableSchema.php b/src/Plugin/GraphQL/Schema/AlterableComposableSchema.php
index 0f50a9f12..e6d659c0e 100644
--- a/src/Plugin/GraphQL/Schema/AlterableComposableSchema.php
+++ b/src/Plugin/GraphQL/Schema/AlterableComposableSchema.php
@@ -157,12 +157,6 @@ protected function getSchemaDocument(array $extensions = []) {
    * @see \Drupal\graphql\Plugin\GraphQL\Schema\ComposableSchema::getSchemaDocument()
    */
   protected function getExtensionDocument(array $extensions = []) {
-    // Only use caching of the parsed document if we aren't in development mode.
-    $cid = "extension:{$this->getPluginId()}";
-    if (empty($this->inDevelopment) && $cache = $this->astCache->get($cid)) {
-      return $cache->data;
-    }
-
     $extensions = array_filter(array_map(function (SchemaExtensionPluginInterface $extension) {
       return $extension->getExtensionDefinition();
     }, $extensions), function ($definition) {
@@ -176,10 +170,8 @@ protected function getExtensionDocument(array $extensions = []) {
       AlterSchemaExtensionDataEvent::EVENT_NAME
     );
     $ast = !empty($extensions) ? Parser::parse(implode("\n\n", $event->getSchemaExtensionData()), ['noLocation' => TRUE]) : NULL;
-    if (empty($this->inDevelopment)) {
-      $this->astCache->set($cid, $ast, CacheBackendInterface::CACHE_PERMANENT, ['graphql']);
-    }
 
+    // No AST caching here as that will be done in getFullSchemaDocument().
     return $ast;
   }
 
diff --git a/src/Plugin/GraphQL/Schema/SdlSchemaPluginBase.php b/src/Plugin/GraphQL/Schema/SdlSchemaPluginBase.php
index cae50d5d0..f97f69165 100644
--- a/src/Plugin/GraphQL/Schema/SdlSchemaPluginBase.php
+++ b/src/Plugin/GraphQL/Schema/SdlSchemaPluginBase.php
@@ -13,12 +13,15 @@
 use Drupal\graphql\Plugin\SchemaExtensionPluginInterface;
 use Drupal\graphql\Plugin\SchemaExtensionPluginManager;
 use Drupal\graphql\Plugin\SchemaPluginInterface;
+use GraphQL\Language\AST\DocumentNode;
 use GraphQL\Language\AST\InterfaceTypeDefinitionNode;
 use GraphQL\Language\AST\TypeDefinitionNode;
 use GraphQL\Language\AST\UnionTypeDefinitionNode;
 use GraphQL\Language\Parser;
+use GraphQL\Type\Schema;
 use GraphQL\Utils\BuildSchema;
 use GraphQL\Utils\SchemaExtender;
+use GraphQL\Utils\SchemaPrinter;
 use Symfony\Component\DependencyInjection\ContainerInterface;
 
 /**
@@ -117,15 +120,8 @@ public function __construct(
    */
   public function getSchema(ResolverRegistryInterface $registry) {
     $extensions = $this->getExtensions();
-    $resolver = [$registry, 'resolveType'];
     $document = $this->getSchemaDocument($extensions);
-    $schema = BuildSchema::build($document, function ($config, TypeDefinitionNode $type) use ($resolver) {
-      if ($type instanceof InterfaceTypeDefinitionNode || $type instanceof UnionTypeDefinitionNode) {
-        $config['resolveType'] = $resolver;
-      }
-
-      return $config;
-    });
+    $schema = $this->buildSchema($document, $registry);
 
     if (empty($extensions)) {
       return $schema;
@@ -135,10 +131,31 @@ public function getSchema(ResolverRegistryInterface $registry) {
       $extension->registerResolvers($registry);
     }
 
-    if ($extendSchema = $this->getExtensionDocument($extensions)) {
-      return SchemaExtender::extend($schema, $extendSchema);
+    $extendedDocument = $this->getFullSchemaDocument($schema, $extensions);
+    if (empty($extendedDocument)) {
+      return $schema;
     }
 
+    return $this->buildSchema($extendedDocument, $registry);
+  }
+
+  /**
+   * Create a GraphQL schema object from the given AST document.
+   *
+   * This method is private for now as the build/cache approach might change.
+   */
+  private function buildSchema(DocumentNode $astDocument, ResolverRegistryInterface $registry): Schema {
+    $resolver = [$registry, 'resolveType'];
+    // Performance: only validate the schema in development mode, skip it in
+    // production on every request.
+    $options = empty($this->inDevelopment) ? ['assumeValid' => TRUE] : [];
+    $schema = BuildSchema::build($astDocument, function ($config, TypeDefinitionNode $type) use ($resolver) {
+      if ($type instanceof InterfaceTypeDefinitionNode || $type instanceof UnionTypeDefinitionNode) {
+        $config['resolveType'] = $resolver;
+      }
+
+      return $config;
+    }, $options);
     return $schema;
   }
 
@@ -185,6 +202,33 @@ protected function getSchemaDocument(array $extensions = []) {
     return $ast;
   }
 
+  /**
+   * Returns the full AST combination of parsed schema with extensions, cached.
+   *
+   * This method is private for now as the build/cache approach might change.
+   */
+  private function getFullSchemaDocument(Schema $schema, array $extensions): ?DocumentNode {
+    // Only use caching of the parsed document if we aren't in development mode.
+    $cid = "full:{$this->getPluginId()}";
+    if (empty($this->inDevelopment) && $cache = $this->astCache->get($cid)) {
+      return $cache->data;
+    }
+
+    $ast = NULL;
+    if ($extendAst = $this->getExtensionDocument($extensions)) {
+      $fullSchema = SchemaExtender::extend($schema, $extendAst);
+      // Performance: export the full schema as string and parse it again. That
+      // way we can cache the full AST.
+      $fullSchemaString = SchemaPrinter::doPrint($fullSchema);
+      $ast = Parser::parse($fullSchemaString, ['noLocation' => TRUE]);
+    }
+
+    if (empty($this->inDevelopment)) {
+      $this->astCache->set($cid, $ast, CacheBackendInterface::CACHE_PERMANENT, ['graphql']);
+    }
+    return $ast;
+  }
+
   /**
    * Retrieves the parsed AST of the schema extension definitions.
    *
@@ -196,12 +240,6 @@ protected function getSchemaDocument(array $extensions = []) {
    * @throws \GraphQL\Error\SyntaxError
    */
   protected function getExtensionDocument(array $extensions = []) {
-    // Only use caching of the parsed document if we aren't in development mode.
-    $cid = "extension:{$this->getPluginId()}";
-    if (empty($this->inDevelopment) && $cache = $this->astCache->get($cid)) {
-      return $cache->data;
-    }
-
     $extensions = array_filter(array_map(function (SchemaExtensionPluginInterface $extension) {
       return $extension->getExtensionDefinition();
     }, $extensions), function ($definition) {
@@ -209,10 +247,7 @@ protected function getExtensionDocument(array $extensions = []) {
     });
 
     $ast = !empty($extensions) ? Parser::parse(implode("\n\n", $extensions), ['noLocation' => TRUE]) : NULL;
-    if (empty($this->inDevelopment)) {
-      $this->astCache->set($cid, $ast, CacheBackendInterface::CACHE_PERMANENT, ['graphql']);
-    }
-
+    // No AST caching here as that will be done in getFullSchemaDocument().
     return $ast;
   }