From b3a3cb84d8d67d1d6e817dc85b9ae0ecdd9908d1 Mon Sep 17 00:00:00 2001
From: Sylvain Lebresne <sylvain@apollographql.com>
Date: Tue, 20 Sep 2022 18:27:31 +0200
Subject: [PATCH] Never merge fragments that have `@defer` (#2143)

Never merge fragments that have `@defer`
---
 internals-js/src/definitions.ts               |  28 ++++-
 .../src/__tests__/buildPlan.defer.test.ts     | 108 ++++++++++++++++++
 2 files changed, 132 insertions(+), 4 deletions(-)

diff --git a/internals-js/src/definitions.ts b/internals-js/src/definitions.ts
index dcb8fc8b4..02125e1df 100644
--- a/internals-js/src/definitions.ts
+++ b/internals-js/src/definitions.ts
@@ -3197,20 +3197,40 @@ export class Directive<
   }
 }
 
-export function sameDirectiveApplication(application1: Directive<any, any>, application2: Directive<any, any>): boolean {
-  return application1.name === application2.name && argumentsEquals(application1.arguments(), application2.arguments());
+/**
+ * Checks if 2 directive applications should be considered equal.
+ *
+ * By default, 2 directive applications are considered equal if they are for the same directive and are passed the same values to
+ * the same arguments. However, some special directive can be excluded so that no 2 applications are ever consider equal. By default,
+ * this is the case of @defer, as never want to merge @defer applications so that each create its own "deferred block".
+ */
+export function sameDirectiveApplication(
+  application1: Directive<any, any>,
+  application2: Directive<any, any>,
+  directivesNeverEqualToThemselves: string[] = [ 'defer' ],
+): boolean {
+  // Note: we check name equality first because this method is most often called with directive that are simply not the same
+  // name and this ensure we exit cheaply more often than not.
+  return application1.name === application2.name 
+    && !directivesNeverEqualToThemselves.includes(application1.name)
+    && !directivesNeverEqualToThemselves.includes(application2.name)
+    && argumentsEquals(application1.arguments(), application2.arguments());
 }
 
 /**
  * Checks whether the 2 provided "set" of directive applications are the same (same applications, regardless or order).
  */
-export function sameDirectiveApplications(applications1: readonly Directive<any, any>[], applications2: readonly Directive<any, any>[]): boolean {
+export function sameDirectiveApplications(
+  applications1: readonly Directive<any, any>[],
+  applications2: readonly Directive<any, any>[],
+  directivesNeverEqualToThemselves: string[] = [ 'defer' ],
+): boolean {
   if (applications1.length !== applications2.length) {
     return false;
   }
 
   for (const directive1 of applications1) {
-    if (!applications2.some(directive2 => sameDirectiveApplication(directive1, directive2))) {
+    if (!applications2.some(directive2 => sameDirectiveApplication(directive1, directive2, directivesNeverEqualToThemselves))) {
       return false;
     }
   }
diff --git a/query-planner-js/src/__tests__/buildPlan.defer.test.ts b/query-planner-js/src/__tests__/buildPlan.defer.test.ts
index 0721b0417..c4fd42520 100644
--- a/query-planner-js/src/__tests__/buildPlan.defer.test.ts
+++ b/query-planner-js/src/__tests__/buildPlan.defer.test.ts
@@ -3386,3 +3386,111 @@ describe('named fragments', () => {
     `);
   });
 });
+
+test('do not merge query branches with @defer', () => {
+  const subgraph1 = {
+    name: 'Subgraph1',
+    typeDefs: gql`
+      type Query {
+        t: T
+      }
+
+      type T @key(fields: "id") {
+        id: ID!
+        a: Int
+        b: Int
+      }
+    `
+  }
+
+  const subgraph2 = {
+    name: 'Subgraph2',
+    typeDefs: gql`
+      type T @key(fields: "id") {
+        id: ID!
+        c: Int
+      }
+    `
+  }
+
+  const [api, queryPlanner] = composeAndCreatePlannerWithDefer(subgraph1, subgraph2);
+  // We have 2 separate @defer, so we should 2 deferred parts, not 1 defer parts with parallel fetches.
+  const operation = operationFromDocument(api, gql`
+    {
+      t {
+        a
+        ... @defer {
+          b
+        }
+        ... @defer {
+          c
+        }
+      }
+    }
+  `);
+
+  const queryPlan = queryPlanner.buildQueryPlan(operation);
+  expect(queryPlan).toMatchInlineSnapshot(`
+    QueryPlan {
+      Defer {
+        Primary {
+          {
+            t {
+              a
+            }
+          }:
+          Fetch(service: "Subgraph1", id: 0) {
+            {
+              t {
+                __typename
+                a
+                id
+              }
+            }
+          }
+        }, [
+          Deferred(depends: [0], path: "t") {
+            {
+              c
+            }:
+            Flatten(path: "t") {
+              Fetch(service: "Subgraph2") {
+                {
+                  ... on T {
+                    __typename
+                    id
+                  }
+                } =>
+                {
+                  ... on T {
+                    c
+                  }
+                }
+              },
+            }
+          },
+          Deferred(depends: [0], path: "t") {
+            {
+              b
+            }:
+            Flatten(path: "t") {
+              Fetch(service: "Subgraph1") {
+                {
+                  ... on T {
+                    __typename
+                    id
+                  }
+                } =>
+                {
+                  ... on T {
+                    b
+                  }
+                }
+              },
+            }
+          },
+        ]
+      },
+    }
+  `);
+});