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

This PR fixes a bug where query planning may unexpectedly error while planning an operation with a satisfiable `@shareable` mutation field. Specifically, this PR: 1. Updates query planning on mutation fields to only consider one initial subgraph at a time for a traversal, and returns the plan with lowest cost. - Previously, this was mixing up options that started in different subgraphs in the same traversal, which would sometimes generate query planning errors if those differing options were chosen. There were a few ways to fix this, but the least invasive/simplest was just to do multiple traversals, each with a limited initial subgraph. - Non-local selection estimation was also appropriately updated to account for this (if it weren't, it would systemically overestimate the number of non-local selections). 2. Updates a particular query planning optimization to avoid discarding an option if another option with less subgraph jumps/cost is found, if that option starts in a different subgraph and the option is for a mutation operation. 3. Updates query planning to ignore `@key` edges at top-level (previously only root type resolution edges were ignored). This PR was originally filed as part of #3303, but was split off so it could be merged faster in the next patch release.

Author: sachindshinde

.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