Extract @cost and @listSize to query planner's subgraph schemas

apollo-federation/src/link/cost_spec_definition.rs

@@ -0,0 +1,165 @@
+use std::collections::HashMap;
+
+use apollo_compiler::ast::Argument;
+use apollo_compiler::ast::Directive;
+use apollo_compiler::name;
+use apollo_compiler::schema::Component;
+use apollo_compiler::schema::EnumType;
+use apollo_compiler::Name;
+use apollo_compiler::Node;
+use lazy_static::lazy_static;
+
+use crate::error::FederationError;
+use crate::link::spec::Identity;
+use crate::link::spec::Url;
+use crate::link::spec::Version;
+use crate::link::spec_definition::SpecDefinition;
+use crate::link::spec_definition::SpecDefinitions;
+use crate::schema::position::EnumTypeDefinitionPosition;
+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");
+
+#[derive(Clone)]
+pub(crate) struct CostSpecDefinition {
+    url: Url,
+    minimum_federation_version: Option<Version>,
+}
+
+impl CostSpecDefinition {
+    pub(crate) fn new(version: Version, minimum_federation_version: Option<Version>) -> Self {
+        Self {
+            url: Url {
+                identity: Identity::cost_identity(),
+                version,
+            },
+            minimum_federation_version,
+        }
+    }
+
+    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);
+
+        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);
+
+        Ok(Directive { name, arguments })
+    }
+
+    pub(crate) fn propagate_demand_control_directives(
+        &self,
+        subgraph_schema: &FederationSchema,
+        source: &apollo_compiler::ast::DirectiveList,
+        dest: &mut apollo_compiler::ast::DirectiveList,
+        original_directive_names: &HashMap<Name, Name>,
+    ) -> Result<(), FederationError> {
+        let cost_directive_name = original_directive_names.get(&COST_DIRECTIVE_NAME_IN_SPEC);
+        if let Some(cost_directive) = source.get(
+            cost_directive_name
+                .unwrap_or(&COST_DIRECTIVE_NAME_IN_SPEC)
+                .as_str(),
+        ) {
+            dest.push(Node::new(self.cost_directive(
+                subgraph_schema,
+                cost_directive.arguments.clone(),
+            )?));
+        }
+
+        let list_size_directive_name =
+            original_directive_names.get(&LIST_SIZE_DIRECTIVE_NAME_IN_SPEC);
+        if let Some(list_size_directive) = source.get(
+            list_size_directive_name
+                .unwrap_or(&LIST_SIZE_DIRECTIVE_NAME_IN_SPEC)
+                .as_str(),
+        ) {
+            dest.push(Node::new(self.list_size_directive(
+                subgraph_schema,
+                list_size_directive.arguments.clone(),
+            )?));
+        }
+
+        Ok(())
+    }
+
+    pub(crate) fn propagate_demand_control_directives_for_enum(
+        &self,
+        subgraph_schema: &mut FederationSchema,
+        source: &Node<EnumType>,
+        dest: &EnumTypeDefinitionPosition,
+        original_directive_names: &HashMap<Name, Name>,
+    ) -> Result<(), FederationError> {
+        let cost_directive_name = original_directive_names.get(&COST_DIRECTIVE_NAME_IN_SPEC);
+        if let Some(cost_directive) = source.directives.get(
+            cost_directive_name
+                .unwrap_or(&COST_DIRECTIVE_NAME_IN_SPEC)
+                .as_str(),
+        ) {
+            dest.insert_directive(
+                subgraph_schema,
+                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);
+        if let Some(list_size_directive) = source.directives.get(
+            list_size_directive_name
+                .unwrap_or(&LIST_SIZE_DIRECTIVE_NAME_IN_SPEC)
+                .as_str(),
+        ) {
+            dest.insert_directive(
+                subgraph_schema,
+                Component::from(
+                    self.list_size_directive(
+                        subgraph_schema,
+                        list_size_directive.arguments.clone(),
+                    )?,
+                ),
+            )?;
+        }
+
+        Ok(())
+    }
+}
+
+impl SpecDefinition for CostSpecDefinition {
+    fn url(&self) -> &Url {
+        &self.url
+    }
+
+    fn minimum_federation_version(&self) -> Option<&Version> {
+        self.minimum_federation_version.as_ref()
+    }
+}
+
+lazy_static! {
+    pub(crate) static ref COST_VERSIONS: SpecDefinitions<CostSpecDefinition> = {
+        let mut definitions = SpecDefinitions::new(Identity::cost_identity());
+        definitions.add(CostSpecDefinition::new(
+            Version { major: 0, minor: 1 },
+            Some(Version { major: 2, minor: 9 }),
+        ));
+        definitions
+    };
+}

apollo-federation/src/link/federation_spec_definition.rs

@@ -13,6 +13,8 @@ use crate::error::FederationError;
 use crate::error::SingleFederationError;
 use crate::link::argument::directive_optional_boolean_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;
@@ -387,6 +389,17 @@ 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 {
@@ -426,6 +439,22 @@ lazy_static! {
             major: 2,
             minor: 5,
         }));
+        definitions.add(FederationSpecDefinition::new(Version {
+            major: 2,
+            minor: 6,
+        }));
+        definitions.add(FederationSpecDefinition::new(Version {
+            major: 2,
+            minor: 7,
+        }));
+        definitions.add(FederationSpecDefinition::new(Version {
+            major: 2,
+            minor: 8,
+        }));
+        definitions.add(FederationSpecDefinition::new(Version {
+            major: 2,
+            minor: 9,
+        }));
         definitions
     };
 }

apollo-federation/src/link/mod.rs

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

apollo-federation/src/link/spec.rs

@@ -88,6 +88,13 @@ impl Identity {
             name: name!("inaccessible"),
         }
     }
+
+    pub fn cost_identity() -> Identity {
+        Identity {
+            domain: APOLLO_SPEC_DOMAIN.to_string(),
+            name: name!("cost"),
+        }
+    }
 }
 
 /// The version of a `@link` specification, in the form of a major and minor version numbers.

apollo-federation/src/link/spec_definition.rs

@@ -182,6 +182,17 @@ impl<T: SpecDefinition> SpecDefinitions<T> {
         self.definitions.get(requested)
     }
 
+    pub(crate) fn find_for_federation_version(&self, federation_version: &Version) -> Option<&T> {
+        for definition in self.definitions.values() {
+            if let Some(minimum_federation_version) = definition.minimum_federation_version() {
+                if minimum_federation_version >= federation_version {
+                    return Some(definition);
+                }
+            }
+        }
+        None
+    }
+
     pub(crate) fn versions(&self) -> Keys<Version, T> {
         self.definitions.keys()
     }