Ensure cost directives are picked up when not explicitly imported

apollo-federation/src/link/cost_spec_definition.rs

@@ -1,16 +1,20 @@
+use std::collections::HashSet;
+
 use apollo_compiler::ast::Argument;
 use apollo_compiler::ast::Directive;
-use apollo_compiler::collections::IndexMap;
+use apollo_compiler::ast::DirectiveList;
+use apollo_compiler::ast::FieldDefinition;
+use apollo_compiler::ast::InputValueDefinition;
 use apollo_compiler::name;
 use apollo_compiler::schema::Component;
-use apollo_compiler::schema::EnumType;
-use apollo_compiler::schema::ObjectType;
-use apollo_compiler::schema::ScalarType;
+use apollo_compiler::schema::ExtendedType;
 use apollo_compiler::Name;
 use apollo_compiler::Node;
 use lazy_static::lazy_static;
 
 use crate::error::FederationError;
+use crate::error::SingleFederationError;
+use crate::link::federation_spec_definition::get_federation_spec_definition_from_subgraph;
 use crate::link::spec::Identity;
 use crate::link::spec::Url;
 use crate::link::spec::Version;
@@ -21,42 +25,42 @@ use crate::schema::position::ObjectTypeDefinitionPosition;
 use crate::schema::position::ScalarTypeDefinitionPosition;
 use crate::schema::FederationSchema;
 
-pub(crate) const COST_DIRECTIVE_NAME_IN_SPEC: Name = name!("cost");
-pub(crate) const COST_DIRECTIVE_NAME_DEFAULT: Name = name!("federation__cost");
-
-pub(crate) const LIST_SIZE_DIRECTIVE_NAME_IN_SPEC: Name = name!("listSize");
-pub(crate) const LIST_SIZE_DIRECTIVE_NAME_DEFAULT: Name = name!("federation__listSize");
+const COST_DIRECTIVE_NAME: Name = name!("cost");
+const COST_DIRECTIVE_WEIGHT_ARGUMENT_NAME: Name = name!("weight");
+const LIST_SIZE_DIRECTIVE_NAME: Name = name!("listSize");
+const LIST_SIZE_DIRECTIVE_ASSUMED_SIZE_ARGUMENT_NAME: Name = name!("assumedSize");
+const LIST_SIZE_DIRECTIVE_SLICING_ARGUMENTS_ARGUMENT_NAME: Name = name!("slicingArguments");
+const LIST_SIZE_DIRECTIVE_SIZED_FIELDS_ARGUMENT_NAME: Name = name!("sizedFields");
+const LIST_SIZE_DIRECTIVE_REQUIRE_ONE_SLICING_ARGUMENT_ARGUMENT_NAME: Name =
+    name!("requireOneSlicingArgument");
 
 #[derive(Clone)]
-pub(crate) struct CostSpecDefinition {
+pub struct CostSpecDefinition {
     url: Url,
     minimum_federation_version: Option<Version>,
 }
 
 macro_rules! propagate_demand_control_directives {
     ($func_name:ident, $directives_ty:ty, $wrap_ty:expr) => {
         pub(crate) fn $func_name(
-            &self,
-            subgraph_schema: &FederationSchema,
+            supergraph_schema: &FederationSchema,
             source: &$directives_ty,
+            subgraph_schema: &FederationSchema,
             dest: &mut $directives_ty,
-            original_directive_names: &IndexMap<Name, Name>,
         ) -> Result<(), FederationError> {
-            let cost_directive_name = original_directive_names.get(&COST_DIRECTIVE_NAME_IN_SPEC);
-            let cost_directive = cost_directive_name.and_then(|name| source.get(name.as_str()));
+            let cost_directive = Self::cost_directive_name(supergraph_schema)?
+                .and_then(|name| source.get(name.as_str()));
             if let Some(cost_directive) = cost_directive {
-                dest.push($wrap_ty(self.cost_directive(
+                dest.push($wrap_ty(Self::cost_directive(
                     subgraph_schema,
                     cost_directive.arguments.clone(),
                 )?));
             }
 
-            let list_size_directive_name =
-                original_directive_names.get(&LIST_SIZE_DIRECTIVE_NAME_IN_SPEC);
-            let list_size_directive =
-                list_size_directive_name.and_then(|name| source.get(name.as_str()));
+            let list_size_directive = Self::list_size_directive_name(supergraph_schema)?
+                .and_then(|name| source.get(name.as_str()));
             if let Some(list_size_directive) = list_size_directive {
-                dest.push($wrap_ty(self.list_size_directive(
+                dest.push($wrap_ty(Self::list_size_directive(
                     subgraph_schema,
                     list_size_directive.arguments.clone(),
                 )?));
@@ -68,34 +72,31 @@ macro_rules! propagate_demand_control_directives {
 }
 
 macro_rules! propagate_demand_control_directives_to_position {
-    ($func_name:ident, $source_ty:ty, $dest_ty:ty) => {
+    ($func_name:ident, $source_ty:ty, $pos_ty:ty) => {
         pub(crate) fn $func_name(
-            &self,
+            supergraph_schema: &FederationSchema,
             subgraph_schema: &mut FederationSchema,
-            source: &Node<$source_ty>,
-            dest: &$dest_ty,
-            original_directive_names: &IndexMap<Name, Name>,
+            pos: &$pos_ty,
         ) -> Result<(), FederationError> {
-            let cost_directive_name = original_directive_names.get(&COST_DIRECTIVE_NAME_IN_SPEC);
-            let cost_directive =
-                cost_directive_name.and_then(|name| source.directives.get(name.as_str()));
+            let source = pos.get(supergraph_schema.schema())?;
+            let cost_directive = Self::cost_directive_name(supergraph_schema)?
+                .and_then(|name| source.directives.get(name.as_str()));
             if let Some(cost_directive) = cost_directive {
-                dest.insert_directive(
+                pos.insert_directive(
                     subgraph_schema,
-                    Component::from(
-                        self.cost_directive(subgraph_schema, cost_directive.arguments.clone())?,
-                    ),
+                    Component::from(Self::cost_directive(
+                        subgraph_schema,
+                        cost_directive.arguments.clone(),
+                    )?),
                 )?;
             }
 
-            let list_size_directive_name =
-                original_directive_names.get(&LIST_SIZE_DIRECTIVE_NAME_IN_SPEC);
-            let list_size_directive =
-                list_size_directive_name.and_then(|name| source.directives.get(name.as_str()));
+            let list_size_directive = Self::list_size_directive_name(supergraph_schema)?
+                .and_then(|name| source.directives.get(name.as_str()));
             if let Some(list_size_directive) = list_size_directive {
-                dest.insert_directive(
+                pos.insert_directive(
                     subgraph_schema,
-                    Component::from(self.list_size_directive(
+                    Component::from(Self::list_size_directive(
                         subgraph_schema,
                         list_size_directive.arguments.clone(),
                     )?),
@@ -119,35 +120,36 @@ impl CostSpecDefinition {
     }
 
     pub(crate) fn cost_directive(
-        &self,
         schema: &FederationSchema,
         arguments: Vec<Node<Argument>>,
     ) -> Result<Directive, FederationError> {
-        let name = self
-            .directive_name_in_schema(schema, &COST_DIRECTIVE_NAME_IN_SPEC)?
-            .unwrap_or(COST_DIRECTIVE_NAME_DEFAULT);
+        let name =
+            Self::cost_directive_name(schema)?.ok_or_else(|| SingleFederationError::Internal {
+                message: "The \"@cost\" directive is undefined in the target schema".to_string(),
+            })?;
 
         Ok(Directive { name, arguments })
     }
 
     pub(crate) fn list_size_directive(
-        &self,
         schema: &FederationSchema,
         arguments: Vec<Node<Argument>>,
     ) -> Result<Directive, FederationError> {
-        let name = self
-            .directive_name_in_schema(schema, &LIST_SIZE_DIRECTIVE_NAME_IN_SPEC)?
-            .unwrap_or(LIST_SIZE_DIRECTIVE_NAME_DEFAULT);
+        let name = Self::list_size_directive_name(schema)?.ok_or_else(|| {
+            SingleFederationError::Internal {
+                message: "The \"@listSize\" directive is undefined in the target schema"
+                    .to_string(),
+            }
+        })?;
 
         Ok(Directive { name, arguments })
     }
 
     propagate_demand_control_directives!(
         propagate_demand_control_directives,
-        apollo_compiler::ast::DirectiveList,
+        DirectiveList,
         Node::new
     );
-
     propagate_demand_control_directives_to_position!(
         propagate_demand_control_directives_for_enum,
         EnumType,
@@ -163,6 +165,81 @@ impl CostSpecDefinition {
         ScalarType,
         ScalarTypeDefinitionPosition
     );
+
+    fn for_federation_schema(schema: &FederationSchema) -> Option<&'static Self> {
+        let link = schema
+            .metadata()?
+            .for_identity(&Identity::cost_identity())?;
+        COST_VERSIONS.find(&link.url.version)
+    }
+
+    /// Returns the name of the `@cost` directive in the given schema, accounting for import aliases or specification name
+    /// prefixes such as `@federation__cost`. This checks the linked cost specification, if there is one, and falls back
+    /// to the federation spec.
+    fn cost_directive_name(schema: &FederationSchema) -> Result<Option<Name>, FederationError> {
+        if let Some(spec) = Self::for_federation_schema(schema) {
+            spec.directive_name_in_schema(schema, &COST_DIRECTIVE_NAME)
+        } else if let Ok(fed_spec) = get_federation_spec_definition_from_subgraph(schema) {
+            fed_spec.directive_name_in_schema(schema, &COST_DIRECTIVE_NAME)
+        } else {
+            Ok(None)
+        }
+    }
+
+    /// Returns the name of the `@listSize` directive in the given schema, accounting for import aliases or specification name
+    /// prefixes such as `@federation__listSize`. This checks the linked cost specification, if there is one, and falls back
+    /// to the federation spec.
+    fn list_size_directive_name(
+        schema: &FederationSchema,
+    ) -> Result<Option<Name>, FederationError> {
+        if let Some(spec) = Self::for_federation_schema(schema) {
+            spec.directive_name_in_schema(schema, &LIST_SIZE_DIRECTIVE_NAME)
+        } else if let Ok(fed_spec) = get_federation_spec_definition_from_subgraph(schema) {
+            fed_spec.directive_name_in_schema(schema, &LIST_SIZE_DIRECTIVE_NAME)
+        } else {
+            Ok(None)
+        }
+    }
+
+    pub fn cost_directive_from_argument(
+        schema: &FederationSchema,
+        argument: &InputValueDefinition,
+        ty: &ExtendedType,
+    ) -> Result<Option<CostDirective>, FederationError> {
+        let directive_name = Self::cost_directive_name(schema)?;
+        if let Some(name) = directive_name.as_ref() {
+            Ok(CostDirective::from_directives(name, &argument.directives)
+                .or(CostDirective::from_schema_directives(name, ty.directives())))
+        } else {
+            Ok(None)
+        }
+    }
+
+    pub fn cost_directive_from_field(
+        schema: &FederationSchema,
+        field: &FieldDefinition,
+        ty: &ExtendedType,
+    ) -> Result<Option<CostDirective>, FederationError> {
+        let directive_name = Self::cost_directive_name(schema)?;
+        if let Some(name) = directive_name.as_ref() {
+            Ok(CostDirective::from_directives(name, &field.directives)
+                .or(CostDirective::from_schema_directives(name, ty.directives())))
+        } else {
+            Ok(None)
+        }
+    }
+
+    pub fn list_size_directive_from_field_definition(
+        schema: &FederationSchema,
+        field: &FieldDefinition,
+    ) -> Result<Option<ListSizeDirective>, FederationError> {
+        let directive_name = Self::list_size_directive_name(schema)?;
+        if let Some(name) = directive_name.as_ref() {
+            Ok(ListSizeDirective::from_field_definition(name, field))
+        } else {
+            Ok(None)
+        }
+    }
 }
 
 impl SpecDefinition for CostSpecDefinition {
@@ -185,3 +262,96 @@ lazy_static! {
         definitions
     };
 }
+
+pub struct CostDirective {
+    weight: i32,
+}
+
+impl CostDirective {
+    pub fn weight(&self) -> f64 {
+        self.weight as f64
+    }
+
+    fn from_directives(directive_name: &Name, directives: &DirectiveList) -> Option<Self> {
+        directives
+            .get(directive_name)?
+            .specified_argument_by_name(&COST_DIRECTIVE_WEIGHT_ARGUMENT_NAME)?
+            .to_i32()
+            .map(|weight| Self { weight })
+    }
+
+    fn from_schema_directives(
+        directive_name: &Name,
+        directives: &apollo_compiler::schema::DirectiveList,
+    ) -> Option<Self> {
+        directives
+            .get(directive_name)?
+            .specified_argument_by_name(&COST_DIRECTIVE_WEIGHT_ARGUMENT_NAME)?
+            .to_i32()
+            .map(|weight| Self { weight })
+    }
+}
+
+pub struct ListSizeDirective {
+    pub assumed_size: Option<i32>,
+    pub slicing_argument_names: Option<HashSet<String>>,
+    pub sized_fields: Option<HashSet<String>>,
+    pub require_one_slicing_argument: bool,
+}
+
+impl ListSizeDirective {
+    pub fn from_field_definition(
+        directive_name: &Name,
+        definition: &FieldDefinition,
+    ) -> Option<Self> {
+        let directive = definition.directives.get(directive_name)?;
+        let assumed_size = Self::assumed_size(directive);
+        let slicing_argument_names = Self::slicing_argument_names(directive);
+        let sized_fields = Self::sized_fields(directive);
+        let require_one_slicing_argument =
+            Self::require_one_slicing_argument(directive).unwrap_or(true);
+
+        Some(Self {
+            assumed_size,
+            slicing_argument_names,
+            sized_fields,
+            require_one_slicing_argument,
+        })
+    }
+
+    fn assumed_size(directive: &Directive) -> Option<i32> {
+        directive
+            .specified_argument_by_name(&LIST_SIZE_DIRECTIVE_ASSUMED_SIZE_ARGUMENT_NAME)?
+            .to_i32()
+    }
+
+    fn slicing_argument_names(directive: &Directive) -> Option<HashSet<String>> {
+        let names = directive
+            .specified_argument_by_name(&LIST_SIZE_DIRECTIVE_SLICING_ARGUMENTS_ARGUMENT_NAME)?
+            .as_list()?
+            .iter()
+            .flat_map(|arg| arg.as_str())
+            .map(String::from)
+            .collect();
+        Some(names)
+    }
+
+    fn sized_fields(directive: &Directive) -> Option<HashSet<String>> {
+        let fields = directive
+            .specified_argument_by_name(&LIST_SIZE_DIRECTIVE_SIZED_FIELDS_ARGUMENT_NAME)?
+            .as_list()?
+            .iter()
+            .flat_map(|arg| arg.as_str())
+            .map(String::from)
+            .collect();
+        Some(fields)
+    }
+
+    fn require_one_slicing_argument(directive: &Directive) -> Option<bool> {
+        directive
+            .specified_argument_by_name(
+                &LIST_SIZE_DIRECTIVE_REQUIRE_ONE_SLICING_ARGUMENT_ARGUMENT_NAME,
+            )?
+            .to_bool()
+    }
+}

apollo-federation/src/link/federation_spec_definition.rs

@@ -14,8 +14,6 @@ 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;
 use crate::link::spec::Identity;
 use crate::link::spec::Url;
 use crate::link::spec::Version;
@@ -538,17 +536,6 @@ impl FederationSpecDefinition {
             arguments,
         })
     }
-
-    pub(crate) fn get_cost_spec_definition(
-        &self,
-        schema: &FederationSchema,
-    ) -> Option<&'static CostSpecDefinition> {
-        schema
-            .metadata()
-            .and_then(|metadata| metadata.for_identity(&Identity::cost_identity()))
-            .and_then(|link| COST_VERSIONS.find(&link.url.version))
-            .or_else(|| COST_VERSIONS.find_for_federation_version(self.version()))
-    }
 }
 
 impl SpecDefinition for FederationSpecDefinition {

apollo-federation/src/link/mod.rs

@@ -23,7 +23,7 @@ use crate::link::spec::Url;
 
 pub(crate) mod argument;
 pub(crate) mod context_spec_definition;
-pub(crate) mod cost_spec_definition;
+pub mod cost_spec_definition;
 pub mod database;
 pub(crate) mod federation_spec_definition;
 pub(crate) mod graphql_definition;