Miscellaneous @context/@fromcontext bugfixes (#6380)

This fixes a few bugs spotted around the native query planner implementation of `@context`/`@fromContext` bringing the implementation to a more consistent parity with the JavaScript implementation.

Co-authored-by: Iryna Shestak <shestak.irina@gmail.com>

apollo-federation/src/lib.rs

@@ -42,12 +42,15 @@ pub(crate) mod utils;
 use apollo_compiler::ast::NamedType;
 use apollo_compiler::validation::Valid;
 use apollo_compiler::Schema;
+use itertools::Itertools;
 use link::join_spec_definition::JOIN_VERSIONS;
 use schema::FederationSchema;
 
 pub use crate::api_schema::ApiSchemaOptions;
 use crate::error::FederationError;
 use crate::error::SingleFederationError;
+use crate::link::context_spec_definition::ContextSpecDefinition;
+use crate::link::context_spec_definition::CONTEXT_VERSIONS;
 use crate::link::join_spec_definition::JoinSpecDefinition;
 use crate::link::link_spec_definition::LinkSpecDefinition;
 use crate::link::spec::Identity;
@@ -59,18 +62,23 @@ use crate::subgraph::ValidSubgraph;
 pub use crate::supergraph::ValidFederationSubgraph;
 pub use crate::supergraph::ValidFederationSubgraphs;
 
-pub(crate) type SupergraphSpecs = (&'static LinkSpecDefinition, &'static JoinSpecDefinition);
+pub(crate) type SupergraphSpecs = (
+    &'static LinkSpecDefinition,
+    &'static JoinSpecDefinition,
+    Option<&'static ContextSpecDefinition>,
+);
 
 pub(crate) fn validate_supergraph_for_query_planning(
     supergraph_schema: &FederationSchema,
 ) -> Result<SupergraphSpecs, FederationError> {
-    validate_supergraph(supergraph_schema, &JOIN_VERSIONS)
+    validate_supergraph(supergraph_schema, &JOIN_VERSIONS, &CONTEXT_VERSIONS)
 }
 
 /// Checks that required supergraph directives are in the schema, and returns which ones were used.
 pub(crate) fn validate_supergraph(
     supergraph_schema: &FederationSchema,
     join_versions: &'static SpecDefinitions<JoinSpecDefinition>,
+    context_versions: &'static SpecDefinitions<ContextSpecDefinition>,
 ) -> Result<SupergraphSpecs, FederationError> {
     let Some(metadata) = supergraph_schema.metadata() else {
         return Err(SingleFederationError::InvalidFederationSupergraph {
@@ -94,7 +102,22 @@ pub(crate) fn validate_supergraph(
             ),
         }.into());
     };
-    Ok((link_spec_definition, join_spec_definition))
+    let context_spec_definition = metadata.for_identity(&Identity::context_identity()).map(|context_link| {
+        context_versions.find(&context_link.url.version).ok_or_else(|| {
+            SingleFederationError::InvalidFederationSupergraph {
+                message: format!(
+                    "Invalid supergraph: uses unsupported context spec version {} (supported versions: {})",
+                    context_link.url.version,
+                    context_versions.versions().join(", "),
+                ),
+            }
+        })
+    }).transpose()?;
+    Ok((
+        link_spec_definition,
+        join_spec_definition,
+        context_spec_definition,
+    ))
 }
 
 pub struct Supergraph {

apollo-federation/src/link/argument.rs

@@ -144,7 +144,7 @@ pub(crate) fn directive_optional_list_argument<'a>(
             Value::Null => Ok(None),
             Value::List(values) => Ok(Some(values.as_slice())),
             _ => bail!(
-                r#"Argument "{name}" of directive "@{}" must be a boolean."#,
+                r#"Argument "{name}" of directive "@{}" must be a list."#,
                 application.name
             ),
         },

apollo-federation/src/link/context_spec_definition.rs

@@ -1,8 +1,13 @@
+use apollo_compiler::ast::Directive;
+use apollo_compiler::ast::DirectiveDefinition;
 use apollo_compiler::name;
 use apollo_compiler::Name;
+use apollo_compiler::Node;
 use lazy_static::lazy_static;
 
 use crate::error::FederationError;
+use crate::internal_error;
+use crate::link::argument::directive_required_string_argument;
 use crate::link::spec::Identity;
 use crate::link::spec::Url;
 use crate::link::spec::Version;
@@ -11,7 +16,11 @@ use crate::link::spec_definition::SpecDefinitions;
 use crate::schema::FederationSchema;
 
 pub(crate) const CONTEXT_DIRECTIVE_NAME_IN_SPEC: Name = name!("context");
-pub(crate) const CONTEXT_DIRECTIVE_NAME_DEFAULT: Name = name!("federation__context");
+pub(crate) const CONTEXT_NAME_ARGUMENT_NAME: Name = name!("name");
+
+pub(crate) struct ContextDirectiveArguments<'doc> {
+    pub(crate) name: &'doc str,
+}
 
 #[derive(Clone)]
 pub(crate) struct ContextSpecDefinition {
@@ -30,13 +39,21 @@ impl ContextSpecDefinition {
         }
     }
 
-    pub(crate) fn context_directive_name_in_schema(
+    pub(crate) fn context_directive_definition<'schema>(
+        &self,
+        schema: &'schema FederationSchema,
+    ) -> Result<&'schema Node<DirectiveDefinition>, FederationError> {
+        self.directive_definition(schema, &CONTEXT_DIRECTIVE_NAME_IN_SPEC)?
+            .ok_or_else(|| internal_error!("Unexpectedly could not find context spec in schema"))
+    }
+
+    pub(crate) fn context_directive_arguments<'doc>(
         &self,
-        schema: &FederationSchema,
-    ) -> Result<Name, FederationError> {
-        Ok(self
-            .directive_name_in_schema(schema, &CONTEXT_DIRECTIVE_NAME_IN_SPEC)?
-            .unwrap_or(CONTEXT_DIRECTIVE_NAME_DEFAULT))
+        application: &'doc Node<Directive>,
+    ) -> Result<ContextDirectiveArguments<'doc>, FederationError> {
+        Ok(ContextDirectiveArguments {
+            name: directive_required_string_argument(application, &CONTEXT_NAME_ARGUMENT_NAME)?,
+        })
     }
 }
 

apollo-federation/src/link/federation_spec_definition.rs

@@ -469,15 +469,16 @@ impl FederationSpecDefinition {
         })
     }
 
-    pub(crate) fn context_directive_arguments(
-        application: &Node<Directive>,
-    ) -> Result<ContextDirectiveArguments, FederationError> {
+    pub(crate) fn context_directive_arguments<'doc>(
+        &self,
+        application: &'doc Node<Directive>,
+    ) -> Result<ContextDirectiveArguments<'doc>, FederationError> {
         Ok(ContextDirectiveArguments {
             name: directive_required_string_argument(application, &FEDERATION_NAME_ARGUMENT_NAME)?,
         })
     }
 
-    // The directive is named `@fromContex`. This is confusing for clippy, as
+    // The directive is named `@fromContext`. This is confusing for clippy, as
     // `from` is a conventional prefix used in conversion methods, which do not
     // take `self` as an argument. This function does **not** perform
     // conversion, but extracts `@fromContext` directive definition.
@@ -495,24 +496,7 @@ impl FederationSpecDefinition {
             })
     }
 
-    // The directive is named `@fromContex`. This is confusing for clippy, as
-    // `from` is a conventional prefix used in conversion methods, which do not
-    // take `self` as an argument. This function does **not** perform
-    // conversion, but extracts `@fromContext` directive arguments.
-    #[allow(clippy::wrong_self_convention)]
-    pub(crate) fn from_context_directive_arguments<'doc>(
-        &self,
-        application: &'doc Node<Directive>,
-    ) -> Result<FromContextDirectiveArguments<'doc>, FederationError> {
-        Ok(FromContextDirectiveArguments {
-            field: directive_required_string_argument(
-                application,
-                &FEDERATION_FIELD_ARGUMENT_NAME,
-            )?,
-        })
-    }
-
-    // The directive is named `@fromContex`. This is confusing for clippy, as
+    // The directive is named `@fromContext`. This is confusing for clippy, as
     // `from` is a conventional prefix used in conversion methods, which do not
     // take `self` as an argument. This function does **not** perform
     // conversion, but extracts `@fromContext` directive.
@@ -539,6 +523,23 @@ impl FederationSpecDefinition {
         })
     }
 
+    // The directive is named `@fromContext`. This is confusing for clippy, as
+    // `from` is a conventional prefix used in conversion methods, which do not
+    // take `self` as an argument. This function does **not** perform
+    // conversion, but extracts `@fromContext` directive arguments.
+    #[allow(clippy::wrong_self_convention)]
+    pub(crate) fn from_context_directive_arguments<'doc>(
+        &self,
+        application: &'doc Node<Directive>,
+    ) -> Result<FromContextDirectiveArguments<'doc>, FederationError> {
+        Ok(FromContextDirectiveArguments {
+            field: directive_required_string_argument(
+                application,
+                &FEDERATION_FIELD_ARGUMENT_NAME,
+            )?,
+        })
+    }
+
     pub(crate) fn get_cost_spec_definition(
         &self,
         schema: &FederationSchema,

apollo-federation/src/link/join_spec_definition.rs

@@ -49,7 +49,7 @@ pub(crate) const JOIN_OVERRIDE_LABEL_ARGUMENT_NAME: Name = name!("overrideLabel"
 pub(crate) const JOIN_USEROVERRIDDEN_ARGUMENT_NAME: Name = name!("usedOverridden");
 pub(crate) const JOIN_INTERFACE_ARGUMENT_NAME: Name = name!("interface");
 pub(crate) const JOIN_MEMBER_ARGUMENT_NAME: Name = name!("member");
-pub(crate) const JOIN_MEMBER_CONTEXT_ARGUMENTS: Name = name!("contextArguments");
+pub(crate) const JOIN_CONTEXTARGUMENTS_ARGUMENT_NAME: Name = name!("contextArguments");
 
 pub(crate) struct GraphDirectiveArguments<'doc> {
     pub(crate) name: &'doc str,
@@ -81,7 +81,9 @@ impl<'doc> TryFrom<&'doc Value> for ContextArgument<'doc> {
             value: &'a Value,
         ) -> Result<(), FederationError> {
             if let Some(first_value) = field {
-                bail!("Duplicate contextArgument for '{name}' field: {first_value} and {value}")
+                bail!(
+                    r#"Input field "{name}" in contextArguments is repeated with value "{value}" (previous value was "{first_value}")"#
+                )
             }
             let _ = field.insert(value);
             Ok(())
@@ -94,36 +96,36 @@ impl<'doc> TryFrom<&'doc Value> for ContextArgument<'doc> {
             field
                 .ok_or_else(|| {
                     FederationError::internal(format!(
-                        "'{field_name}' field was missing from contextArgument"
+                        r#"Input field "{field_name}" is missing from contextArguments"#
                     ))
                 })?
                 .as_str()
                 .ok_or_else(|| {
                     FederationError::internal(format!(
-                        "'{field_name}' field of contextArgument was not a string"
+                        r#"Input field "{field_name}" in contextArguments is not a string"#
                     ))
                 })
         }
 
-        let Value::Object(names) = value else {
-            bail!("Item in contextArgument list is not an object {value}")
+        let Value::Object(input_object) = value else {
+            bail!(r#"Item "{value}" in contextArguments list is not an object"#)
         };
         let mut name = None;
         let mut type_ = None;
         let mut context = None;
         let mut selection = None;
-        for (arg_name, value) in names.as_slice() {
-            match arg_name.as_str() {
-                "name" => insert_value(arg_name, &mut name, value)?,
-                "type" => insert_value(arg_name, &mut type_, value)?,
-                "context" => insert_value(arg_name, &mut context, value)?,
-                "selection" => insert_value(arg_name, &mut selection, value)?,
-                _ => bail!("Found unknown contextArgument {arg_name}"),
+        for (input_field_name, value) in input_object {
+            match input_field_name.as_str() {
+                "name" => insert_value(input_field_name, &mut name, value)?,
+                "type" => insert_value(input_field_name, &mut type_, value)?,
+                "context" => insert_value(input_field_name, &mut context, value)?,
+                "selection" => insert_value(input_field_name, &mut selection, value)?,
+                _ => bail!(r#"Found unknown contextArguments input field "{input_field_name}""#),
             }
         }
 
         let name = field_or_else("name", name)?;
-        let type_ = field_or_else("type_", type_)?;
+        let type_ = field_or_else("type", type_)?;
         let context = field_or_else("context", context)?;
         let selection = field_or_else("selection", selection)?;
 
@@ -308,7 +310,7 @@ impl JoinSpecDefinition {
             )?,
             context_arguments: directive_optional_list_argument(
                 application,
-                &JOIN_MEMBER_CONTEXT_ARGUMENTS,
+                &JOIN_CONTEXTARGUMENTS_ARGUMENT_NAME,
             )?
             .map(|values| {
                 values

apollo-federation/src/operation/rebase.rs

@@ -22,7 +22,6 @@ use super::SelectionSet;
 use super::TYPENAME_FIELD;
 use crate::ensure;
 use crate::error::FederationError;
-use crate::link::federation_spec_definition::get_federation_spec_definition_from_subgraph;
 use crate::schema::position::CompositeTypeDefinitionPosition;
 use crate::schema::position::OutputTypeDefinitionPosition;
 use crate::schema::ValidFederationSchema;
@@ -276,15 +275,17 @@ impl Field {
         else {
             return Ok(None);
         };
-        if let Ok(federation_spec_definition) = get_federation_spec_definition_from_subgraph(schema)
+        if let Some(federation_spec_definition) = schema
+            .subgraph_metadata()
+            .map(|d| d.federation_spec_definition())
         {
             let from_context_directive_definition_name = &federation_spec_definition
                 .from_context_directive_definition(schema)?
                 .name;
-            // We need to prevent arguments with `@fromContext` from being lost/overwriten. If the
-            // would-be parent type's field has `@fromContext` and one (or more) of this field's
-            // arguments doesn't exist in the would-be parent's field, rebasing would loose that
-            // context.
+            // We need to ensure that all arguments with `@fromContext` are provided. If the
+            // would-be parent type's field has an argument with `@fromContext` and that argument
+            // has no value/data in this field, then we return `None` to indicate the rebase isn't
+            // possible.
             if field_definition.arguments.iter().any(|arg_definition| {
                 arg_definition
                     .directives