Add test for QP optimization of __typename

The optimization of apollographql#2137 is meant to have no impact on the
generated plans: it's an optimization that ensure the query
planner does not needlessly evaluate ineffient options just
because __typename happens to be resolvable from any subgraph.
But this make testing that optimization a bit more complicated.
To work around this, this patch adds a new behaviour to the
query planner whereby along the generation of the plan, it
also exposes some statistics on the plan generation. As of
this commit, the only thing exposed is the number of plan
that were evaluated under the hood, as that is what we care
to check here (but it is also one of the main contributor
to time spent query planning, so arguably an important
statistic).

There was probably more hacky ways to test this optimization,
but it while this statistic is currently only use for testing,
this feels like something that could have other uses and be
useful to enrich later.

Author: Sylvain Lebresne

query-planner-js/src/__tests__/buildPlan.test.ts

@@ -4055,4 +4055,157 @@ describe('__typename handling', () => {
       }
     `);
   });
+
+  it('does not needlessly consider options for __typename', () => {
+    const subgraph1 = {
+      name: 'Subgraph1',
+      typeDefs: gql`
+        type Query {
+          s: S
+        }
+
+        type S @key(fields: "id") {
+          id: ID
+        }
+      `
+    }
+
+    const subgraph2 = {
+      name: 'Subgraph2',
+      typeDefs: gql`
+        type S @key(fields: "id") {
+          id: ID
+          t: T @shareable
+        }
+
+        type T @key(fields: "id") {
+          id: ID!
+          x: Int
+        }
+      `
+    }
+
+    const subgraph3 = {
+      name: 'Subgraph3',
+      typeDefs: gql`
+        type S @key(fields: "id") {
+          id: ID
+          t: T @shareable
+        }
+
+        type T @key(fields: "id") {
+          id: ID!
+          y: Int
+        }
+      `
+    }
+
+    const [api, queryPlanner] = composeAndCreatePlanner(subgraph1, subgraph2, subgraph3);
+    // This tests the patch from https://github.com/apollographql/federation/pull/2137.
+    // Namely, the schema is such that `x` can only be fetched from one subgraph, but
+    // technically __typename can be fetched from 2 subgraphs. However, the optimization
+    // we test for is that we actually don't consider both choices for __typename and
+    // instead only evaluate a single query plan (the assertion on `evaluatePlanCount`)
+    let operation = operationFromDocument(api, gql`
+      query {
+        s {
+          t {
+            __typename
+            x
+          }
+        }
+      }
+    `);
+
+    let plan = queryPlanner.buildQueryPlan(operation);
+    expect(queryPlanner.lastGeneratedPlanStatistics()?.evaluatedPlanCount).toBe(1);
+    expect(plan).toMatchInlineSnapshot(`
+      QueryPlan {
+        Sequence {
+          Fetch(service: "Subgraph1") {
+            {
+              s {
+                __typename
+                id
+              }
+            }
+          },
+          Flatten(path: "s") {
+            Fetch(service: "Subgraph2") {
+              {
+                ... on S {
+                  __typename
+                  id
+                }
+              } =>
+              {
+                ... on S {
+                  t {
+                    __typename
+                    x
+                  }
+                }
+              }
+            },
+          },
+        },
+      }
+    `);
+
+    // Almost the same test, but we artificially create a case where the result set
+    // for `s` has a __typename alongside just an inline fragments. This should
+    // change nothing to the example (the __typename on `s` is trivially fetched
+    // from the 1st subgraph and does not create new choices), but an early bug
+    // in the implementation made this example forgo the optimization of the
+    // __typename within `t`. We make sure this is not case (that we still only
+    // consider a single choice of plan).
+    operation = operationFromDocument(api, gql`
+      query {
+        s {
+          __typename
+          ... on S {
+            t {
+              __typename
+              x
+            }
+          }
+        }
+      }
+    `);
+
+    plan = queryPlanner.buildQueryPlan(operation);
+    expect(queryPlanner.lastGeneratedPlanStatistics()?.evaluatedPlanCount).toBe(1);
+    expect(plan).toMatchInlineSnapshot(`
+      QueryPlan {
+        Sequence {
+          Fetch(service: "Subgraph1") {
+            {
+              s {
+                __typename
+                id
+              }
+            }
+          },
+          Flatten(path: "s") {
+            Fetch(service: "Subgraph2") {
+              {
+                ... on S {
+                  __typename
+                  id
+                }
+              } =>
+              {
+                ... on S {
+                  t {
+                    __typename
+                    x
+                  }
+                }
+              }
+            },
+          },
+        },
+      }
+    `);
+  });
 });

query-planner-js/src/buildPlan.ts

@@ -230,6 +230,7 @@ class QueryPlanningTaversal<RV extends Vertex> {
     readonly startFetchIdGen: number,
     readonly hasDefers: boolean,
     readonly variableDefinitions: VariableDefinitions,
+    private readonly statistics: PlanningStatistics | undefined,
     private readonly startVertex: RV,
     private readonly rootKind: SchemaRootKind,
     readonly costFunction: CostFunction,
@@ -445,6 +446,10 @@ class QueryPlanningTaversal<RV extends Vertex> {
       debug.log(() => `Reduced plans to consider to ${planCount} plans`);
     }
 
+    if (this.statistics) {
+      this.statistics.evaluatedPlanCount += planCount;
+    }
+
     debug.log(() => `All branches:${this.closedBranches.map((opts, i) => `\n${i}:${opts.map((opt => `\n - ${simultaneousPathsToString(opt)}`))}`)}`);
 
     // Note that usually, we'll have a majority of branches with just one option. We can group them in
@@ -528,6 +533,7 @@ class QueryPlanningTaversal<RV extends Vertex> {
       0,
       false,
       this.variableDefinitions,
+      undefined,
       edge.head,
       'query',
       this.costFunction,
@@ -2012,6 +2018,10 @@ function withSiblingTypenameOptimizedAway(operation: Operation): Operation {
   );
 }
 
+export type PlanningStatistics = {
+  evaluatedPlanCount: number,
+}
+
 export function computeQueryPlan({
   config,
   supergraphSchema,
@@ -2022,14 +2032,21 @@ export function computeQueryPlan({
   supergraphSchema: Schema,
   federatedQueryGraph: QueryGraph,
   operation: Operation,
-}): QueryPlan {
+}): {
+  plan: QueryPlan,
+  statistics: PlanningStatistics,
+} {
   if (operation.rootKind === 'subscription') {
     throw ERRORS.UNSUPPORTED_FEATURE.err(
       'Query planning does not currently support subscriptions.',
       { nodes: [parse(operation.toString())] },
     );
   }
 
+  const statistics: PlanningStatistics = {
+    evaluatedPlanCount: 0,
+  };
+
   const reuseQueryFragments = config.reuseQueryFragments ?? true;
   let fragments = operation.selectionSet.fragments
   if (fragments && reuseQueryFragments) {
@@ -2062,7 +2079,10 @@ export function computeQueryPlan({
   debug.group(() => `Computing plan for\n${operation}`);
   if (operation.selectionSet.isEmpty()) {
     debug.groupEnd('Empty plan');
-    return { kind: 'QueryPlan' };
+    return {
+      plan: { kind: 'QueryPlan' },
+      statistics,
+    };
   }
 
   const root = federatedQueryGraph.root(operation.rootKind);
@@ -2086,6 +2106,7 @@ export function computeQueryPlan({
       processor,
       root,
       deferConditions,
+      statistics,
     })
   } else {
     rootNode = computePlanInternal({
@@ -2095,11 +2116,15 @@ export function computeQueryPlan({
       processor,
       root,
       hasDefers,
+      statistics,
     });
   }
 
   debug.groupEnd('Query plan computed');
-  return { kind: 'QueryPlan', node: rootNode };
+  return {
+    plan: { kind: 'QueryPlan', node: rootNode },
+    statistics,
+  };
 }
 
 function computePlanInternal({
@@ -2109,16 +2134,18 @@ function computePlanInternal({
   processor,
   root,
   hasDefers,
+  statistics,
 }: {
   supergraphSchema: Schema,
   federatedQueryGraph: QueryGraph,
   operation: Operation,
   processor: FetchGroupProcessor<PlanNode | undefined, DeferredNode>
   root: RootVertex,
   hasDefers: boolean,
+  statistics: PlanningStatistics,
 }): PlanNode | undefined {
   if (operation.rootKind === 'mutation') {
-    const dependencyGraphs = computeRootSerialDependencyGraph(supergraphSchema, operation, federatedQueryGraph, root, hasDefers);
+    const dependencyGraphs = computeRootSerialDependencyGraph(supergraphSchema, operation, federatedQueryGraph, root, hasDefers, statistics);
     let allMain: (PlanNode | undefined)[] = [];
     let allDeferred: DeferredNode[] = [];
     let primarySelection: SelectionSet | undefined = undefined;
@@ -2143,7 +2170,7 @@ function computePlanInternal({
       deferred: allDeferred,
     });
   } else {
-    const dependencyGraph =  computeRootParallelDependencyGraph(supergraphSchema, operation, federatedQueryGraph, root, 0, hasDefers);
+    const dependencyGraph =  computeRootParallelDependencyGraph(supergraphSchema, operation, federatedQueryGraph, root, 0, hasDefers, statistics);
     const { main, deferred } = dependencyGraph.process(processor);
     return processRootNodes({
       processor,
@@ -2162,13 +2189,15 @@ function computePlanForDeferConditionals({
   processor,
   root,
   deferConditions,
+  statistics,
 }: {
   supergraphSchema: Schema,
   federatedQueryGraph: QueryGraph,
   operation: Operation,
   processor: FetchGroupProcessor<PlanNode | undefined, DeferredNode>
   root: RootVertex,
   deferConditions: SetMultiMap<string, string>,
+  statistics: PlanningStatistics,
 }): PlanNode | undefined {
   return generateConditionNodes(
     operation,
@@ -2181,6 +2210,7 @@ function computePlanForDeferConditionals({
       processor,
       root,
       hasDefers: true,
+      statistics,
     }),
   );
 }
@@ -2296,6 +2326,7 @@ function computeRootParallelDependencyGraph(
   root: RootVertex,
   startFetchIdGen: number,
   hasDefer: boolean,
+  statistics: PlanningStatistics,
 ): FetchDependencyGraph {
   return computeRootParallelBestPlan(
     supergraphSchema,
@@ -2305,6 +2336,7 @@ function computeRootParallelDependencyGraph(
     root,
     startFetchIdGen,
     hasDefer,
+    statistics,
   )[0];
 }
 
@@ -2316,6 +2348,7 @@ function computeRootParallelBestPlan(
   root: RootVertex,
   startFetchIdGen: number,
   hasDefers: boolean,
+  statistics: PlanningStatistics,
 ): [FetchDependencyGraph, OpPathTree<RootVertex>, number] {
   const planningTraversal = new QueryPlanningTaversal(
     supergraphSchema,
@@ -2324,10 +2357,11 @@ function computeRootParallelBestPlan(
     startFetchIdGen,
     hasDefers,
     variables,
+    statistics,
     root,
     root.rootKind,
     defaultCostFunction,
-    emptyContext
+    emptyContext,
   );
   const plan = planningTraversal.findBestPlan();
   // Getting no plan means the query is essentially unsatisfiable (it's a valid query, but we can prove it will never return a result),
@@ -2358,16 +2392,17 @@ function computeRootSerialDependencyGraph(
   federatedQueryGraph: QueryGraph,
   root: RootVertex,
   hasDefers: boolean,
+  statistics: PlanningStatistics,
 ): FetchDependencyGraph[] {
   const rootType = hasDefers ? supergraphSchema.schemaDefinition.rootType(root.rootKind) : undefined;
   // We have to serially compute a plan for each top-level selection.
   const splittedRoots = splitTopLevelFields(operation.selectionSet);
   const graphs: FetchDependencyGraph[] = [];
   let startingFetchId: number = 0;
-  let [prevDepGraph, prevPaths] = computeRootParallelBestPlan(supergraphSchema, splittedRoots[0], operation.variableDefinitions, federatedQueryGraph, root, startingFetchId, hasDefers);
+  let [prevDepGraph, prevPaths] = computeRootParallelBestPlan(supergraphSchema, splittedRoots[0], operation.variableDefinitions, federatedQueryGraph, root, startingFetchId, hasDefers, statistics);
   let prevSubgraph = onlyRootSubgraph(prevDepGraph);
   for (let i = 1; i < splittedRoots.length; i++) {
-    const [newDepGraph, newPaths] = computeRootParallelBestPlan(supergraphSchema, splittedRoots[i], operation.variableDefinitions, federatedQueryGraph, root, prevDepGraph.nextFetchId(), hasDefers);
+    const [newDepGraph, newPaths] = computeRootParallelBestPlan(supergraphSchema, splittedRoots[i], operation.variableDefinitions, federatedQueryGraph, root, prevDepGraph.nextFetchId(), hasDefers, statistics);
     const newSubgraph = onlyRootSubgraph(newDepGraph);
     if (prevSubgraph === newSubgraph) {
       // The new operation (think 'mutation' operation) is on the same subgraph than the previous one, so we can concat them in a single fetch