Implemented extract_defer_from_operation function (apollographql/fe…

…deration-next#284)

- ported `extractDeferFromOperation`

apollo-federation/src/link/mod.rs

@@ -10,7 +10,7 @@ use std::str;
 use std::{collections::HashMap, sync::Arc};
 use thiserror::Error;
 
-mod argument;
+pub(crate) mod argument;
 pub mod database;
 pub(crate) mod federation_spec_definition;
 pub(crate) mod graphql_definition;

apollo-federation/src/query_graph/graph_path.rs

@@ -330,6 +330,49 @@ impl OpPathElement {
             OpPathElement::InlineFragment(inline_fragment) => inline_fragment.as_path_element(),
         }
     }
+
+    pub(crate) fn defer_directive_args(&self) -> Option<DeferDirectiveArguments> {
+        match self {
+            OpPathElement::Field(_) => None, // @defer cannot be on field at the moment
+            OpPathElement::InlineFragment(inline_fragment) => inline_fragment
+                .data()
+                .defer_directive_arguments()
+                .ok()
+                .flatten(),
+        }
+    }
+
+    /// Returns this fragment element but with any @defer directive on it removed.
+    ///
+    /// This method will return `None` if, upon removing @defer, the fragment has no conditions nor
+    /// any remaining applied directives (meaning that it carries no information whatsoever and can be
+    /// ignored).
+    pub(crate) fn without_defer(&self) -> Option<Self> {
+        match self {
+            Self::Field(_) => Some(self.clone()), // unchanged
+            Self::InlineFragment(inline_fragment) => {
+                let updated_directives: DirectiveList = inline_fragment
+                    .data()
+                    .directives
+                    .get_all("defer")
+                    .cloned()
+                    .collect();
+                if inline_fragment.data().type_condition_position.is_none()
+                    && updated_directives.is_empty()
+                {
+                    return None;
+                }
+                if inline_fragment.data().directives.len() == updated_directives.len() {
+                    Some(self.clone())
+                } else {
+                    // PORT_NOTE: We won't need to port `this.copyAttachementsTo(updated);` line here
+                    // since `with_updated_directives` clones the whole `self` and thus sibling
+                    // type names should be copied as well.
+                    Some(self.with_updated_directives(updated_directives))
+                }
+            }
+        }
+    }
 }
 
 pub(crate) fn selection_of_element(

apollo-federation/src/query_plan/fetch_dependency_graph.rs

@@ -192,23 +192,16 @@ pub(crate) struct DeferTracking {
 #[derive(Debug, Clone)]
 pub(crate) struct DeferredInfo {
     pub(crate) label: DeferRef,
-    pub(crate) path: FetchDependencyGraphPath,
+    pub(crate) path: FetchDependencyGraphNodePath,
     pub(crate) sub_selection: NormalizedSelectionSet,
     pub(crate) deferred: IndexSet<DeferRef>,
     pub(crate) dependencies: IndexSet<DeferRef>,
 }
 
 // TODO: Write docstrings
-#[derive(Debug, Clone)]
-pub(crate) struct FetchDependencyGraphPath {
-    pub(crate) full_path: OpPath,
-    pub(crate) path_in_node: OpPath,
-    pub(crate) response_path: Vec<FetchDataPathElement>,
-}
-
 #[derive(Debug, Clone, Default)]
 pub(crate) struct FetchDependencyGraphNodePath {
-    full_path: Arc<OpPath>,
+    pub(crate) full_path: Arc<OpPath>,
     path_in_node: Arc<OpPath>,
     response_path: Vec<FetchDataPathElement>,
 }
@@ -1640,7 +1633,7 @@ impl DeferTracking {
         &mut self,
         defer_context: &DeferContext,
         defer_args: &DeferDirectiveArguments,
-        path: FetchDependencyGraphPath,
+        path: FetchDependencyGraphNodePath,
         parent_type: CompositeTypeDefinitionPosition,
     ) {
         // Having the primary selection undefined means that @defer handling is actually disabled, so there's no need to track anything.
@@ -1734,7 +1727,7 @@ impl DeferredInfo {
     fn empty(
         schema: ValidFederationSchema,
         label: DeferRef,
-        path: FetchDependencyGraphPath,
+        path: FetchDependencyGraphNodePath,
         parent_type: CompositeTypeDefinitionPosition,
     ) -> Self {
         Self {
@@ -2120,7 +2113,7 @@ fn compute_nodes_for_op_path_element<'a>(
             operation,
             &stack_item.defer_context,
             &stack_item.node_path,
-        );
+        )?;
         // We've now removed any @defer.
         // If the operation contains other directives or a non-trivial type condition,
         // we need to preserve it and so we add operation.
@@ -2192,7 +2185,7 @@ fn compute_nodes_for_op_path_element<'a>(
             None,
         )
     }
-    let (Some(updated_operation), updated_defer_context) = extract_defer_from_operation(
+    let Ok((Some(updated_operation), updated_defer_context)) = extract_defer_from_operation(
         dependency_graph,
         operation,
         &stack_item.defer_context,
@@ -2429,13 +2422,55 @@ fn compute_input_rewrites_on_key_fetch(
     }
 }
 
+/// Returns an updated pair of (`operation`, `defer_context`) after the `defer` directive removed.
+/// - The updated operation can be `None`, if operation is no longer necessary.
 fn extract_defer_from_operation(
-    _dependency_graph: &mut FetchDependencyGraph,
-    _operation: &OpPathElement,
-    _defer_context: &DeferContext,
-    _node_path: &FetchDependencyGraphNodePath,
-) -> (Option<OpPathElement>, DeferContext) {
-    todo!() // Port `extractDeferFromOperation`
+    dependency_graph: &mut FetchDependencyGraph,
+    operation: &OpPathElement,
+    defer_context: &DeferContext,
+    node_path: &FetchDependencyGraphNodePath,
+) -> Result<(Option<OpPathElement>, DeferContext), FederationError> {
+    let defer_args = operation.defer_directive_args();
+    let Some(defer_args) = defer_args else {
+        let updated_path_to_defer_parent = defer_context
+            .path_to_defer_parent
+            .with_pushed(operation.clone().into());
+        let updated_context = DeferContext {
+            path_to_defer_parent: updated_path_to_defer_parent.into(),
+            // Following fields are identical to those of `defer_context`.
+            current_defer_ref: defer_context.current_defer_ref.clone(),
+            active_defer_ref: defer_context.active_defer_ref.clone(),
+            is_part_of_query: defer_context.is_part_of_query,
+        };
+        return Ok((Some(operation.clone()), updated_context));
+    };
+
+    let updated_defer_ref = defer_args.label().ok_or_else(||
+        // PORT_NOTE: The original TypeScript code has an assertion here.
+        FederationError::internal(
+                    "All defers should have a label at this point",
+                ))?;
+    let updated_operation = operation.without_defer();
+    let updated_path_to_defer_parent = match updated_operation {
+        None => Default::default(), // empty OpPath
+        Some(ref updated_operation) => OpPath(vec![Arc::new(updated_operation.clone())]),
+    };
+
+    dependency_graph.defer_tracking.register_defer(
+        defer_context,
+        &defer_args,
+        node_path.clone(),
+        operation.parent_type_position(),
+    );
+
+    let updated_context = DeferContext {
+        current_defer_ref: Some(updated_defer_ref.into()),
+        path_to_defer_parent: updated_path_to_defer_parent.into(),
+        // Following fields are identical to those of `defer_context`.
+        active_defer_ref: defer_context.active_defer_ref.clone(),
+        is_part_of_query: defer_context.is_part_of_query,
+    };
+    Ok((updated_operation, updated_context))
 }
 
 fn handle_requires(

apollo-federation/src/query_plan/fetch_dependency_graph_processor.rs

@@ -346,7 +346,7 @@ impl FetchDependencyGraphProcessor<Option<PlanNode>, DeferredDeferBlock>
             } else {
                 Some(defer_info.label.clone())
             },
-            query_path: (&defer_info.path.full_path).try_into()?,
+            query_path: defer_info.path.full_path.as_ref().try_into()?,
             // Note that if the deferred block has nested @defer,
             // then the `value` is going to be a `DeferNode`
             // and we'll use it's own `subselection`, so we don't need it here.