Skip to content

Commit

Permalink
fix: support fragments from fetch nodes in connectors (#6166)
Browse files Browse the repository at this point in the history
  • Loading branch information
@lennyburdette
lennyburdette authored Oct 17, 2024
1 parent 445534b commit 4f7931d
Show file tree
Hide file tree
Showing 3 changed files with 913 additions and 135 deletions.
199 changes: 163 additions & 36 deletions apollo-federation/src/sources/connect/json_selection/selection_set.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ use apollo_compiler::collections::IndexSet;
use apollo_compiler::executable::Field;
use apollo_compiler::executable::Selection;
use apollo_compiler::executable::SelectionSet;
use apollo_compiler::ExecutableDocument;
use apollo_compiler::Node;
use multimap::MultiMap;

Expand All @@ -40,21 +41,29 @@ use crate::sources::connect::SubSelection;

impl JSONSelection {
/// Apply a selection set to create a new [`JSONSelection`]
pub fn apply_selection_set(&self, selection_set: &SelectionSet) -> Self {
pub fn apply_selection_set(
&self,
document: &ExecutableDocument,
selection_set: &SelectionSet,
) -> Self {
match self {
Self::Named(sub) => Self::Named(sub.apply_selection_set(selection_set)),
Self::Path(path) => Self::Path(path.apply_selection_set(selection_set)),
Self::Named(sub) => Self::Named(sub.apply_selection_set(document, selection_set)),
Self::Path(path) => Self::Path(path.apply_selection_set(document, selection_set)),
}
}
}

impl SubSelection {
/// Apply a selection set to create a new [`SubSelection`]
pub fn apply_selection_set(&self, selection_set: &SelectionSet) -> Self {
pub fn apply_selection_set(
&self,
document: &ExecutableDocument,
selection_set: &SelectionSet,
) -> Self {
let mut new_selections = Vec::new();
let mut dropped_fields = IndexSet::default();
let mut referenced_fields = HashSet::new();
let field_map = map_fields_by_name(selection_set);
let field_map = map_fields_by_name(document, selection_set);

// When the operation contains __typename, it might be used to complete
// an entity reference (e.g. `__typename id`) for a subsequent fetch.
Expand Down Expand Up @@ -112,8 +121,9 @@ impl SubSelection {
Some(Alias::new(field_response_key))
},
name.clone(),
sub.as_ref()
.map(|sub| sub.apply_selection_set(&field.selection_set)),
sub.as_ref().map(|sub| {
sub.apply_selection_set(document, &field.selection_set)
}),
));
}
} else if self.star.is_some() {
Expand All @@ -130,7 +140,8 @@ impl SubSelection {
for field in fields {
new_selections.push(NamedSelection::Path(
Some(Alias::new(field.response_key().as_str())),
path_selection.apply_selection_set(&field.selection_set),
path_selection
.apply_selection_set(document, &field.selection_set),
));
}
}
Expand All @@ -141,7 +152,7 @@ impl SubSelection {
// the new PathWithSubSelection to the new_selections.
new_selections.push(NamedSelection::Path(
None,
path_selection.apply_selection_set(selection_set),
path_selection.apply_selection_set(document, selection_set),
));
}

Expand All @@ -159,7 +170,7 @@ impl SubSelection {
for field in fields {
new_selections.push(NamedSelection::Group(
Alias::new(field.response_key().as_str()),
sub.apply_selection_set(&field.selection_set),
sub.apply_selection_set(document, &field.selection_set),
));
}
}
Expand Down Expand Up @@ -194,37 +205,59 @@ impl SubSelection {

impl PathSelection {
/// Apply a selection set to create a new [`PathSelection`]
pub fn apply_selection_set(&self, selection_set: &SelectionSet) -> Self {
pub fn apply_selection_set(
&self,
document: &ExecutableDocument,
selection_set: &SelectionSet,
) -> Self {
Self {
path: WithRange::new(
self.path.apply_selection_set(selection_set),
self.path.apply_selection_set(document, selection_set),
self.path.range(),
),
}
}
}

impl PathList {
pub(crate) fn apply_selection_set(&self, selection_set: &SelectionSet) -> Self {
pub(crate) fn apply_selection_set(
&self,
document: &ExecutableDocument,
selection_set: &SelectionSet,
) -> Self {
match self {
Self::Var(name, path) => Self::Var(
name.clone(),
WithRange::new(path.apply_selection_set(selection_set), path.range()),
WithRange::new(
path.apply_selection_set(document, selection_set),
path.range(),
),
),
Self::Key(key, path) => Self::Key(
key.clone(),
WithRange::new(path.apply_selection_set(selection_set), path.range()),
WithRange::new(
path.apply_selection_set(document, selection_set),
path.range(),
),
),
Self::Expr(expr, path) => Self::Expr(
expr.clone(),
WithRange::new(path.apply_selection_set(selection_set), path.range()),
WithRange::new(
path.apply_selection_set(document, selection_set),
path.range(),
),
),
Self::Method(method_name, args, path) => Self::Method(
method_name.clone(),
args.clone(),
WithRange::new(path.apply_selection_set(selection_set), path.range()),
WithRange::new(
path.apply_selection_set(document, selection_set),
path.range(),
),
),
Self::Selection(sub) => Self::Selection(sub.apply_selection_set(selection_set)),
Self::Selection(sub) => {
Self::Selection(sub.apply_selection_set(document, selection_set))
}
Self::Empty => Self::Empty,
}
}
Expand All @@ -238,23 +271,32 @@ fn key_name(path_selection: &PathSelection) -> Option<&str> {
}
}

fn map_fields_by_name(set: &SelectionSet) -> MultiMap<String, &Node<Field>> {
fn map_fields_by_name<'a>(
document: &'a ExecutableDocument,
set: &'a SelectionSet,
) -> MultiMap<String, &'a Node<Field>> {
let mut map = MultiMap::new();
map_fields_by_name_impl(set, &mut map);
map_fields_by_name_impl(document, set, &mut map);
map
}

fn map_fields_by_name_impl<'a>(set: &'a SelectionSet, map: &mut MultiMap<String, &'a Node<Field>>) {
fn map_fields_by_name_impl<'a>(
document: &'a ExecutableDocument,
set: &'a SelectionSet,
map: &mut MultiMap<String, &'a Node<Field>>,
) {
for selection in &set.selections {
match selection {
Selection::Field(field) => {
map.insert(field.name.to_string(), field);
}
Selection::FragmentSpread(_) => {
// Ignore - these are always expanded into inline fragments
Selection::FragmentSpread(f) => {
if let Some(fragment) = f.fragment_def(document) {
map_fields_by_name_impl(document, &fragment.selection_set, map);
}
}
Selection::InlineFragment(fragment) => {
map_fields_by_name_impl(&fragment.selection_set, map);
map_fields_by_name_impl(document, &fragment.selection_set, map);
}
}
}
Expand All @@ -264,12 +306,14 @@ fn map_fields_by_name_impl<'a>(set: &'a SelectionSet, map: &mut MultiMap<String,
mod tests {
use apollo_compiler::executable::SelectionSet;
use apollo_compiler::validation::Valid;
use apollo_compiler::ExecutableDocument;
use apollo_compiler::Schema;
use pretty_assertions::assert_eq;

fn selection_set(schema: &Valid<Schema>, s: &str) -> SelectionSet {
apollo_compiler::ExecutableDocument::parse_and_validate(schema, s, "./")
.unwrap()
fn selection_set(schema: &Valid<Schema>, s: &str) -> (ExecutableDocument, SelectionSet) {
let document =
apollo_compiler::ExecutableDocument::parse_and_validate(schema, s, "./").unwrap();
let selection_set = document
.operations
.anonymous
.as_ref()
Expand All @@ -279,7 +323,8 @@ mod tests {
.next()
.unwrap()
.selection_set
.clone()
.clone();
(document.into_inner(), selection_set)
}

#[test]
Expand Down Expand Up @@ -323,12 +368,12 @@ mod tests {
)
.unwrap();

let selection_set = selection_set(
let (document, selection_set) = selection_set(
&schema,
"{ t { z: a, y: b, x: d, w: h v: k { u: l t: m } } }",
);

let transformed = json.apply_selection_set(&selection_set);
let transformed = json.apply_selection_set(&document, &selection_set);
assert_eq!(
transformed.to_string(),
r###"$.result {
Expand Down Expand Up @@ -400,12 +445,12 @@ mod tests {
)
.unwrap();

let selection_set = selection_set(
let (document, selection_set) = selection_set(
&schema,
"{ t { a b_alias c { e: e_alias h group { j } } path_to_f } }",
);

let transformed = json_selection.apply_selection_set(&selection_set);
let transformed = json_selection.apply_selection_set(&document, &selection_set);
assert_eq!(
transformed.to_string(),
r###"$.result {
Expand Down Expand Up @@ -511,9 +556,9 @@ mod tests {
)
.unwrap();

let selection_set = selection_set(&schema, "{ t { a { b { renamed } } } }");
let (document, selection_set) = selection_set(&schema, "{ t { a { b { renamed } } } }");

let transformed = json.apply_selection_set(&selection_set);
let transformed = json.apply_selection_set(&document, &selection_set);
assert_eq!(
transformed.to_string(),
r###"$.result {
Expand Down Expand Up @@ -579,10 +624,10 @@ mod tests {
)
.unwrap();

let selection_set =
let (document, selection_set) =
selection_set(&schema, "{ t { id __typename author { __typename id } } }");

let transformed = json.apply_selection_set(&selection_set);
let transformed = json.apply_selection_set(&document, &selection_set);
assert_eq!(
transformed.to_string(),
r###"$.result {
Expand All @@ -592,6 +637,88 @@ mod tests {
__typename: $->echo("A")
id: authorId
}
}"###
);
}

#[test]
fn test_fragments() {
let json = super::JSONSelection::parse(
r###"
reviews: result {
id
product: { upc: product_upc }
author: { id: author_id }
}
"###,
)
.unwrap()
.1;

let schema = Schema::parse_and_validate(
r###"
type Query {
_entities(representations: [_Any!]!): [_Entity]
}
scalar _Any
union _Entity = Product
type Product {
upc: String
reviews: [Review]
}
type Review {
id: ID
product: Product
author: User
}
type User {
id: ID
}
"###,
"./",
)
.unwrap();

let (document, selection_set) = selection_set(
&schema,
"query ($representations: [_Any!]!) {
_entities(representations: $representations) {
..._generated_onProduct1_0
}
}
fragment _generated_onProduct1_0 on Product {
reviews {
id
product {
__typename
upc
}
author {
__typename
id
}
}
}",
);

let transformed = json.apply_selection_set(&document, &selection_set);
assert_eq!(
transformed.to_string(),
r###"reviews: result {
id
product: {
__typename: $->echo("Product")
upc: product_upc
}
author: {
__typename: $->echo("User")
id: author_id
}
}"###
);
}
Expand Down
Loading

0 comments on commit 4f7931d

Please sign in to comment.