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

fix: cleanup defer handling within the QP #6470

Merged
merged 4 commits into from
Dec 17, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
15 changes: 15 additions & 0 deletions apollo-federation/src/operation/directive_list.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,12 +7,15 @@ use std::ops::Deref;
use std::sync::Arc;
use std::sync::OnceLock;

use apollo_compiler::collections::IndexSet;
use apollo_compiler::executable;
use apollo_compiler::Name;
use apollo_compiler::Node;
use serde::Serialize;

use super::sort_arguments;
use super::DEFER_DIRECTIVE_NAME;
use super::DEFER_LABEL_ARGUMENT_NAME;

/// Compare sorted input values, which means specifically establishing an order between the variants
/// of input values, and comparing values for the same variants accordingly.
Expand Down Expand Up @@ -332,6 +335,18 @@ impl DirectiveList {
inner.rehash();
Some(item)
}

/// Removes @defer directive from self if it has a matching label.
pub(crate) fn remove_defer(&mut self, defer_labels: &IndexSet<String>) {
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This was lifted out of the old DeferFilter. We don't need to pass schema anymore as at this stage all @defer applications will already have labels normalized.

let label = self
.get(&DEFER_DIRECTIVE_NAME)
.and_then(|directive| directive.specified_argument_by_name(&DEFER_LABEL_ARGUMENT_NAME))
.and_then(|arg| arg.as_str());

if label.is_some_and(|label| defer_labels.contains(label)) {
self.remove_one(&DEFER_DIRECTIVE_NAME);
}
}
}

/// Iterate over a [`DirectiveList`] in a consistent sort order.
Expand Down
113 changes: 37 additions & 76 deletions apollo-federation/src/operation/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2896,47 +2896,21 @@ impl DeferNormalizer {
}
}

#[derive(Debug, Clone, Copy)]
enum DeferFilter<'a> {
All,
Labels(&'a IndexSet<String>),
}

impl DeferFilter<'_> {
fn remove_defer(&self, directive_list: &mut DirectiveList, schema: &apollo_compiler::Schema) {
match self {
Self::All => {
directive_list.remove_one(&DEFER_DIRECTIVE_NAME);
}
Self::Labels(set) => {
let label = directive_list
.get(&DEFER_DIRECTIVE_NAME)
.and_then(|directive| {
directive
.argument_by_name(&DEFER_LABEL_ARGUMENT_NAME, schema)
.ok()
})
.and_then(|arg| arg.as_str());
if label.is_some_and(|label| set.contains(label)) {
directive_list.remove_one(&DEFER_DIRECTIVE_NAME);
}
}
}
}
}

impl Fragment {
/// Returns true if the fragment's selection set contains the @defer directive.
fn has_defer(&self) -> bool {
self.selection_set.has_defer()
}

fn without_defer(
/// Create a new fragment without @defer directive applications that have a matching label.
fn reduce_defer(
&self,
filter: DeferFilter<'_>,
defer_labels: &IndexSet<String>,
named_fragments: &NamedFragments,
) -> Result<Self, FederationError> {
let selection_set = self.selection_set.without_defer(filter, named_fragments)?;
let selection_set = self
.selection_set
.reduce_defer(defer_labels, named_fragments)?;
Ok(Fragment {
schema: self.schema.clone(),
name: self.name.clone(),
Expand All @@ -2948,8 +2922,8 @@ impl Fragment {
}

impl NamedFragments {
/// Creates new fragment definitions with the @defer directive removed.
fn without_defer(&self, filter: DeferFilter<'_>) -> Result<Self, FederationError> {
/// Creates new fragment definitions by removing all @defer directives that had a matching label.
fn reduce_defer(&self, defer_labels: &IndexSet<String>) -> Result<Self, FederationError> {
let mut new_fragments = NamedFragments {
fragments: Default::default(),
};
Expand All @@ -2960,7 +2934,7 @@ impl NamedFragments {
// if a fragment doesn't actually use @defer itself, to make sure that the `.selection_set`
// values on each selection are up to date.
for fragment in self.iter() {
let fragment = fragment.without_defer(filter, &new_fragments)?;
let fragment = fragment.reduce_defer(defer_labels, &new_fragments)?;
new_fragments.insert(fragment);
}
Ok(new_fragments)
Expand All @@ -2981,10 +2955,11 @@ impl FragmentSpread {
self.directives.has(&DEFER_DIRECTIVE_NAME)
}

fn without_defer(&self, filter: DeferFilter<'_>) -> Result<Self, FederationError> {
let mut without_defer = self.clone();
filter.remove_defer(&mut without_defer.directives, without_defer.schema.schema());
Ok(without_defer)
/// Create a new fragment spread without @defer directive applications that have a matching label.
fn reduce_defer(&self, defer_labels: &IndexSet<String>) -> Result<Self, FederationError> {
let mut reduce_defer = self.clone();
reduce_defer.directives.remove_defer(defer_labels);
Ok(reduce_defer)
}
}

Expand All @@ -3000,10 +2975,11 @@ impl InlineFragment {
self.directives.has(&DEFER_DIRECTIVE_NAME)
}

fn without_defer(&self, filter: DeferFilter<'_>) -> Result<Self, FederationError> {
let mut without_defer = self.clone();
filter.remove_defer(&mut without_defer.directives, without_defer.schema.schema());
Ok(without_defer)
/// Create a new inline fragment without @defer directive applications that have a matching label.
fn reduce_defer(&self, defer_labels: &IndexSet<String>) -> Result<Self, FederationError> {
let mut reduce_defer = self.clone();
reduce_defer.directives.remove_defer(defer_labels);
Ok(reduce_defer)
}
}

Expand Down Expand Up @@ -3098,9 +3074,10 @@ impl Selection {
}
}

fn without_defer(
/// Create a new selection without @defer directive applications that have a matching label.
fn reduce_defer(
&self,
filter: DeferFilter<'_>,
defer_labels: &IndexSet<String>,
named_fragments: &NamedFragments,
) -> Result<Self, FederationError> {
match self {
Expand All @@ -3115,17 +3092,19 @@ impl Selection {

Ok(field
.with_updated_selection_set(Some(
selection_set.without_defer(filter, named_fragments)?,
selection_set.reduce_defer(defer_labels, named_fragments)?,
))
.into())
}
Selection::FragmentSpread(frag) => {
let spread = frag.spread.without_defer(filter)?;
let spread = frag.spread.reduce_defer(defer_labels)?;
Ok(FragmentSpreadSelection::new(spread, named_fragments)?.into())
}
Selection::InlineFragment(frag) => {
let inline_fragment = frag.inline_fragment.without_defer(filter)?;
let selection_set = frag.selection_set.without_defer(filter, named_fragments)?;
let inline_fragment = frag.inline_fragment.reduce_defer(defer_labels)?;
let selection_set = frag
.selection_set
.reduce_defer(defer_labels, named_fragments)?;
Ok(InlineFragmentSelection::new(inline_fragment, selection_set).into())
}
}
Expand Down Expand Up @@ -3156,19 +3135,18 @@ impl Selection {
}

impl SelectionSet {
/// Create a new selection set without @defer directive applications.
fn without_defer(
/// Create a new selection set without @defer directive applications that have a matching label.
fn reduce_defer(
&self,
filter: DeferFilter<'_>,
defer_labels: &IndexSet<String>,
named_fragments: &NamedFragments,
) -> Result<Self, FederationError> {
let mut without_defer =
SelectionSet::empty(self.schema.clone(), self.type_position.clone());
let mut reduce_defer = SelectionSet::empty(self.schema.clone(), self.type_position.clone());
for selection in self.selections.values() {
without_defer
.add_local_selection(&selection.without_defer(filter, named_fragments)?)?;
reduce_defer
.add_local_selection(&selection.reduce_defer(defer_labels, named_fragments)?)?;
}
Ok(without_defer)
Ok(reduce_defer)
}

fn has_defer(&self) -> bool {
Expand Down Expand Up @@ -3203,31 +3181,14 @@ impl Operation {
.any(|f| f.has_defer())
}

/// Create a new operation without @defer directive applications.
pub(crate) fn without_defer(mut self) -> Result<Self, FederationError> {
if self.has_defer() {
let named_fragments = self.named_fragments.without_defer(DeferFilter::All)?;
self.selection_set = self
.selection_set
.without_defer(DeferFilter::All, &named_fragments)?;
self.named_fragments = named_fragments;
}
debug_assert!(!self.has_defer());
Ok(self)
}

/// Create a new operation without specific @defer(label:) directive applications.
pub(crate) fn reduce_defer(
mut self,
labels: &IndexSet<String>,
) -> Result<Self, FederationError> {
if self.has_defer() {
let named_fragments = self
.named_fragments
.without_defer(DeferFilter::Labels(labels))?;
self.selection_set = self
.selection_set
.without_defer(DeferFilter::Labels(labels), &named_fragments)?;
let named_fragments = self.named_fragments.reduce_defer(labels)?;
self.selection_set = self.selection_set.reduce_defer(labels, &named_fragments)?;
self.named_fragments = named_fragments;
}
Ok(self)
Expand Down
Loading
Loading