Use a single special-purpose walk in subscription validation

Author: goto-bus-stop

crates/apollo-compiler/src/validation/mod.rs

@@ -1,11 +1,6 @@
 //! Supporting APIs for [GraphQL validation](https://spec.graphql.org/October2021/#sec-Validation)
 //! and other kinds of errors.
 
-use crate::coordinate::SchemaCoordinate;
-#[cfg(doc)]
-use crate::ExecutableDocument;
-use crate::Schema;
-
 pub(crate) mod argument;
 pub(crate) mod diagnostics;
 pub(crate) mod directive;
@@ -26,6 +21,7 @@ pub(crate) mod variable;
 use crate::collections::HashMap;
 use crate::collections::HashSet;
 use crate::collections::IndexSet;
+use crate::coordinate::SchemaCoordinate;
 use crate::diagnostic::CliReport;
 use crate::diagnostic::Diagnostic;
 use crate::diagnostic::ToCliReport;
@@ -41,6 +37,7 @@ use crate::schema::BuildError as SchemaBuildError;
 use crate::schema::Implementers;
 use crate::Name;
 use crate::Node;
+use crate::Schema;
 use std::fmt;
 use std::sync::Arc;
 use std::sync::OnceLock;

crates/apollo-compiler/src/validation/operation.rs

@@ -1,52 +1,117 @@
+use crate::collections::HashSet;
 use crate::executable;
+use crate::validation::diagnostics::DiagnosticData;
 use crate::validation::DiagnosticList;
 use crate::validation::ExecutableValidationContext;
+use crate::validation::RecursionLimitError;
 use crate::ExecutableDocument;
+use crate::Name;
 use crate::Node;
 
+/// Iterate all selections in the selection set.
+///
+/// This includes fields, fragment spreads, and inline fragments. For fragments, both the spread
+/// and the fragment's nested selections are reported.
+///
+/// Does not recurse into nested fields.
+fn walk_selections<'doc>(
+    document: &'doc ExecutableDocument,
+    selections: &'doc executable::SelectionSet,
+    mut f: impl FnMut(&'doc executable::Selection),
+) -> Result<(), RecursionLimitError> {
+    fn walk_selections_inner<'doc>(
+        document: &'doc ExecutableDocument,
+        selection_set: &'doc executable::SelectionSet,
+        seen: &mut HashSet<&'doc Name>,
+        f: &mut dyn FnMut(&'doc executable::Selection),
+    ) -> Result<(), RecursionLimitError> {
+        for selection in &selection_set.selections {
+            f(selection);
+            match selection {
+                executable::Selection::Field(_) => {
+                    // Nothing to do
+                }
+                executable::Selection::FragmentSpread(fragment) => {
+                    let new = seen.insert(&fragment.fragment_name);
+                    if !new {
+                        continue;
+                    }
+
+                    // If the fragment doesn't exist, that error is reported elsewhere.
+                    if let Some(fragment_definition) =
+                        document.fragments.get(&fragment.fragment_name)
+                    {
+                        walk_selections_inner(
+                            document,
+                            &fragment_definition.selection_set,
+                            seen,
+                            f,
+                        )?;
+                    }
+                }
+                executable::Selection::InlineFragment(fragment) => {
+                    walk_selections_inner(document, &fragment.selection_set, seen, f)?;
+                }
+            }
+        }
+        Ok(())
+    }
+
+    walk_selections_inner(document, selections, &mut HashSet::default(), &mut f)
+}
+
 pub(crate) fn validate_subscription(
     document: &executable::ExecutableDocument,
     operation: &Node<executable::Operation>,
     diagnostics: &mut DiagnosticList,
 ) {
-    if operation.is_subscription() {
-        let fields = super::selection::expand_selections(
-            &document.fragments,
-            std::iter::once(&operation.selection_set),
-            Some((operation, diagnostics)),
-        );
+    if !operation.is_subscription() {
+        return;
+    }
 
-        if fields.len() > 1 {
-            diagnostics.push(
-                operation.location(),
-                executable::BuildError::SubscriptionUsesMultipleFields {
-                    name: operation.name.clone(),
-                    fields: fields
-                        .iter()
-                        .map(|field| field.field.name.clone())
-                        .collect(),
-                },
-            );
+    let mut field_names = vec![];
+
+    let walked = walk_selections(document, &operation.selection_set, |selection| {
+        if let executable::Selection::Field(field) = selection {
+            field_names.push(field.name.clone());
+            if matches!(field.name.as_str(), "__type" | "__schema" | "__typename") {
+                diagnostics.push(
+                    field.location(),
+                    executable::BuildError::SubscriptionUsesIntrospection {
+                        name: operation.name.clone(),
+                        field: field.name.clone(),
+                    },
+                );
+            }
         }
 
-        let has_introspection_fields = fields
+        if let Some(conditional_directive) = selection
+            .directives()
             .iter()
-            .find(|field| {
-                matches!(
-                    field.field.name.as_str(),
-                    "__type" | "__schema" | "__typename"
-                )
-            })
-            .map(|field| &field.field);
-        if let Some(field) = has_introspection_fields {
+            .find(|d| matches!(d.name.as_str(), "skip" | "include"))
+        {
             diagnostics.push(
-                field.location(),
-                executable::BuildError::SubscriptionUsesIntrospection {
+                conditional_directive.location(),
+                executable::BuildError::SubscriptionUsesConditionalSelection {
                     name: operation.name.clone(),
-                    field: field.name.clone(),
                 },
             );
         }
+    });
+
+    if walked.is_err() {
+        diagnostics.push(None, DiagnosticData::RecursionError {});
+        return;
+    }
+
+    if field_names.len() > 1 {
+        diagnostics.push(
+            operation.location(),
+            executable::BuildError::SubscriptionUsesMultipleFields {
+                name: operation.name.clone(),
+                fields: field_names,
+            },
+        );
     }
 }
 

crates/apollo-compiler/src/validation/selection.rs

@@ -64,31 +64,16 @@ impl<'a> FieldSelection<'a> {
 }
 
 /// Expand one or more selection sets to a list of all fields selected.
-pub(crate) fn expand_selections<'doc>(
+fn expand_selections<'doc>(
     fragments: &'doc IndexMap<Name, Node<executable::Fragment>>,
     selection_sets: impl Iterator<Item = &'doc executable::SelectionSet>,
-    mut for_subscription_top_level: Option<(&executable::Operation, &mut DiagnosticList)>,
 ) -> Vec<FieldSelection<'doc>> {
     let mut selections = vec![];
     let mut queue: VecDeque<&executable::SelectionSet> = selection_sets.collect();
     let mut seen_fragments = HashSet::with_hasher(ahash::RandomState::new());
 
     while let Some(next_set) = queue.pop_front() {
         for selection in &next_set.selections {
-            if let Some((operation, diagnostics)) = &mut for_subscription_top_level {
-                if let Some(conditional_directive) = selection
-                    .directives()
-                    .iter()
-                    .find(|d| matches!(d.name.as_str(), "skip" | "include"))
-                {
-                    diagnostics.push(
-                        conditional_directive.location(),
-                        executable::BuildError::SubscriptionUsesConditionalSelection {
-                            name: operation.name.clone(),
-                        },
-                    );
-                }
-            }
             match selection {
                 executable::Selection::Field(field) => {
                     selections.push(FieldSelection::new(&next_set.ty, field))
@@ -535,11 +520,8 @@ impl<'alloc, 's, 'doc> FieldsInSetCanMerge<'alloc, 's, 'doc> {
         &self,
         selection_sets: impl Iterator<Item = &'doc executable::SelectionSet>,
     ) -> &'alloc [FieldSelection<'doc>] {
-        self.alloc.alloc(expand_selections(
-            &self.document.fragments,
-            selection_sets,
-            None,
-        ))
+        self.alloc
+            .alloc(expand_selections(&self.document.fragments, selection_sets))
     }
 
     pub(crate) fn validate_operation(

crates/apollo-compiler/src/validation/variable.rs

@@ -80,7 +80,14 @@ pub(crate) fn validate_variable_definitions(
     }
 }
 
-/// Named fragments are "deduplicated": only visited once even if spread multiple times
+/// Call a function for every selection that is reachable from the given selection set.
+///
+/// This includes fields, fragment spreads, and inline fragments. For fragments, both the spread
+/// and the fragment's nested selections are reported. For fields, nested selections are also
+/// reported.
+///
+/// Named fragments are "deduplicated": only visited once even if spread multiple times *in
+/// different locations*. This is only appropriate for certain kinds of validations, so reuser beware.
 fn walk_selections_with_deduped_fragments<'doc>(
     document: &'doc ExecutableDocument,
     selections: &'doc executable::SelectionSet,

crates/apollo-compiler/test_data/ok/0117_subscription_conditions_not_at_root.graphql

@@ -0,0 +1,21 @@
+subscription ConditionalSub($includeContent: Boolean = true, $small: Boolean = true) {
+    messages {
+        username
+        text @include(if: $includeContent)
+        avatar @skip(if: $small)
+    }
+}
+
+type Query {
+    hello: String
+}
+
+type Message {
+    username: String
+    text: String
+    avatar: String
+}
+
+type Subscription {
+    messages: Message
+}