fix(query-planning): Prevent error for satisfiable @shareable mutation fields

.changeset/many-rings-glow.md

@@ -0,0 +1,6 @@
+---
+"@apollo/query-planner": patch
+"@apollo/query-graphs": patch
+---
+
+Fixes a bug where query planning may unexpectedly error due to attempting to generate a plan where a `@shareable` mutation field is called more than once across multiple subgraphs. ([#3304](https://github.com/apollographql/federation/pull/3304))

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

@@ -1490,8 +1490,11 @@ function advancePathWithNonCollectingAndTypePreservingTransitions<TTrigger, V ex
       // We have edges between Query objects so that if a field returns a query object, we can jump to any subgraph
       // at that point. However, there is no point of using those edges at the beginning of a path, except for when
       // we have a @defer, in which case we want to allow re-jumping to the same subgraph.
-      if (isTopLevelPath && edge.transition.kind === 'RootTypeResolution' && !(toAdvance.deferOnTail && edge.isKeyOrRootTypeEdgeToSelf())) {
-        debug.groupEnd(`Ignored: edge is a top-level "RootTypeResolution"`);
+      if (isTopLevelPath
+        && (edge.transition.kind === 'RootTypeResolution' || edge.transition.kind === 'KeyResolution')
+        && !(toAdvance.deferOnTail && edge.isKeyOrRootTypeEdgeToSelf())
+      ) {
+        debug.groupEnd(`Ignored: edge is a top-level "RootTypeResolution" or "KeyResolution"`);
         continue;
       }
 
@@ -1573,7 +1576,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 +2390,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 +2401,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-graphs-js/src/nonLocalSelectionsEstimation.ts

@@ -600,13 +600,18 @@ export class NonLocalSelectionsMetadata {
    * This calls {@link checkNonLocalSelectionsLimitExceeded} for each of the
    * selections in the open branches stack; see that function's doc comment for
    * more information.
+   *
+   * To support mutations, we allow indicating the initial subgraph is
+   * constrained, in which case indirect options will be ignored until the first
+   * field (similar to query planning).
    */
   checkNonLocalSelectionsLimitExceededAtRoot(
     stack: [Selection, SimultaneousPathsWithLazyIndirectPaths[]][],
     state: NonLocalSelectionsState,
     supergraphSchema: Schema,
     inconsistentAbstractTypesRuntimes: Set<string>,
     overrideConditions: Map<string, boolean>,
+    isInitialSubgraphConstrained: boolean,
   ): boolean {
     for (const [selection, simultaneousPaths] of stack) {
       const tailVertices = new Set<Vertex>();
@@ -616,7 +621,10 @@ export class NonLocalSelectionsMetadata {
         }
       }
       const tailVerticesInfo =
-        this.estimateVerticesWithIndirectOptions(tailVertices);
+        this.estimateVerticesWithIndirectOptions(
+          tailVertices,
+          isInitialSubgraphConstrained,
+        );
 
       // Note that top-level selections aren't avoided via fully-local selection
       // set optimization, so we always add them here.
@@ -626,12 +634,16 @@ export class NonLocalSelectionsMetadata {
 
       if (selection.selectionSet) {
         const selectionHasDefer = selection.hasDefer();
+        const isInitialSubgraphConstrainedAfterElement =
+          isInitialSubgraphConstrained
+            && selection.kind === 'FragmentSelection';
         const nextVertices = this.estimateNextVerticesForSelection(
           selection.element,
           tailVerticesInfo,
           state,
           supergraphSchema,
           overrideConditions,
+          isInitialSubgraphConstrainedAfterElement,
         );
         if (this.checkNonLocalSelectionsLimitExceeded(
           selection.selectionSet,
@@ -641,6 +653,7 @@ export class NonLocalSelectionsMetadata {
           supergraphSchema,
           inconsistentAbstractTypesRuntimes,
           overrideConditions,
+          isInitialSubgraphConstrainedAfterElement,
         )) {
           return true;
         }
@@ -674,6 +687,10 @@ export class NonLocalSelectionsMetadata {
    * Note that this function takes in whether the parent selection of the
    * selection set has @defer, as that affects whether the optimization is
    * disabled for that selection set.
+   *
+   * To support mutations, we allow indicating the initial subgraph is
+   * constrained, in which case indirect options will be ignored until the first
+   * field (similar to query planning).
    */
   private checkNonLocalSelectionsLimitExceeded(
     selectionSet: SelectionSet,
@@ -683,6 +700,7 @@ export class NonLocalSelectionsMetadata {
     supergraphSchema: Schema,
     inconsistentAbstractTypesRuntimes: Set<string>,
     overrideConditions: Map<string, boolean>,
+    isInitialSubgraphConstrained: boolean,
   ): boolean {
     // Compute whether the selection set is non-local, and if so, add its
     // selections to the count. Any of the following causes the selection set to
@@ -709,12 +727,16 @@ export class NonLocalSelectionsMetadata {
 
       const oldCount = state.count;
       if (selection.selectionSet) {
+        const isInitialSubgraphConstrainedAfterElement =
+          isInitialSubgraphConstrained
+            && selection.kind === 'FragmentSelection';
         const nextVertices = this.estimateNextVerticesForSelection(
           element,
           parentVertices,
           state,
           supergraphSchema,
           overrideConditions,
+          isInitialSubgraphConstrainedAfterElement,
         );
         if (this.checkNonLocalSelectionsLimitExceeded(
           selection.selectionSet,
@@ -724,6 +746,7 @@ export class NonLocalSelectionsMetadata {
           supergraphSchema,
           inconsistentAbstractTypesRuntimes,
           overrideConditions,
+          isInitialSubgraphConstrainedAfterElement,
         )) {
           return true;
         }
@@ -822,13 +845,20 @@ export class NonLocalSelectionsMetadata {
    * selection for a set of parent vertices (including indirect options), this
    * function can be used to estimate an upper bound on the next vertices after
    * taking the selection (also with indirect options).
+   *
+   * To support mutations, we allow indicating the initial subgraph will be
+   * constrained after taking the element, in which case indirect options will
+   * be ignored (and caching will be skipped). This is to ensure that top-level
+   * mutation fields are not executed on a different subgraph than the initial
+   * one during query planning.
    */
   private estimateNextVerticesForSelection(
     element: OperationElement,
     parentVertices: NextVerticesInfo,
     state: NonLocalSelectionsState,
     supergraphSchema: Schema,
     overrideConditions: Map<string, boolean>,
+    isInitialSubgraphConstrainedAfterElement: boolean,
   ): NextVerticesInfo {
     const selectionKey = element.kind === 'Field'
       ? element.definition.name
@@ -837,6 +867,28 @@ export class NonLocalSelectionsMetadata {
       // For empty type condition, the vertices don't change.
       return parentVertices;
     }
+    if (isInitialSubgraphConstrainedAfterElement) {
+      // When the initial subgraph is constrained, skip caching entirely. Note
+      // that caching is not skipped when the initial subgraph is constrained
+      // before this element but not after. Because of that, there may be cache
+      // entries for remaining vertices that were actually part of a complete
+      // digraph, but this is only a slight caching inefficiency and doesn't
+      // affect the computation's result.
+      assert(
+        parentVertices.nextVerticesWithIndirectOptions.types.size === 0,
+        () => 'Initial subgraph was constrained which indicates no indirect'
+          + ' options should be taken, but the parent vertices unexpectedly had'
+          + ' a complete digraph which indicates indirect options were taken'
+          + ' upstream in the path.',
+      );
+      return this.estimateNextVerticesForSelectionWithoutCaching(
+        element,
+        parentVertices.nextVerticesWithIndirectOptions.remainingVertices,
+        supergraphSchema,
+        overrideConditions,
+        true,
+      );
+    }
     let cache = state.nextVerticesCache.get(selectionKey);
     if (!cache) {
       cache = {
@@ -866,6 +918,7 @@ export class NonLocalSelectionsMetadata {
           indirectOptions.sameTypeOptions,
           supergraphSchema,
           overrideConditions,
+          false,
         );
         cache.typesToNextVertices.set(typeName, cacheEntry);
       }
@@ -879,6 +932,7 @@ export class NonLocalSelectionsMetadata {
           [vertex],
           supergraphSchema,
           overrideConditions,
+          false,
         );
         cache.remainingVerticesToNextVertices.set(vertex, cacheEntry);
       }
@@ -922,16 +976,23 @@ export class NonLocalSelectionsMetadata {
    * (We do account for override conditions, which are relatively
    * straightforward.)
    *
-   * Since we're iterating through next vertices in the process, for efficiency
-   * sake we also compute whether there are any reachable cross-subgraph edges
-   * from the next vertices (without indirect options). This method assumes that
-   * inline fragments have type conditions.
+   * Since we're iterating through next vertices in the process, for
+   * efficiency's sake we also compute whether there are any reachable
+   * cross-subgraph edges from the next vertices (without indirect options).
+   * This method assumes that inline fragments have type conditions.
+   *
+   * To support mutations, we allow indicating the initial subgraph will be
+   * constrained after taking the element, in which case indirect options will
+   * be ignored. This is to ensure that top-level mutation fields are not
+   * executed on a different subgraph than the initial one during query
+   * planning.
    */
   private estimateNextVerticesForSelectionWithoutCaching(
     element: OperationElement,
     parentVertices: Iterable<Vertex>,
     supergraphSchema: Schema,
     overrideConditions: Map<string, boolean>,
+    isInitialSubgraphConstrainedAfterElement: boolean,
   ): NextVerticesInfo {
     const nextVertices = new Set<Vertex>();
     switch (element.kind) {
@@ -955,7 +1016,7 @@ export class NonLocalSelectionsMetadata {
           }
         };
         for (const vertex of parentVertices) {
-          // As an upper bound for efficiency sake, we consider both
+          // As an upper bound for efficiency's sake, we consider both
           // non-type-exploded and type-exploded options.
           processHeadVertex(vertex);
           const downcasts = this.verticesToObjectTypeDowncasts.get(vertex);
@@ -1041,16 +1102,33 @@ export class NonLocalSelectionsMetadata {
         assertUnreachable(element);
     }
 
-    return this.estimateVerticesWithIndirectOptions(nextVertices);
+    return this.estimateVerticesWithIndirectOptions(
+      nextVertices,
+      isInitialSubgraphConstrainedAfterElement,
+    );
   }
 
   /**
-   * Estimate the indirect options for the given next vertices, and add them to
-   * the given vertices. As an upper bound for efficiency's sake, we assume we
-   * can take any indirect option (i.e. ignore any edge conditions).
+   * Estimate the indirect options for the given next vertices, and return the
+   * given next vertices along with `nextVerticesWithIndirectOptions` which
+   * contains these direct and indirect options. As an upper bound for
+   * efficiency's sake, we assume we can take any indirect option (i.e. ignore
+   * any edge conditions).
+   *
+   * Since we're iterating through next vertices in the process, for
+   * efficiency's sake we also compute whether there are any reachable
+   * cross-subgraph edges from the next vertices (without indirect options).
+   *
+   * To support mutations, we allow ignoring indirect options, as we don't want
+   * top-level mutation fields to be executed on a different subgraph than the
+   * initial one. In that case, `nextVerticesWithIndirectOptions` will not have
+   * any `types`, and the given vertices will be added to `remainingVertices`
+   * (despite them potentially being part of the complete digraph for their
+   * type). This is fine, as caching logic accounts for this accordingly.
    */
   private estimateVerticesWithIndirectOptions(
     nextVertices: Set<Vertex>,
+    ignoreIndirectOptions: boolean,
   ): NextVerticesInfo {
     const nextVerticesInfo: NextVerticesInfo = {
       nextVertices,
@@ -1063,7 +1141,16 @@ export class NonLocalSelectionsMetadata {
     for (const nextVertex of nextVertices) {
       nextVerticesInfo.nextVerticesHaveReachableCrossSubgraphEdges ||=
         nextVertex.hasReachableCrossSubgraphEdges;
-
+
+      // As noted above, we don't want top-level mutation fields to be executed
+      // on a different subgraph than the initial one, so we support ignoring
+      // indirect options here.
+      if (ignoreIndirectOptions) {
+        nextVerticesInfo.nextVerticesWithIndirectOptions.remainingVertices
+          .add(nextVertex);
+        continue;
+      }
+
       const typeName = nextVertex.type.name
       const optionsMetadata = this.typesToIndirectOptions.get(typeName);
       if (optionsMetadata) {
@@ -1085,7 +1172,7 @@ export class NonLocalSelectionsMetadata {
           continue;
         }
       }
-      // We need to add the remaining vertex, and if its our first time seeing
+      // We need to add the remaining vertex, and if it's our first time seeing
       // it, we also add any of its interface object options.
       if (
         !nextVerticesInfo.nextVerticesWithIndirectOptions.remainingVertices