Fix bug with prefetching on different nesting levels

Allow Promise to be returned from prefetchMethod with returnRequested: true

Author: Andrii Sudiev

src/Annotations/Prefetch.php

@@ -10,8 +10,14 @@
 #[Attribute(Attribute::TARGET_PARAMETER)]
 class Prefetch implements ParameterAnnotationInterface
 {
-    /** @param string|(callable&array{class-string, string}) $callable */
-    public function __construct(public readonly string|array $callable)
+    /**
+     * @param string|(callable&array{class-string, string}) $callable
+     * @param bool $returnRequested return value mapped to requested method
+     */
+    public function __construct(
+        public readonly string|array $callable,
+        public readonly bool $returnRequested = false,
+    )
     {
     }
 

src/Mappers/Parameters/PrefetchParameterMiddleware.php

@@ -52,6 +52,7 @@ public function mapParameter(ReflectionParameter $parameter, DocBlock $docBlock,
             fieldName: $method->getName(),
             resolver: $resolver,
             parameters: $parameters,
+            returnRequested: $prefetch->returnRequested,
         );
     }
 }

src/Parameters/PrefetchDataParameter.php

@@ -12,6 +12,7 @@
 use TheCodingMachine\GraphQLite\PrefetchBuffer;
 use TheCodingMachine\GraphQLite\QueryField;
 
+use function array_key_exists;
 use function assert;
 
 /**
@@ -27,6 +28,7 @@ public function __construct(
         private readonly string $fieldName,
         private readonly mixed $resolver,
         public readonly array $parameters,
+        public readonly bool $returnRequested = false,
     )
     {
     }
@@ -52,24 +54,40 @@ public function resolve(object|null $source, array $args, mixed $context, Resolv
         // So we record all of these ->resolve() calls, collect them together and when a value is actually
         // needed, GraphQL calls the callback of Deferred below. That's when we call the prefetch method,
         // already knowing all the requested fields (source-arguments combinations).
-        return new Deferred(function () use ($info, $context, $args, $prefetchBuffer) {
-            if (! $prefetchBuffer->hasResult($args, $info)) {
-                $prefetchResult = $this->computePrefetch($args, $context, $info, $prefetchBuffer);
-
-                $prefetchBuffer->storeResult($prefetchResult, $args, $info);
+        return new Deferred(function () use ($source, $info, $context, $args, $prefetchBuffer) {
+            if (! $prefetchBuffer->hasResult($source)) {
+                $this->computePrefetch($args, $context, $info, $prefetchBuffer);
             }
 
-            return $prefetchResult ?? $prefetchBuffer->getResult($args, $info);
+            return $prefetchBuffer->getResult($source);
         });
     }
 
     /** @param array<string, mixed> $args */
-    private function computePrefetch(array $args, mixed $context, ResolveInfo $info, PrefetchBuffer $prefetchBuffer): mixed
+    private function computePrefetch(array $args, mixed $context, ResolveInfo $info, PrefetchBuffer $prefetchBuffer): void
     {
         $sources = $prefetchBuffer->getObjectsByArguments($args, $info);
+        $prefetchBuffer->purge($args, $info);
         $toPassPrefetchArgs = QueryField::paramsToArguments($this->fieldName, $this->parameters, null, $args, $context, $info, $this->resolver);
 
-        return ($this->resolver)($sources, ...$toPassPrefetchArgs);
+        $resolvedValues = ($this->resolver)($sources, ...$toPassPrefetchArgs);
+        if ($this->returnRequested) {
+            foreach ($resolvedValues as $key => $resolvedValue) {
+                if (! array_key_exists($key, $sources)) {
+                    throw new GraphQLRuntimeException(
+                        'Called by Prefetch function should accept ' .
+                        'Array<key> and return Array<value>, but the function did ' .
+                        'not return an Array of the same length as the Array of keys.',
+                    );
+                }
+                $prefetchBuffer->storeResult($sources[$key], $resolvedValue);
+            }
+        } else {
+            foreach ($sources as $source) {
+                // map results to each source to support old prefetch behavior
+                $prefetchBuffer->storeResult($source, $resolvedValues);
+            }
+        }
     }
 
     /** @inheritDoc */

src/PrefetchBuffer.php

@@ -5,8 +5,8 @@
 namespace TheCodingMachine\GraphQLite;
 
 use GraphQL\Type\Definition\ResolveInfo;
+use SplObjectStorage;
 
-use function array_key_exists;
 use function md5;
 use function serialize;
 
@@ -18,8 +18,13 @@ class PrefetchBuffer
     /** @var array<string, array<int, object>> An array of buffered, indexed by hash of arguments. */
     private array $objects = [];
 
-    /** @var array<string, mixed> An array of prefetch method results, indexed by hash of arguments. */
-    private array $results = [];
+    /** @var SplObjectStorage A Storage of prefetch method results, holds source to resolved values. */
+    private SplObjectStorage $results;
+
+    public function __construct()
+    {
+        $this->results = new SplObjectStorage();
+    }
 
     /** @param array<array-key, mixed> $arguments The input arguments passed from GraphQL to the field. */
     public function register(
@@ -66,28 +71,22 @@ public function purge(
         unset($this->objects[$this->computeHash($arguments, $info)]);
     }
 
-    /** @param array<array-key, mixed> $arguments The input arguments passed from GraphQL to the field. */
     public function storeResult(
+        object $source,
         mixed $result,
-        array $arguments,
-        ResolveInfo|null $info = null,
     ): void {
-        $this->results[$this->computeHash($arguments, $info)] = $result;
+        $this->results->offsetSet($source, $result);
     }
 
-    /** @param array<array-key, mixed> $arguments The input arguments passed from GraphQL to the field. */
     public function hasResult(
-        array $arguments,
-        ResolveInfo|null $info = null,
+        object $source,
     ): bool {
-        return array_key_exists($this->computeHash($arguments, $info), $this->results);
+        return $this->results->offsetExists($source);
     }
 
-    /** @param array<array-key, mixed> $arguments The input arguments passed from GraphQL to the field. */
     public function getResult(
-        array $arguments,
-        ResolveInfo|null $info = null,
+        object $source,
     ): mixed {
-        return $this->results[$this->computeHash($arguments, $info)];
+        return $this->results->offsetGet($source);
     }
 }

src/QueryField.php

@@ -4,11 +4,9 @@
 
 namespace TheCodingMachine\GraphQLite;
 
-use GraphQL\Deferred;
+use Closure;
 use GraphQL\Error\ClientAware;
 use GraphQL\Executor\Promise\Adapter\SyncPromise;
-use GraphQL\Executor\Promise\Adapter\SyncPromiseAdapter;
-use GraphQL\Executor\Promise\Promise;
 use GraphQL\Language\AST\FieldDefinitionNode;
 use GraphQL\Type\Definition\FieldDefinition;
 use GraphQL\Type\Definition\ListOfType;
@@ -22,9 +20,6 @@
 use TheCodingMachine\GraphQLite\Parameters\ParameterInterface;
 use TheCodingMachine\GraphQLite\Parameters\SourceParameter;
 
-use function array_filter;
-use function array_map;
-
 /**
  * A GraphQL field that maps to a PHP method automatically.
  *
@@ -90,24 +85,12 @@ public function __construct(
                 return $result;
             };
 
-            $deferred = (bool) array_filter($toPassArgs, static fn (mixed $value) => $value instanceof SyncPromise);
-
             // GraphQL allows deferring resolving the field's value using promises, i.e. they call the resolve
             // function ahead of time for all of the fields (allowing us to gather all calls and do something
             // in batch, like prefetch) and then resolve the promises as needed. To support that for prefetch,
             // we're checking if any of the resolved parameters returned a promise. If they did, we know
             // that the value should also be resolved using a promise, so we're wrapping it in one.
-            return $deferred ? new Deferred(static function () use ($toPassArgs, $callResolver) {
-                $syncPromiseAdapter = new SyncPromiseAdapter();
-
-                // Wait for every deferred parameter.
-                $toPassArgs = array_map(
-                    static fn (mixed $value) => $value instanceof SyncPromise ? $syncPromiseAdapter->wait(new Promise($value, $syncPromiseAdapter)) : $value,
-                    $toPassArgs,
-                );
-
-                return $callResolver(...$toPassArgs);
-            }) : $callResolver(...$toPassArgs);
+            return $this->deferred($toPassArgs, $callResolver);
         };
 
         $config += $additionalConfig;
@@ -204,4 +187,21 @@ public static function paramsToArguments(string $name, array $parameters, object
 
         return $toPassArgs;
     }
+
+    /**
+     * @param array<int, mixed> $toPassArgs
+     * Create Deferred if any of arguments contain promise
+     */
+    public static function deferred(array $toPassArgs, Closure $callResolver): mixed
+    {
+        foreach ($toPassArgs as $position => $toPassArg) {
+            if ($toPassArg instanceof SyncPromise) {
+                return $toPassArg->then(static function ($resolvedValue) use ($toPassArgs, $position, $callResolver) {
+                    $toPassArgs[$position] = $resolvedValue;
+                    return self::deferred($toPassArgs, $callResolver);
+                });
+            }
+        }
+        return $callResolver(...$toPassArgs);
+    }
 }

tests/Fixtures/Integration/Controllers/BlogController.php

@@ -0,0 +1,21 @@
+<?php
+
+declare(strict_types=1);
+
+namespace TheCodingMachine\GraphQLite\Fixtures\Integration\Controllers;
+
+use TheCodingMachine\GraphQLite\Annotations\Query;
+use TheCodingMachine\GraphQLite\Fixtures\Integration\Models\Blog;
+
+class BlogController
+{
+    /** @return Blog[] */
+    #[Query]
+    public function getBlogs(): array
+    {
+        return [
+            new Blog(1),
+            new Blog(2),
+        ];
+    }
+}

tests/Fixtures/Integration/Models/Blog.php

@@ -0,0 +1,86 @@
+<?php
+
+declare(strict_types=1);
+
+namespace TheCodingMachine\GraphQLite\Fixtures\Integration\Models;
+
+use TheCodingMachine\GraphQLite\Annotations\Field;
+use TheCodingMachine\GraphQLite\Annotations\Prefetch;
+use TheCodingMachine\GraphQLite\Annotations\Type;
+
+#[Type]
+class Blog
+{
+    public function __construct(
+        private readonly int $id,
+    ) {
+    }
+
+    #[Field(outputType: 'ID!')]
+    public function getId(): int
+    {
+        return $this->id;
+    }
+
+    /**
+     * @param Post[] $prefetchedPosts
+     *
+     * @return Post[]
+     */
+    #[Field]
+    public function getPosts(
+        #[Prefetch('prefetchPosts', true)]
+        array $prefetchedPosts,
+    ): array {
+        return $prefetchedPosts;
+    }
+
+    /**
+     * @param Blog[] $prefetchedSubBlogs
+     *
+     * @return Blog[]
+     */
+    #[Field]
+    public function getSubBlogs(
+        #[Prefetch('prefetchSubBlogs', true)]
+        array $prefetchedSubBlogs,
+    ): array {
+        return $prefetchedSubBlogs;
+    }
+
+    /**
+     * @param Blog[] $blogs
+     *
+     * @return Post[][]
+     */
+    public static function prefetchPosts(iterable $blogs): array
+    {
+        $posts = [];
+        foreach ($blogs as $key => $blog) {
+            $blogId = $blog->getId();
+            $posts[$key] = [
+                new Post('post-' . $blogId . '.1'),
+                new Post('post-' . $blogId . '.2'),
+            ];
+        }
+
+        return $posts;
+    }
+
+    /**
+     * @param Blog[] $blogs
+     *
+     * @return Blog[][]
+     */
+    public static function prefetchSubBlogs(iterable $blogs): array
+    {
+        $subBlogs = [];
+        foreach ($blogs as $key => $blog) {
+            $blogId = $blog->getId();
+            $subBlogId = $blogId * 10;
+            $subBlogs[$key] = [new Blog($subBlogId)];
+        }
+
+        return $subBlogs;
+    }
+}

tests/Fixtures/Integration/Models/Comment.php

@@ -0,0 +1,23 @@
+<?php
+
+declare(strict_types=1);
+
+namespace TheCodingMachine\GraphQLite\Fixtures\Integration\Models;
+
+use TheCodingMachine\GraphQLite\Annotations\Field;
+use TheCodingMachine\GraphQLite\Annotations\Type;
+
+#[Type]
+class Comment
+{
+    public function __construct(
+        private readonly string $text,
+    ) {
+    }
+
+    #[Field]
+    public function getText(): string
+    {
+        return $this->text;
+    }
+}