Fix query planning for @Shareable mutation fields

Author: sachindshinde

apollo-federation/src/query_graph/graph_path.rs

@@ -1801,10 +1801,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,6 +307,7 @@ impl<'a: 'b, 'b> QueryPlanningTraversal<'a, 'b> {
             excluded_destinations,
             excluded_conditions,
             &parameters.override_conditions,
+            initial_subgraph_constraint,
             &parameters.disabled_subgraphs,
         )?;
 
@@ -519,8 +523,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 +1166,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),

apollo-federation/tests/query_plan/build_query_plan_tests.rs

@@ -478,6 +478,77 @@ fn it_executes_mutation_operations_in_sequence() {
     );
 }
 
+#[test]
+fn it_executes_a_single_mutation_operation_on_a_shareable_field() {
+    let planner = planner!(
+        Subgraph1: r#"
+          type Query {
+            dummy: Int
+          }
+
+          type Mutation {
+            f: F @shareable
+          }
+
+          type F @key(fields: "id") {
+            id: ID!
+            x: Int
+          }
+        "#,
+        Subgraph2: r#"
+          type Mutation {
+            f: F @shareable
+          }
+
+          type F @key(fields: "id", resolvable: false) {
+            id: ID!
+            y: Int
+          }
+        "#,
+    );
+    assert_plan!(
+        &planner,
+        r#"
+          mutation {
+            f {
+              x
+              y
+            }
+          }
+        "#,
+        @r###"
+        QueryPlan {
+          Sequence {
+            Fetch(service: "Subgraph2") {
+              {
+                f {
+                  __typename
+                  id
+                  y
+                }
+              }
+            },
+            Flatten(path: "f") {
+              Fetch(service: "Subgraph1") {
+                {
+                  ... on F {
+                    __typename
+                    id
+                  }
+                } =>
+                {
+                  ... on F {
+                    x
+                  }
+                }
+              },
+            },
+          },
+        }
+      "###
+    );
+}
+
 /// @requires references external field indirectly
 #[test]
 fn key_where_at_external_is_not_at_top_level_of_selection_of_requires() {

apollo-federation/tests/query_plan/supergraphs/it_executes_a_single_mutation_operation_on_a_shareable_field.graphql

@@ -0,0 +1,79 @@
+# Composed from subgraphs with hash: 7206fa0505b6ab6548c5d6460861c3fcdc746b91
+schema
+  @link(url: "https://specs.apollo.dev/link/v1.0")
+  @link(url: "https://specs.apollo.dev/join/v0.5", for: EXECUTION)
+{
+  query: Query
+  mutation: Mutation
+}
+
+directive @join__directive(graphs: [join__Graph!], name: String!, args: join__DirectiveArguments) repeatable on SCHEMA | OBJECT | INTERFACE | FIELD_DEFINITION
+
+directive @join__enumValue(graph: join__Graph!) repeatable on ENUM_VALUE
+
+directive @join__field(graph: join__Graph, requires: join__FieldSet, provides: join__FieldSet, type: String, external: Boolean, override: String, usedOverridden: Boolean, overrideLabel: String, contextArguments: [join__ContextArgument!]) repeatable on FIELD_DEFINITION | INPUT_FIELD_DEFINITION
+
+directive @join__graph(name: String!, url: String!) on ENUM_VALUE
+
+directive @join__implements(graph: join__Graph!, interface: String!) repeatable on OBJECT | INTERFACE
+
+directive @join__type(graph: join__Graph!, key: join__FieldSet, extension: Boolean! = false, resolvable: Boolean! = true, isInterfaceObject: Boolean! = false) repeatable on OBJECT | INTERFACE | UNION | ENUM | INPUT_OBJECT | SCALAR
+
+directive @join__unionMember(graph: join__Graph!, member: String!) repeatable on UNION
+
+directive @link(url: String, as: String, for: link__Purpose, import: [link__Import]) repeatable on SCHEMA
+
+type F
+  @join__type(graph: SUBGRAPH1, key: "id")
+  @join__type(graph: SUBGRAPH2, key: "id", resolvable: false)
+{
+  id: ID!
+  x: Int @join__field(graph: SUBGRAPH1)
+  y: Int @join__field(graph: SUBGRAPH2)
+}
+
+input join__ContextArgument {
+  name: String!
+  type: String!
+  context: String!
+  selection: join__FieldValue!
+}
+
+scalar join__DirectiveArguments
+
+scalar join__FieldSet
+
+scalar join__FieldValue
+
+enum join__Graph {
+  SUBGRAPH1 @join__graph(name: "Subgraph1", url: "none")
+  SUBGRAPH2 @join__graph(name: "Subgraph2", url: "none")
+}
+
+scalar link__Import
+
+enum link__Purpose {
+  """
+  `SECURITY` features provide metadata necessary to securely resolve fields.
+  """
+  SECURITY
+
+  """
+  `EXECUTION` features provide metadata necessary for operation execution.
+  """
+  EXECUTION
+}
+
+type Mutation
+  @join__type(graph: SUBGRAPH1)
+  @join__type(graph: SUBGRAPH2)
+{
+  f: F
+}
+
+type Query
+  @join__type(graph: SUBGRAPH1)
+  @join__type(graph: SUBGRAPH2)
+{
+  dummy: Int @join__field(graph: SUBGRAPH1)
+}