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

Assorted bug fixes in new QP, found by a newly-ported test #5157

Merged
merged 12 commits into from
May 14, 2024
1 change: 1 addition & 0 deletions apollo-federation/src/query_graph/build_query_graph.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2009,6 +2009,7 @@ struct FederatedQueryGraphBuilderSubgraphData {
interface_object_directive_definition_name: Name,
}

#[derive(Debug)]
struct QueryGraphEdgeData {
head: NodeIndex,
tail: NodeIndex,
Expand Down
21 changes: 10 additions & 11 deletions apollo-federation/src/query_graph/graph_path.rs
Original file line number Diff line number Diff line change
Expand Up @@ -794,10 +794,12 @@ impl std::fmt::Display for ClosedPath {

/// A list of the options generated during query planning for a specific "closed branch", which is a
/// full/closed path in a GraphQL operation (i.e. one that ends in a leaf field).
#[derive(Debug)]
pub(crate) struct ClosedBranch(pub(crate) Vec<Arc<ClosedPath>>);

/// A list of the options generated during query planning for a specific "open branch", which is a
/// partial/open path in a GraphQL operation (i.e. one that does not end in a leaf field).
#[derive(Debug)]
pub(crate) struct OpenBranch(pub(crate) Vec<SimultaneousPathsWithLazyIndirectPaths>);

impl<TTrigger, TEdge> GraphPath<TTrigger, TEdge>
Expand Down Expand Up @@ -1169,6 +1171,7 @@ where
return Ok(Box::new(
self.graph
.out_edges_with_federation_self_edges(self.tail)
.into_iter()
.map(|edge_ref| edge_ref.id()),
));
}
Expand All @@ -1194,6 +1197,7 @@ where
Ok(Box::new(
self.graph
.out_edges(self.tail)
.into_iter()
.map(|edge_ref| edge_ref.id()),
))
}
Expand Down Expand Up @@ -1988,14 +1992,11 @@ impl OpGraphPath {
else {
return Ok(path);
};
let typename_field = Field::new(FieldData {
schema: self.graph.schema_by_source(&tail_weight.source)?.clone(),
field_position: tail_type_pos.introspection_typename_field(),
alias: None,
arguments: Arc::new(vec![]),
directives: Arc::new(Default::default()),
sibling_typename: None,
});
let typename_field = Field::new_introspection_typename(
self.graph.schema_by_source(&tail_weight.source)?,
&tail_type_pos,
None,
);
let Some(edge) = self.graph.edge_for_field(path.tail, &typename_field) else {
return Err(FederationError::internal(
"Unexpectedly missing edge for __typename field",
Expand Down Expand Up @@ -3438,9 +3439,7 @@ impl ClosedBranch {
}

// Keep track of which options should be kept.
let mut keep_options = std::iter::repeat_with(|| true)
.take(self.0.len())
.collect::<Vec<_>>();
let mut keep_options = vec![true; self.0.len()];
for option_index in 0..(self.0.len()) {
if !keep_options[option_index] {
continue;
Expand Down
40 changes: 28 additions & 12 deletions apollo-federation/src/query_graph/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -458,26 +458,42 @@ impl QueryGraph {
pub(crate) fn out_edges_with_federation_self_edges(
&self,
node: NodeIndex,
) -> impl Iterator<Item = EdgeReference<QueryGraphEdge>> {
self.graph.edges_directed(node, Direction::Outgoing)
) -> Vec<EdgeReference<QueryGraphEdge>> {
Self::sorted_edges(self.graph.edges_directed(node, Direction::Outgoing))
}

/// The outward edges from the given node, minus self-key and self-root-type-resolution edges,
/// as they're rarely useful (currently only used by `@defer`).
pub(crate) fn out_edges(
&self,
node: NodeIndex,
) -> impl Iterator<Item = EdgeReference<QueryGraphEdge>> {
self.graph
.edges_directed(node, Direction::Outgoing)
.filter(|edge_ref| {
pub(crate) fn out_edges(&self, node: NodeIndex) -> Vec<EdgeReference<QueryGraphEdge>> {
Self::sorted_edges(self.graph.edges_directed(node, Direction::Outgoing).filter(
|edge_ref| {
!(edge_ref.source() == edge_ref.target()
&& matches!(
edge_ref.weight().transition,
QueryGraphEdgeTransition::KeyResolution
| QueryGraphEdgeTransition::RootTypeResolution { .. }
))
})
},
))
}

/// Edge iteration order is unspecified in petgraph, but appears to be
/// *reverse* insertion order in practice.
/// This can affect generated query plans, such as when two options have the same cost.
/// To match the JS code base, we want to iterate in insertion order.
///
/// Sorting by edge indices relies on documented behavior:
/// <https://docs.rs/petgraph/latest/petgraph/graph/struct.Graph.html#graph-indices>
///
/// As of this writing, edges of the query graph are removed
/// in `FederatedQueryGraphBuilder::update_edge_tail` which specifically preserves indices
/// by pairing with an insertion.
fn sorted_edges<'graph>(
edges: impl Iterator<Item = EdgeReference<'graph, QueryGraphEdge>>,
) -> Vec<EdgeReference<'graph, QueryGraphEdge>> {
let mut edges: Vec<_> = edges.collect();
edges.sort_by_key(|e| -> EdgeIndex { e.id() });
edges
}

pub(crate) fn is_self_key_or_root_edge(
Expand Down Expand Up @@ -538,7 +554,7 @@ impl QueryGraph {
}

pub(crate) fn edge_for_field(&self, node: NodeIndex, field: &Field) -> Option<EdgeIndex> {
let mut candidates = self.out_edges(node).filter_map(|edge_ref| {
let mut candidates = self.out_edges(node).into_iter().filter_map(|edge_ref| {
let edge_weight = edge_ref.weight();
let QueryGraphEdgeTransition::FieldCollection {
field_definition_position,
Expand Down Expand Up @@ -578,7 +594,7 @@ impl QueryGraph {
// No type condition means the type hasn't changed, meaning there is no edge to take.
return None;
};
let mut candidates = self.out_edges(node).filter_map(|edge_ref| {
let mut candidates = self.out_edges(node).into_iter().filter_map(|edge_ref| {
let edge_weight = edge_ref.weight();
let QueryGraphEdgeTransition::Downcast {
to_type_position, ..
Expand Down
24 changes: 23 additions & 1 deletion apollo-federation/src/query_graph/path_tree.rs
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ use crate::query_plan::operation::SelectionSet;
// Typescript doesn't have a native way of associating equality/hash functions with types, so they
// were passed around manually. This isn't the case with Rust, where we instead implement trigger
// equality via `PartialEq` and `Hash`.
#[derive(Debug, Clone)]
#[derive(Clone)]
pub(crate) struct PathTree<TTrigger, TEdge>
where
TTrigger: Eq + Hash,
Expand Down Expand Up @@ -361,6 +361,27 @@ fn merge_conditions(
}
}

impl<TTrigger: std::fmt::Debug, TEdge: std::fmt::Debug> std::fmt::Debug
for PathTree<TTrigger, TEdge>
where
TTrigger: Eq + Hash,
TEdge: Copy + Into<Option<EdgeIndex>>,
{
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
let Self {
graph: _, // skip
node,
local_selection_sets,
childs,
} = self;
f.debug_struct("PathTree")
.field("node", node)
.field("local_selection_sets", local_selection_sets)
.field("childs", childs)
.finish()
}
}

#[cfg(test)]
mod tests {
use std::sync::Arc;
Expand Down Expand Up @@ -418,6 +439,7 @@ mod tests {
// find the edge that matches `field_name`
let (edge_ref, field_def) = query_graph
.out_edges(curr_node_idx)
.into_iter()
.find_map(|e_ref| {
let edge = e_ref.weight();
match &edge.transition {
Expand Down
33 changes: 17 additions & 16 deletions apollo-federation/src/query_plan/display.rs
Original file line number Diff line number Diff line change
Expand Up @@ -57,17 +57,17 @@ impl SubscriptionNode {

state.write("Primary: {")?;
primary.write_indented(state)?;
state.write("}")?;
state.write("},")?;

if let Some(rest) = rest {
state.new_line()?;
state.write("Rest: {")?;
rest.write_indented(state)?;
state.write("}")?;
state.write("},")?;
}

state.dedent()?;
state.write("}")
state.write("},")
}
}

Expand All @@ -94,13 +94,14 @@ impl FetchNode {
if let Some(v) = requires.as_ref() {
if !v.is_empty() {
write_selections(state, v)?;
state.write(" => ")?;
state.write(" =>")?;
state.new_line()?;
}
}
write_operation(state, operation_document)?;

state.dedent()?;
state.write("}")
state.write("},")
}
}

Expand All @@ -111,7 +112,7 @@ impl SequenceNode {

write_indented_lines(state, nodes, |state, node| node.write_indented(state))?;

state.write("}")
state.write("},")
}
}

Expand All @@ -122,7 +123,7 @@ impl ParallelNode {

write_indented_lines(state, nodes, |state, node| node.write_indented(state))?;

state.write("}")
state.write("},")
}
}

Expand All @@ -143,7 +144,7 @@ impl FlattenNode {
node.write_indented(state)?;

state.dedent()?;
state.write("}")
state.write("},")
}
}

Expand All @@ -163,16 +164,16 @@ impl ConditionNode {
state.indent()?;
if_clause.write_indented(state)?;
state.dedent()?;
state.write("}")?;
state.write("},")?;

state.write("Else {")?;
state.indent()?;
else_clause.write_indented(state)?;
state.dedent()?;
state.write("}")?;
state.write("},")?;

state.dedent()?;
state.write("}")
state.write("},")
}

(Some(if_clause), None) => {
Expand All @@ -182,7 +183,7 @@ impl ConditionNode {
if_clause.write_indented(state)?;

state.dedent()?;
state.write("}")
state.write("},")
}

(None, Some(else_clause)) => {
Expand All @@ -192,7 +193,7 @@ impl ConditionNode {
else_clause.write_indented(state)?;

state.dedent()?;
state.write("}")
state.write("},")
}

// Shouldn’t happen?
Expand All @@ -217,7 +218,7 @@ impl DeferNode {
}

state.dedent()?;
state.write("}")
state.write("},")
}
}

Expand Down Expand Up @@ -248,7 +249,7 @@ impl PrimaryDeferBlock {

state.dedent()?;
}
state.write("}")
state.write("},")
}
}

Expand Down Expand Up @@ -296,7 +297,7 @@ impl DeferredDeferBlock {

state.dedent()?;
}
state.write("}")
state.write("},")
}
}

Expand Down
41 changes: 15 additions & 26 deletions apollo-federation/src/query_plan/fetch_dependency_graph.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2016,7 +2016,7 @@ fn compute_nodes_for_key_resolution<'a>(
// (and so no one subgraph has a type definition with all the proper fields,
// only the supergraph does).
let input_type = dependency_graph.type_for_fetch_inputs(source_type.type_name())?;
let mut input_selections = SelectionSet::empty(
let mut input_selections = SelectionSet::for_composite_type(
dependency_graph.supergraph_schema.clone(),
input_type.clone(),
);
Expand Down Expand Up @@ -2057,14 +2057,11 @@ fn compute_nodes_for_key_resolution<'a>(
// We also ensure to get the __typename of the current type in the "original" node.
let node =
FetchDependencyGraph::node_weight_mut(&mut dependency_graph.graph, stack_item.node_id)?;
let typename_field = Arc::new(OpPathElement::Field(Field::new(FieldData {
schema: dependency_graph.supergraph_schema.clone(),
field_position: source_type.introspection_typename_field(),
alias: None,
arguments: Default::default(),
directives: Default::default(),
sibling_typename: None,
})));
let typename_field = Arc::new(OpPathElement::Field(Field::new_introspection_typename(
&dependency_graph.supergraph_schema,
&source_type,
None,
)));
let typename_path = stack_item
.node_path
.path_in_node
Expand Down Expand Up @@ -2131,14 +2128,11 @@ fn compute_nodes_for_root_type_resolution<'a>(
let node =
FetchDependencyGraph::node_weight_mut(&mut dependency_graph.graph, stack_item.node_id)?;
if !stack_item.node_path.path_in_node.is_empty() {
let typename_field = Arc::new(OpPathElement::Field(Field::new(FieldData {
schema: dependency_graph.supergraph_schema.clone(),
field_position: source_type.introspection_typename_field().into(),
alias: None,
arguments: Default::default(),
directives: Default::default(),
sibling_typename: None,
})));
let typename_field = Arc::new(OpPathElement::Field(Field::new_introspection_typename(
&dependency_graph.supergraph_schema,
&source_type.into(),
None,
)));
let typename_path = stack_item
.node_path
.path_in_node
Expand Down Expand Up @@ -2239,16 +2233,11 @@ fn compute_nodes_for_op_path_element<'a>(
} else {
Some(name.clone())
};
let typename_field = Arc::new(OpPathElement::Field(Field::new(FieldData {
schema: dependency_graph.supergraph_schema.clone(),
field_position: operation
.parent_type_position()
.introspection_typename_field(),
let typename_field = Arc::new(OpPathElement::Field(Field::new_introspection_typename(
&dependency_graph.supergraph_schema,
&operation.parent_type_position(),
alias,
arguments: Default::default(),
directives: Default::default(),
sibling_typename: None,
})));
)));
let typename_path = stack_item
.node_path
.path_in_node
Expand Down
1 change: 1 addition & 0 deletions apollo-federation/src/query_plan/generate.rs
Original file line number Diff line number Diff line change
Expand Up @@ -107,6 +107,7 @@ where
Element: Clone,
// Uncomment to use `dbg!()`
// Element: std::fmt::Debug,
// Plan: std::fmt::Debug,
{
if to_add.is_empty() {
let cost = plan_builder.compute_plan_cost(&mut initial)?;
Expand Down
Loading
Loading