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

.changesets/fix_sachin_fix_shareable_mutation_fields_in_qp.md

@@ -0,0 +1,5 @@
+### Error when query planning a satisfiable @shareable mutation field ([PR #8352](https://github.com/apollographql/router/pull/8352))
+
+When query planning a mutation operation that executes a `@shareable` mutation field at top-level, query planning may unexpectedly error due to attempting to generate a plan where that mutation field is called more than once across multiple subgraphs. Query planning has now been updated to avoid generating such plans.
+
+By [@sachindshinde](https://github.com/sachindshinde) in https://github.com/apollographql/router/pull/8352

apollo-federation/src/query_graph/graph_path.rs

@@ -1607,11 +1607,14 @@ where
                     && matches!(
                         edge_weight.transition,
                         QueryGraphEdgeTransition::RootTypeResolution { .. }
+                            | QueryGraphEdgeTransition::KeyResolution
                     )
                     && !(to_advance.defer_on_tail.is_some()
                         && self.graph.is_self_key_or_root_edge(edge)?)
                 {
-                    debug!(r#"Ignored: edge is a top-level "RootTypeResolution""#);
+                    debug!(
+                        r#"Ignored: edge is a top-level "RootTypeResolution" or "KeyResolution""#
+                    );
                     continue;
                 }
 
@@ -1801,10 +1804,17 @@ where
                                     "Encountered non-root path with a subgraph-entering transition",
                                 ));
                             };
-                            self.graph
-                                .root_kinds_to_nodes_by_source(&edge_tail_weight.source)?
-                                .get(root_kind)
-                                .copied()
+                            // Since mutation options need to originate from the same subgraph, we
+                            // pretend we cannot find a root node  in another subgraph (effectively
+                            // skipping the optimization).
+                            if *root_kind == SchemaRootDefinitionKind::Mutation {
+                                None
+                            } else {
+                                self.graph
+                                    .root_kinds_to_nodes_by_source(&edge_tail_weight.source)?
+                                    .get(root_kind)
+                                    .copied()
+                            }
                         } else {
                             Some(last_subgraph_entering_edge_head)
                         };

apollo-federation/src/query_graph/graph_path/operation.rs

@@ -2289,6 +2289,7 @@ pub(crate) fn create_initial_options(
     excluded_edges: ExcludedDestinations,
     excluded_conditions: ExcludedConditions,
     override_conditions: &OverrideConditions,
+    initial_subgraph_constraint: Option<Arc<str>>,
     disabled_subgraphs: &IndexSet<Arc<str>>,
 ) -> Result<Vec<SimultaneousPathsWithLazyIndirectPaths>, FederationError> {
     let initial_paths = SimultaneousPaths::from(initial_path);
@@ -2310,8 +2311,18 @@ pub(crate) fn create_initial_options(
             .paths
             .iter()
             .cloned()
-            .map(SimultaneousPaths::from)
-            .collect();
+            .map(|path| {
+                let Some(initial_subgraph_constraint) = &initial_subgraph_constraint else {
+                    return Ok::<_, FederationError>(Some(path));
+                };
+                let tail_weight = path.graph.node_weight(path.tail)?;
+                Ok(if &tail_weight.source == initial_subgraph_constraint {
+                    Some(path)
+                } else {
+                    None
+                })
+            })
+            .process_results(|iter| iter.flatten().map(SimultaneousPaths::from).collect())?;
         Ok(lazy_initial_path.create_lazy_options(options, initial_context))
     } else {
         Ok(vec![lazy_initial_path])

apollo-federation/src/query_plan/query_planner.rs

@@ -9,6 +9,7 @@ use apollo_compiler::collections::IndexMap;
 use apollo_compiler::collections::IndexSet;
 use apollo_compiler::validation::Valid;
 use itertools::Itertools;
+use petgraph::visit::EdgeRef;
 use serde::Deserialize;
 use serde::Serialize;
 use tracing::trace;
@@ -576,7 +577,7 @@ impl QueryPlanner {
     }
 }
 
-fn compute_root_serial_dependency_graph(
+fn compute_root_serial_dependency_graph_for_mutation(
     parameters: &QueryPlanningParameters,
     has_defers: bool,
     non_local_selection_state: &mut Option<non_local_selections_estimation::State>,
@@ -606,7 +607,7 @@ fn compute_root_serial_dependency_graph(
         mut fetch_dependency_graph,
         path_tree: mut prev_path,
         ..
-    } = compute_root_parallel_best_plan(
+    } = compute_root_parallel_best_plan_for_mutation(
         parameters,
         selection_set,
         has_defers,
@@ -618,7 +619,7 @@ fn compute_root_serial_dependency_graph(
             fetch_dependency_graph: new_dep_graph,
             path_tree: new_path,
             ..
-        } = compute_root_parallel_best_plan(
+        } = compute_root_parallel_best_plan_for_mutation(
             parameters,
             selection_set,
             has_defers,
@@ -777,6 +778,7 @@ fn compute_root_parallel_best_plan(
         parameters.operation.root_kind,
         FetchDependencyGraphToCostProcessor,
         non_local_selection_state.as_mut(),
+        None,
     )?;
 
     // Getting no plan means the query is essentially unsatisfiable (it's a valid query, but we can prove it will never return a result),
@@ -786,6 +788,43 @@ fn compute_root_parallel_best_plan(
         .unwrap_or_else(|| BestQueryPlanInfo::empty(parameters)))
 }
 
+fn compute_root_parallel_best_plan_for_mutation(
+    parameters: &QueryPlanningParameters,
+    selection: SelectionSet,
+    has_defers: bool,
+    non_local_selection_state: &mut Option<non_local_selections_estimation::State>,
+) -> Result<BestQueryPlanInfo, FederationError> {
+    parameters.federated_query_graph.out_edges(parameters.head).into_iter().map(|edge_ref| {
+        let mutation_subgraph = parameters.federated_query_graph.node_weight(edge_ref.target())?.source.clone();
+        let planning_traversal = QueryPlanningTraversal::new(
+            parameters,
+            selection.clone(),
+            has_defers,
+            parameters.operation.root_kind,
+            FetchDependencyGraphToCostProcessor,
+            non_local_selection_state.as_mut(),
+            Some(mutation_subgraph),
+        )?;
+        planning_traversal.find_best_plan()
+    }).process_results(|iter| iter
+        .flatten()
+        .min_by(|a, b| a.cost.total_cmp(&b.cost))
+        .map(Ok)
+        .unwrap_or_else(|| {
+            if parameters.disabled_subgraphs.is_empty() {
+                Err(FederationError::internal(format!(
+                    "Was not able to plan {} starting from a single subgraph: This shouldn't have happened.",
+                    parameters.operation,
+                )))
+            } else {
+                // If subgraphs were disabled, this could be expected, and we indicate this in
+                // the error accordingly.
+                Err(SingleFederationError::NoPlanFoundWithDisabledSubgraphs.into())
+            }
+        })
+    )?
+}
+
 fn compute_plan_internal(
     parameters: &mut QueryPlanningParameters,
     processor: &mut FetchDependencyGraphToQueryPlanProcessor,
@@ -797,7 +836,7 @@ fn compute_plan_internal(
     let (main, deferred, primary_selection, cost) = if root_kind
         == SchemaRootDefinitionKind::Mutation
     {
-        let dependency_graphs = compute_root_serial_dependency_graph(
+        let dependency_graphs = compute_root_serial_dependency_graph_for_mutation(
             parameters,
             has_defers,
             non_local_selection_state,

apollo-federation/src/query_plan/query_planning_traversal.rs

@@ -222,6 +222,7 @@ impl<'a: 'b, 'b> QueryPlanningTraversal<'a, 'b> {
         root_kind: SchemaRootDefinitionKind,
         cost_processor: FetchDependencyGraphToCostProcessor,
         non_local_selection_state: Option<&mut non_local_selections_estimation::State>,
+        initial_subgraph_constraint: Option<Arc<str>>,
     ) -> Result<Self, FederationError> {
         Self::new_inner(
             parameters,
@@ -231,6 +232,7 @@ impl<'a: 'b, 'b> QueryPlanningTraversal<'a, 'b> {
             root_kind,
             cost_processor,
             non_local_selection_state,
+            initial_subgraph_constraint,
             Default::default(),
             Default::default(),
             Default::default(),
@@ -251,6 +253,7 @@ impl<'a: 'b, 'b> QueryPlanningTraversal<'a, 'b> {
         root_kind: SchemaRootDefinitionKind,
         cost_processor: FetchDependencyGraphToCostProcessor,
         non_local_selection_state: Option<&mut non_local_selections_estimation::State>,
+        initial_subgraph_constraint: Option<Arc<str>>,
         initial_context: OpGraphPathContext,
         excluded_destinations: ExcludedDestinations,
         excluded_conditions: ExcludedConditions,
@@ -304,14 +307,17 @@ impl<'a: 'b, 'b> QueryPlanningTraversal<'a, 'b> {
             excluded_destinations,
             excluded_conditions,
             &parameters.override_conditions,
+            initial_subgraph_constraint.clone(),
             &parameters.disabled_subgraphs,
         )?;
 
         traversal.open_branches = map_options_to_selections(selection_set, initial_options);
 
         if let Some(non_local_selection_state) = non_local_selection_state
-            && traversal
-                .check_non_local_selections_limit_exceeded_at_root(non_local_selection_state)?
+            && traversal.check_non_local_selections_limit_exceeded_at_root(
+                non_local_selection_state,
+                initial_subgraph_constraint.is_some(),
+            )?
         {
             return Err(SingleFederationError::QueryPlanComplexityExceeded {
                 message: format!(
@@ -519,8 +525,9 @@ impl<'a: 'b, 'b> QueryPlanningTraversal<'a, 'b> {
             // 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 process will generate 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.
-            return if self.is_top_level {
+            // 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.
+            return if self.is_top_level && self.root_kind != SchemaRootDefinitionKind::Mutation {
                 if self.parameters.disabled_subgraphs.is_empty() {
                     Err(FederationError::internal(format!(
                         "Was not able to find any options for {selection}: This shouldn't have happened.",
@@ -1161,6 +1168,7 @@ impl<'a: 'b, 'b> QueryPlanningTraversal<'a, 'b> {
             self.root_kind,
             self.cost_processor,
             None,
+            None,
             context.clone(),
             excluded_destinations.clone(),
             excluded_conditions.add_item(edge_conditions),