Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Ensure cost directives are picked up when not explicitly imported #6328

Merged
merged 24 commits into from
Dec 9, 2024
Merged
Changes from 1 commit
Commits
Show all changes
24 commits
Select commit Hold shift + click to select a range
574087c
Use SpecDefinition::directive_name_in_spec instead of a custom-built …
tninesling Nov 20, 2024
9545c53
If the cost spec is not present, check if cost and listsize names are…
tninesling Nov 20, 2024
4ccddea
Fix name lookups for cost directive transfer during subgraph extraction
tninesling Nov 20, 2024
16619c3
Undo some changes to apollo-federation public API
tninesling Nov 21, 2024
9224105
Remove unnecessary argument from position-based demand control direct…
tninesling Nov 21, 2024
3cdc462
Remove unnecessary self references from CostSpecDefinition to ensure …
tninesling Nov 21, 2024
a704085
Push details of directive name lookup into the cost spec so abstracti…
tninesling Nov 22, 2024
c4acdad
Fix linter errors
tninesling Nov 22, 2024
454c0fb
Merge branch 'dev' into tninesling/cost-with-updated-composition
tninesling Nov 22, 2024
2fa54dd
Stop cloning federation schema when checking directives
tninesling Nov 22, 2024
0e35b2f
Lint errors again
tninesling Nov 22, 2024
faa79ba
Hack to switch schema types.get() and type_field() to use ahash::Hash…
tninesling Nov 23, 2024
d9e0e99
Force disambiguation between default type_field and the ahash replace…
tninesling Dec 2, 2024
5764937
Revert "Force disambiguation between default type_field and the ahash…
tninesling Dec 2, 2024
a2843d9
Revert "Hack to switch schema types.get() and type_field() to use aha…
tninesling Dec 2, 2024
fac6098
Propagate federation errors up from cost spec
tninesling Dec 2, 2024
40e5382
Handle federation error propagation in demand control plugin
tninesling Dec 2, 2024
679f1e6
Add changeset
tninesling Dec 3, 2024
fdd035a
Merge branch 'dev' into tninesling/cost-with-updated-composition
tninesling Dec 3, 2024
b1d857c
Merge branch 'dev' into tninesling/cost-with-updated-composition
tninesling Dec 3, 2024
8127488
Merge branch 'dev' into tninesling/cost-with-updated-composition
tninesling Dec 4, 2024
aaf84a7
Use internal_error! for Federation errors
tninesling Dec 9, 2024
2de9e2e
Fix committed suggestion batch
tninesling Dec 9, 2024
4d39f57
Merge branch 'dev' into tninesling/cost-with-updated-composition
tninesling Dec 9, 2024
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Propagate federation errors up from cost spec
tninesling committed Dec 2, 2024
commit fac60989525e87b56c84384928116247ae04bc7b
179 changes: 96 additions & 83 deletions apollo-federation/src/link/cost_spec_definition.rs
Original file line number Diff line number Diff line change
@@ -13,6 +13,7 @@ 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;
@@ -122,8 +123,10 @@ impl CostSpecDefinition {
schema: &FederationSchema,
arguments: Vec<Node<Argument>>,
) -> Result<Directive, FederationError> {
// TODO: Handle no directive name
let name = Self::cost_directive_name(schema)?.expect("has name");
let name =
Self::cost_directive_name(schema)?.ok_or_else(|| SingleFederationError::Internal {
message: "The \"@cost\" directive is undefined in the target schema".to_string(),
})?;
tninesling marked this conversation as resolved.
Show resolved Hide resolved

Ok(Directive { name, arguments })
}
@@ -132,8 +135,12 @@ impl CostSpecDefinition {
schema: &FederationSchema,
arguments: Vec<Node<Argument>>,
) -> Result<Directive, FederationError> {
// TODO: Handle no directive name
let name = Self::list_size_directive_name(schema)?.expect("has name");
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(),
}
})?;
tninesling marked this conversation as resolved.
Show resolved Hide resolved

Ok(Directive { name, arguments })
}
@@ -159,27 +166,19 @@ impl CostSpecDefinition {
ScalarTypeDefinitionPosition
);

fn for_federation_schema(
schema: &FederationSchema,
) -> Result<Option<&'static Self>, FederationError> {
let cost_link = schema
.metadata()
.as_ref()
.and_then(|metadata| metadata.for_identity(&Identity::cost_identity()));
let cost_spec = cost_link.and_then(|link| COST_VERSIONS.find(&link.url.version));
Ok(cost_spec)
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(name) = Self::for_federation_schema(schema)?.and_then(|spec| {
if let Some(spec) = Self::for_federation_schema(schema) {
spec.directive_name_in_schema(schema, &COST_DIRECTIVE_NAME)
.ok()
.flatten()
}) {
Ok(Some(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 {
@@ -193,12 +192,8 @@ impl CostSpecDefinition {
fn list_size_directive_name(
schema: &FederationSchema,
) -> Result<Option<Name>, FederationError> {
if let Some(name) = Self::for_federation_schema(schema)?.and_then(|spec| {
if let Some(spec) = Self::for_federation_schema(schema) {
spec.directive_name_in_schema(schema, &LIST_SIZE_DIRECTIVE_NAME)
.ok()
.flatten()
}) {
Ok(Some(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 {
@@ -210,30 +205,40 @@ impl CostSpecDefinition {
schema: &FederationSchema,
argument: &InputValueDefinition,
ty: &ExtendedType,
) -> Option<CostDirective> {
let directive_name = Self::cost_directive_name(schema).ok().flatten()?;
CostDirective::from_directives(&directive_name, &argument.directives).or(
CostDirective::from_schema_directives(&directive_name, ty.directives()),
)
) -> 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,
) -> Option<CostDirective> {
let directive_name = Self::cost_directive_name(schema).ok().flatten()?;
CostDirective::from_directives(&directive_name, &field.directives).or(
CostDirective::from_schema_directives(&directive_name, ty.directives()),
)
) -> 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,
) -> Option<ListSizeDirective> {
let directive_name = Self::list_size_directive_name(schema).ok().flatten()?;
ListSizeDirective::from_field_definition(&directive_name, field)
) -> 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)
}
}
}

@@ -269,9 +274,9 @@ impl CostDirective {

fn from_directives(directive_name: &Name, directives: &DirectiveList) -> Option<Self> {
directives
.get(directive_name)
.and_then(|cost| cost.specified_argument_by_name(&COST_DIRECTIVE_WEIGHT_ARGUMENT_NAME))
.and_then(|weight| weight.to_i32())
.get(directive_name)?
.specified_argument_by_name(&COST_DIRECTIVE_WEIGHT_ARGUMENT_NAME)?
.to_i32()
.map(|weight| Self { weight })
}

@@ -280,9 +285,9 @@ impl CostDirective {
directives: &apollo_compiler::schema::DirectiveList,
) -> Option<Self> {
directives
.get(directive_name)
.and_then(|cost| cost.specified_argument_by_name(&COST_DIRECTIVE_WEIGHT_ARGUMENT_NAME))
.and_then(|weight| weight.to_i32())
.get(directive_name)?
.specified_argument_by_name(&COST_DIRECTIVE_WEIGHT_ARGUMENT_NAME)?
.to_i32()
.map(|weight| Self { weight })
}
}
@@ -299,46 +304,54 @@ impl ListSizeDirective {
directive_name: &Name,
definition: &FieldDefinition,
) -> Option<Self> {
let directive = definition.directives.get(directive_name);
if let Some(directive) = directive {
let assumed_size = directive
.specified_argument_by_name(&LIST_SIZE_DIRECTIVE_ASSUMED_SIZE_ARGUMENT_NAME)
.and_then(|arg| arg.to_i32());
let slicing_argument_names = directive
.specified_argument_by_name(&LIST_SIZE_DIRECTIVE_SLICING_ARGUMENTS_ARGUMENT_NAME)
.and_then(|arg| arg.as_list())
.map(|arg_list| {
arg_list
.iter()
.flat_map(|arg| arg.as_str())
.map(String::from)
.collect()
});
let sized_fields = directive
.specified_argument_by_name(&LIST_SIZE_DIRECTIVE_SIZED_FIELDS_ARGUMENT_NAME)
.and_then(|arg| arg.as_list())
.map(|arg_list| {
arg_list
.iter()
.flat_map(|arg| arg.as_str())
.map(String::from)
.collect()
});
let require_one_slicing_argument = directive
.specified_argument_by_name(
&LIST_SIZE_DIRECTIVE_REQUIRE_ONE_SLICING_ARGUMENT_ARGUMENT_NAME,
)
.and_then(|arg| arg.to_bool())
.unwrap_or(true);

Some(Self {
assumed_size,
slicing_argument_names,
sized_fields,
require_one_slicing_argument,
})
} else {
None
}
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()
}
}
Original file line number Diff line number Diff line change
@@ -59,21 +59,21 @@ impl DemandControlledSchema {
&fed_schema,
field_definition,
field_type,
) {
)? {
field_cost_directives.insert(field_name.clone(), cost_directive);
}

if let Some(list_size_directive) =
CostSpecDefinition::list_size_directive_from_field_definition(
&fed_schema,
field_definition,
)
)?
{
field_list_size_directives
.insert(field_name.clone(), list_size_directive);
}

// TODO: Need to handle renaming for @requires also
// TODO: We should also handle renaming for @requires
if let Some(requires_directive) = RequiresDirective::from_field_definition(
field_definition,
type_name,
@@ -98,21 +98,21 @@ impl DemandControlledSchema {
&fed_schema,
field_definition,
field_type,
) {
)? {
field_cost_directives.insert(field_name.clone(), cost_directive);
}

if let Some(list_size_directive) =
CostSpecDefinition::list_size_directive_from_field_definition(
&fed_schema,
field_definition,
)
)?
{
field_list_size_directives
.insert(field_name.clone(), list_size_directive);
}

// TODO: Need to handle renaming for @requires also
// TODO: We should also handle renaming for @requires
if let Some(requires_directive) = RequiresDirective::from_field_definition(
field_definition,
type_name,
@@ -172,7 +172,12 @@ impl DemandControlledSchema {
definition: &InputValueDefinition,
ty: &ExtendedType,
) -> Option<CostDirective> {
// For now, we ignore FederationError and return None because this should not block the whole scoring
// process at runtime. Later, this should be pushed into the constructor and propagate any federation
// errors encountered when parsing.
CostSpecDefinition::cost_directive_from_argument(&self.inner, definition, ty)
.ok()
.flatten()
}
}