feat(federation): add support for progressive overrides (#6011)

Implement support for progressive `@override`s (apollographql/federation#2902).

Co-authored-by: Renée <[email protected]>

apollo-federation/cli/src/bench.rs

@@ -54,7 +54,7 @@ pub(crate) fn run_bench(
             }
         };
         let now = Instant::now();
-        let plan = planner.build_query_plan(&document, None);
+        let plan = planner.build_query_plan(&document, None, Default::default());
         let elapsed = now.elapsed().as_secs_f64() * 1000.0;
         let mut eval_plans = None;
         let mut error = None;

apollo-federation/cli/src/main.rs

@@ -254,7 +254,10 @@ fn cmd_plan(
 
     let query_doc =
         ExecutableDocument::parse_and_validate(planner.api_schema().schema(), query, query_path)?;
-    print!("{}", planner.build_query_plan(&query_doc, None)?);
+    print!(
+        "{}",
+        planner.build_query_plan(&query_doc, None, Default::default())?
+    );
     Ok(())
 }
 

apollo-federation/src/error/mod.rs

@@ -33,8 +33,6 @@ impl From<SchemaRootKind> for String {
 
 #[derive(Clone, Debug, strum_macros::Display, PartialEq, Eq)]
 pub enum UnsupportedFeatureKind {
-    #[strum(to_string = "progressive overrides")]
-    ProgressiveOverrides,
     #[strum(to_string = "defer")]
     Defer,
     #[strum(to_string = "context")]

apollo-federation/src/link/federation_spec_definition.rs

@@ -12,6 +12,7 @@ use lazy_static::lazy_static;
 use crate::error::FederationError;
 use crate::error::SingleFederationError;
 use crate::link::argument::directive_optional_boolean_argument;
+use crate::link::argument::directive_optional_string_argument;
 use crate::link::argument::directive_required_string_argument;
 use crate::link::cost_spec_definition::CostSpecDefinition;
 use crate::link::cost_spec_definition::COST_VERSIONS;
@@ -51,6 +52,11 @@ pub(crate) struct ProvidesDirectiveArguments<'doc> {
     pub(crate) fields: &'doc str,
 }
 
+pub(crate) struct OverrideDirectiveArguments<'doc> {
+    pub(crate) from: &'doc str,
+    pub(crate) label: Option<&'doc str>,
+}
+
 #[derive(Debug)]
 pub(crate) struct FederationSpecDefinition {
     url: Url,
@@ -361,6 +367,19 @@ impl FederationSpecDefinition {
         })
     }
 
+    pub(crate) fn override_directive_definition<'schema>(
+        &self,
+        schema: &'schema FederationSchema,
+    ) -> Result<&'schema Node<DirectiveDefinition>, FederationError> {
+        self.directive_definition(schema, &FEDERATION_OVERRIDE_DIRECTIVE_NAME_IN_SPEC)?
+            .ok_or_else(|| {
+                FederationError::internal(format!(
+                    "Unexpectedly could not find federation spec's \"@{}\" directive definition",
+                    FEDERATION_OVERRIDE_DIRECTIVE_NAME_IN_SPEC
+                ))
+            })
+    }
+
     pub(crate) fn override_directive(
         &self,
         schema: &FederationSchema,
@@ -390,6 +409,19 @@ impl FederationSpecDefinition {
         })
     }
 
+    pub(crate) fn override_directive_arguments<'doc>(
+        &self,
+        application: &'doc Node<Directive>,
+    ) -> Result<OverrideDirectiveArguments<'doc>, FederationError> {
+        Ok(OverrideDirectiveArguments {
+            from: directive_required_string_argument(application, &FEDERATION_FROM_ARGUMENT_NAME)?,
+            label: directive_optional_string_argument(
+                application,
+                &FEDERATION_OVERRIDE_LABEL_ARGUMENT_NAME,
+            )?,
+        })
+    }
+
     pub(crate) fn get_cost_spec_definition(
         &self,
         schema: &FederationSchema,

apollo-federation/src/query_graph/build_query_graph.rs

@@ -1,5 +1,6 @@
 use std::sync::Arc;
 
+use apollo_compiler::collections::HashMap;
 use apollo_compiler::collections::IndexMap;
 use apollo_compiler::collections::IndexSet;
 use apollo_compiler::schema::DirectiveList as ComponentDirectiveList;
@@ -21,6 +22,7 @@ use crate::link::federation_spec_definition::KeyDirectiveArguments;
 use crate::operation::merge_selection_sets;
 use crate::operation::Selection;
 use crate::operation::SelectionSet;
+use crate::query_graph::OverrideCondition;
 use crate::query_graph::QueryGraph;
 use crate::query_graph::QueryGraphEdge;
 use crate::query_graph::QueryGraphEdgeTransition;
@@ -140,6 +142,7 @@ impl BaseQueryGraphBuilder {
             QueryGraphEdge {
                 transition,
                 conditions,
+                override_condition: None,
             },
         );
         let head_weight = self.query_graph.node_weight(head)?;
@@ -982,6 +985,7 @@ impl FederatedQueryGraphBuilder {
         self.add_root_edges()?;
         self.handle_key()?;
         self.handle_requires()?;
+        self.handle_progressive_overrides()?;
         // Note that @provides must be handled last when building since it requires copying nodes
         // and their edges, and it's easier to reason about this if we know previous
         self.handle_provides()?;
@@ -1374,6 +1378,102 @@ impl FederatedQueryGraphBuilder {
         Ok(())
     }
 
+    /// Handling progressive overrides here. For each progressive @override
+    /// application (with a label), we want to update the edges to the overridden
+    /// field within the "to" and "from" subgraphs with their respective override
+    /// condition (the label and a T/F value). The "from" subgraph will have an
+    /// override condition of `false`, whereas the "to" subgraph will have an
+    /// override condition of `true`.
+    fn handle_progressive_overrides(&mut self) -> Result<(), FederationError> {
+        let mut edge_to_conditions: HashMap<EdgeIndex, OverrideCondition> = Default::default();
+
+        fn collect_edge_condition(
+            query_graph: &QueryGraph,
+            target_graph: &str,
+            target_field: &ObjectFieldDefinitionPosition,
+            label: &str,
+            condition: bool,
+            edge_to_conditions: &mut HashMap<EdgeIndex, OverrideCondition>,
+        ) -> Result<(), FederationError> {
+            let target_field = FieldDefinitionPosition::Object(target_field.clone());
+            let subgraph_nodes = query_graph
+                .types_to_nodes_by_source
+                .get(target_graph)
+                .unwrap();
+            let parent_node = subgraph_nodes
+                .get(target_field.type_name())
+                .unwrap()
+                .first()
+                .unwrap();
+            for edge in query_graph.out_edges(*parent_node) {
+                let edge_weight = query_graph.edge_weight(edge.id())?;
+                let QueryGraphEdgeTransition::FieldCollection {
+                    field_definition_position,
+                    ..
+                } = &edge_weight.transition
+                else {
+                    continue;
+                };
+
+                if &target_field == field_definition_position {
+                    edge_to_conditions.insert(
+                        edge.id(),
+                        OverrideCondition {
+                            label: label.to_string(),
+                            condition,
+                        },
+                    );
+                }
+            }
+            Ok(())
+        }
+
+        for (to_subgraph_name, subgraph) in &self.base.query_graph.subgraphs_by_name {
+            let subgraph_data = self.subgraphs.get(to_subgraph_name)?;
+            if let Some(override_referencers) = subgraph
+                .referencers()
+                .directives
+                .get(&subgraph_data.overrides_directive_definition_name)
+            {
+                for field_definition_position in &override_referencers.object_fields {
+                    let field = field_definition_position.get(subgraph.schema())?;
+                    for directive in field
+                        .directives
+                        .get_all(&subgraph_data.overrides_directive_definition_name)
+                    {
+                        let application = subgraph_data
+                            .federation_spec_definition
+                            .override_directive_arguments(directive)?;
+                        if let Some(label) = application.label {
+                            collect_edge_condition(
+                                &self.base.query_graph,
+                                to_subgraph_name,
+                                field_definition_position,
+                                label,
+                                true,
+                                &mut edge_to_conditions,
+                            )?;
+                            collect_edge_condition(
+                                &self.base.query_graph,
+                                application.from,
+                                field_definition_position,
+                                label,
+                                false,
+                                &mut edge_to_conditions,
+                            )?;
+                        }
+                    }
+                }
+            }
+        }
+
+        for (edge, condition) in edge_to_conditions {
+            let mutable_edge = self.base.query_graph.edge_weight_mut(edge)?;
+            mutable_edge.override_condition = Some(condition);
+        }
+        Ok(())
+    }
+
     /// Handle @provides by copying the appropriate nodes/edges.
     fn handle_provides(&mut self) -> Result<(), FederationError> {
         let mut provide_id = 0;
@@ -1987,6 +2087,10 @@ impl FederatedQueryGraphBuilderSubgraphs {
                         ),
                     }
                 })?;
+            let overrides_directive_definition_name = federation_spec_definition
+                .override_directive_definition(schema)?
+                .name
+                .clone();
             subgraphs.map.insert(
                 source.clone(),
                 FederatedQueryGraphBuilderSubgraphData {
@@ -1995,6 +2099,7 @@ impl FederatedQueryGraphBuilderSubgraphs {
                     requires_directive_definition_name,
                     provides_directive_definition_name,
                     interface_object_directive_definition_name,
+                    overrides_directive_definition_name,
                 },
             );
         }
@@ -2020,6 +2125,7 @@ struct FederatedQueryGraphBuilderSubgraphData {
     requires_directive_definition_name: Name,
     provides_directive_definition_name: Name,
     interface_object_directive_definition_name: Name,
+    overrides_directive_definition_name: Name,
 }
 
 #[derive(Debug)]