Fix query planning for @Shareable mutation fields

Author: sachindshinde

query-graphs-js/src/__tests__/graphPath.test.ts

@@ -72,6 +72,7 @@ function createOptions(supergraph: Schema, queryGraph: QueryGraph): Simultaneous
     [],
     [],
     new Map(),
+    null,
   );
 }
 

query-graphs-js/src/graphPath.ts

@@ -1573,7 +1573,11 @@ function advancePathWithNonCollectingAndTypePreservingTransitions<TTrigger, V ex
           let backToPreviousSubgraph: boolean;
           if (subgraphEnteringEdge.edge.transition.kind === 'SubgraphEnteringTransition') {
             assert(toAdvance.root instanceof RootVertex, () => `${toAdvance} should be a root path if it starts with subgraph entering edge ${subgraphEnteringEdge.edge}`);
-            prevSubgraphEnteringVertex = rootVertexForSubgraph(toAdvance.graph, edge.tail.source, toAdvance.root.rootKind);
+            // Since mutation options need to originate from the same subgraph, we pretend we cannot find a root vertex
+            // in another subgraph (effectively skipping the optimization).
+            prevSubgraphEnteringVertex = toAdvance.root.rootKind !== 'mutation'
+              ? rootVertexForSubgraph(toAdvance.graph, edge.tail.source, toAdvance.root.rootKind)
+              : undefined;
             // If the entering edge is the root entering of subgraphs, then the "prev subgraph" is really `edge.tail.source` and
             // so `edge` always get us back to that (but `subgraphEnteringEdge.edge.head.source` would be `FEDERATED_GRAPH_ROOT_SOURCE`,
             // so the test we do in the `else` branch would not work here).
@@ -2383,6 +2387,7 @@ export function createInitialOptions<V extends Vertex>(
   excludedEdges: ExcludedDestinations,
   excludedConditions: ExcludedConditions,
   overrideConditions: Map<string, boolean>,
+  initialSubgraphConstraint: string | null,
 ): SimultaneousPathsWithLazyIndirectPaths<V>[] {
   const lazyInitialPath = new SimultaneousPathsWithLazyIndirectPaths(
     [initialPath],
@@ -2393,7 +2398,12 @@ export function createInitialOptions<V extends Vertex>(
     overrideConditions,
   );
   if (isFederatedGraphRootType(initialPath.tail.type)) {
-    const initialOptions = lazyInitialPath.indirectOptions(initialContext, 0);
+    let initialOptions = lazyInitialPath.indirectOptions(initialContext, 0);
+    if (initialSubgraphConstraint !== null) {
+      initialOptions.paths = initialOptions
+        .paths
+        .filter((path) => path.tail.source === initialSubgraphConstraint);
+    }
     return createLazyOptions(initialOptions.paths.map(p => [p]), lazyInitialPath, initialContext, overrideConditions);
   } else {
     return [lazyInitialPath];

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

@@ -6127,6 +6127,85 @@ describe('mutations', () => {
       }
     `);
   });
+
+  it('executes a single mutation operation on a @shareable field', () => {
+    const subgraph1 = {
+      name: 'Subgraph1',
+      typeDefs: gql`
+        type Query {
+          dummy: Int
+        }
+
+        type Mutation {
+          f: F @shareable
+        }
+
+        type F @key(fields: "id") {
+          id: ID!
+          x: Int
+        }
+      `,
+    };
+
+    const subgraph2 = {
+      name: 'Subgraph2',
+      typeDefs: gql`
+        type Mutation {
+          f: F @shareable
+        }
+
+        type F @key(fields: "id", resolvable: false) {
+          id: ID!
+          y: Int
+        }
+      `,
+    };
+
+    const [api, queryPlanner] = composeAndCreatePlanner(subgraph1, subgraph2);
+    const operation = operationFromDocument(
+      api,
+      gql`
+        mutation {
+          f {
+            x
+            y
+          }
+        }
+      `,
+    );
+
+    const plan = queryPlanner.buildQueryPlan(operation);
+    expect(plan).toMatchInlineSnapshot(`
+      QueryPlan {
+        Sequence {
+          Fetch(service: "Subgraph2") {
+            {
+              f {
+                __typename
+                id
+                y
+              }
+            }
+          },
+          Flatten(path: "f") {
+            Fetch(service: "Subgraph1") {
+              {
+                ... on F {
+                  __typename
+                  id
+                }
+              } =>
+              {
+                ... on F {
+                  x
+                }
+              }
+            },
+          },
+        },
+      }
+    `);
+  });
 });
 
 describe('interface type-explosion', () => {

query-planner-js/src/buildPlan.ts

@@ -398,6 +398,7 @@ class QueryPlanningTraversal<RV extends Vertex> {
     initialContext: PathContext,
     typeConditionedFetching: boolean,
     nonLocalSelectionsState: NonLocalSelectionsState | null,
+    initialSubgraphConstraint: string | null,
     excludedDestinations: ExcludedDestinations = [],
     excludedConditions: ExcludedConditions = [],
   ) {
@@ -418,6 +419,7 @@ class QueryPlanningTraversal<RV extends Vertex> {
       excludedDestinations,
       excludedConditions,
       parameters.overrideConditions,
+      initialSubgraphConstraint,
     );
     this.stack = mapOptionsToSelections(selectionSet, initialOptions);
     if (
@@ -527,8 +529,9 @@ class QueryPlanningTraversal<RV extends Vertex> {
       // If we have no options, it means there is no way to build a plan for that branch, and
       // that means the whole query planning has no plan.
       // This should never happen for a top-level query planning (unless the supergraph has *not* been
-      // validated), but can happen when computing sub-plans for a key condition.
-      if (this.isTopLevel) {
+      // validated), but can happen when computing sub-plans for a key condition and when computing
+      // a top-level plan for a mutation field on a specific subgraph.
+      if (this.isTopLevel && this.rootKind !== 'mutation') {
         debug.groupEnd(() => `No valid options to advance ${selection} from ${advanceOptionsToString(options)}`);
         throw new Error(`Was not able to find any options for ${selection}: This shouldn't have happened.`);
       } else {
@@ -794,6 +797,7 @@ class QueryPlanningTraversal<RV extends Vertex> {
       context,
       this.typeConditionedFetching,
       null,
+      null,
       excludedDestinations,
       addConditionExclusion(excludedConditions, edge.conditions),
     ).findBestPlan();
@@ -3593,7 +3597,7 @@ function computePlanInternal({
 
   const { operation, processor } = parameters;
   if (operation.rootKind === 'mutation') {
-    const dependencyGraphs = computeRootSerialDependencyGraph(
+    const dependencyGraphs = computeRootSerialDependencyGraphForMutation(
       parameters,
       hasDefers,
       nonLocalSelectionsState,
@@ -3770,13 +3774,52 @@ function computeRootParallelBestPlan(
     emptyContext,
     parameters.config.typeConditionedFetching,
     nonLocalSelectionsState,
+    null,
   );
   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),
   // so we just return an empty plan.
   return plan ?? createEmptyPlan(parameters);
 }
 
+function computeRootParallelBestPlanForMutation(
+  parameters: PlanningParameters<RootVertex>,
+  selection: SelectionSet,
+  startFetchIdGen: number,
+  hasDefers: boolean,
+  nonLocalSelectionsState: NonLocalSelectionsState | null,
+): [FetchDependencyGraph, OpPathTree<RootVertex>, number] {
+  let bestPlan:
+    | [FetchDependencyGraph, OpPathTree<RootVertex>, number]
+    | undefined;
+  const mutationSubgraphs = parameters.federatedQueryGraph
+    .outEdges(parameters.root).map((edge) => edge.tail.source);
+  for (const mutationSubgraph of mutationSubgraphs) {
+    const planningTraversal = new QueryPlanningTraversal(
+      parameters,
+      selection,
+      startFetchIdGen,
+      hasDefers,
+      parameters.root.rootKind,
+      defaultCostFunction,
+      emptyContext,
+      parameters.config.typeConditionedFetching,
+      nonLocalSelectionsState,
+      mutationSubgraph,
+    );
+    const plan = planningTraversal.findBestPlan();
+    if (!bestPlan || (plan && plan[2] < bestPlan[2])) {
+      bestPlan = plan;
+    }
+  }
+  if (!bestPlan) {
+    throw new Error(
+      `Was not able to plan ${parameters.operation.toString(false, false)} starting from a single subgraph: This shouldn't have happened.`,
+    );
+  }
+  return bestPlan;
+}
+
 function createEmptyPlan(
   parameters: PlanningParameters<RootVertex>,
 ): [FetchDependencyGraph, OpPathTree<RootVertex>, number] {
@@ -3794,7 +3837,7 @@ function onlyRootSubgraph(graph: FetchDependencyGraph): string {
   return subgraphs[0];
 }
 
-function computeRootSerialDependencyGraph(
+function computeRootSerialDependencyGraphForMutation(
   parameters: PlanningParameters<RootVertex>,
   hasDefers: boolean,
   nonLocalSelectionsState: NonLocalSelectionsState | null,
@@ -3805,7 +3848,7 @@ function computeRootSerialDependencyGraph(
   const splittedRoots = splitTopLevelFields(operation.selectionSet);
   const graphs: FetchDependencyGraph[] = [];
   let startingFetchId = 0;
-  let [prevDepGraph, prevPaths] = computeRootParallelBestPlan(
+  let [prevDepGraph, prevPaths] = computeRootParallelBestPlanForMutation(
     parameters,
     splittedRoots[0],
     startingFetchId,
@@ -3814,7 +3857,7 @@ function computeRootSerialDependencyGraph(
   );
   let prevSubgraph = onlyRootSubgraph(prevDepGraph);
   for (let i = 1; i < splittedRoots.length; i++) {
-    const [newDepGraph, newPaths] = computeRootParallelBestPlan(
+    const [newDepGraph, newPaths] = computeRootParallelBestPlanForMutation(
       parameters,
       splittedRoots[i],
       prevDepGraph.nextFetchId(),