feat(federation): query plan correctness checker (#6498)

- This is the first version that implements the completeness part.

(cherry picked from commit 1724f0f)

apollo-federation/Cargo.toml

@@ -15,6 +15,8 @@ autotests = false                                      # Integration tests are m
 # This logging is gated behind a feature to avoid any unnecessary (even if
 # small) runtime costs where this data will not be desired.
 snapshot_tracing = ["ron"]
+# `correctness` feature enables the `correctness` module.
+correctness = []
 
 [dependencies]
 apollo-compiler.workspace = true

apollo-federation/cli/Cargo.toml

@@ -5,7 +5,7 @@ edition = "2021"
 
 [dependencies]
 apollo-compiler.workspace = true
-apollo-federation = { path = ".." }
+apollo-federation = { path = "..", features = ["correctness"] }
 clap = { version = "4.5.1", features = ["derive"] }
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }
 

apollo-federation/cli/src/main.rs

@@ -6,8 +6,10 @@ use std::path::PathBuf;
 use std::process::ExitCode;
 
 use apollo_compiler::ExecutableDocument;
+use apollo_federation::correctness::CorrectnessError;
 use apollo_federation::error::FederationError;
 use apollo_federation::error::SingleFederationError;
+use apollo_federation::internal_error;
 use apollo_federation::query_graph;
 use apollo_federation::query_plan::query_planner::QueryPlanner;
 use apollo_federation::query_plan::query_planner::QueryPlannerConfig;
@@ -262,11 +264,31 @@ fn cmd_plan(
 
     let query_doc =
         ExecutableDocument::parse_and_validate(planner.api_schema().schema(), query, query_path)?;
-    print!(
-        "{}",
-        planner.build_query_plan(&query_doc, None, Default::default())?
+    let query_plan = planner.build_query_plan(&query_doc, None, Default::default())?;
+    println!("{query_plan}");
+
+    // Check the query plan
+    let subgraphs_by_name = supergraph
+        .extract_subgraphs()
+        .unwrap()
+        .into_iter()
+        .map(|(name, subgraph)| (name, subgraph.schema))
+        .collect();
+    let result = apollo_federation::correctness::check_plan(
+        planner.api_schema(),
+        &supergraph.schema,
+        &subgraphs_by_name,
+        &query_doc,
+        &query_plan,
     );
-    Ok(())
+    match result {
+        Ok(_) => Ok(()),
+        Err(CorrectnessError::FederationError(e)) => Err(e),
+        Err(CorrectnessError::ComparisonError(e)) => Err(internal_error!(
+            "Response shape from query plan does not match response shape from input operation:\n{}",
+            e.description()
+        )),
+    }
 }
 
 fn cmd_validate(file_paths: &[PathBuf]) -> Result<(), FederationError> {

apollo-federation/src/correctness/mod.rs

@@ -0,0 +1,83 @@
+pub mod query_plan_analysis;
+#[cfg(test)]
+pub mod query_plan_analysis_test;
+pub mod response_shape;
+pub mod response_shape_compare;
+#[cfg(test)]
+pub mod response_shape_test;
+mod subgraph_constraint;
+
+use std::sync::Arc;
+
+use apollo_compiler::collections::IndexMap;
+use apollo_compiler::validation::Valid;
+use apollo_compiler::ExecutableDocument;
+
+use crate::compat::coerce_executable_values;
+use crate::correctness::response_shape_compare::compare_response_shapes_with_constraint;
+use crate::correctness::response_shape_compare::ComparisonError;
+use crate::query_plan::QueryPlan;
+use crate::schema::ValidFederationSchema;
+use crate::FederationError;
+
+//==================================================================================================
+// Public API
+
+#[derive(derive_more::From)]
+pub enum CorrectnessError {
+    /// Correctness checker's own error
+    FederationError(FederationError),
+    /// Error in the input that is subject to comparison
+    ComparisonError(ComparisonError),
+}
+
+/// Check if `this`'s response shape is a subset of `other`'s response shape.
+pub fn compare_operations(
+    schema: &ValidFederationSchema,
+    this: &Valid<ExecutableDocument>,
+    other: &Valid<ExecutableDocument>,
+) -> Result<(), CorrectnessError> {
+    let this_rs = response_shape::compute_response_shape_for_operation(this, schema)?;
+    let other_rs = response_shape::compute_response_shape_for_operation(other, schema)?;
+    tracing::debug!(
+        "compare_operations:\nResponse shape (left): {this_rs}\nResponse shape (right): {other_rs}"
+    );
+    Ok(response_shape_compare::compare_response_shapes(
+        &this_rs, &other_rs,
+    )?)
+}
+
+/// Check the correctness of the query plan against the schema and input operation by comparing
+/// the response shape of the input operation and the response shape of the query plan.
+/// - The input operation's response shape is supposed to be a subset of the input operation's.
+pub fn check_plan(
+    api_schema: &ValidFederationSchema,
+    supergraph_schema: &ValidFederationSchema,
+    subgraphs_by_name: &IndexMap<Arc<str>, ValidFederationSchema>,
+    operation_doc: &Valid<ExecutableDocument>,
+    plan: &QueryPlan,
+) -> Result<(), CorrectnessError> {
+    // Coerce constant expressions in the input operation document since query planner does it for
+    // subgraph fetch operations. But, this may be unnecessary in the future (see ROUTER-816).
+    let mut operation_doc = operation_doc.clone().into_inner();
+    coerce_executable_values(api_schema.schema(), &mut operation_doc);
+    let operation_doc = operation_doc
+        .validate(api_schema.schema())
+        .map_err(FederationError::from)?;
+
+    let op_rs = response_shape::compute_response_shape_for_operation(&operation_doc, api_schema)?;
+    let root_type = response_shape::compute_the_root_type_condition_for_operation(&operation_doc)?;
+    let plan_rs = query_plan_analysis::interpret_query_plan(supergraph_schema, &root_type, plan)
+        .map_err(|e| {
+            ComparisonError::new(format!(
+                "Failed to compute the response shape from query plan:\n{e}"
+            ))
+        })?;
+    tracing::debug!(
+        "check_plan:\nOperation response shape: {op_rs}\nQuery plan response shape: {plan_rs}"
+    );
+
+    let path_constraint = subgraph_constraint::SubgraphConstraint::at_root(subgraphs_by_name);
+    compare_response_shapes_with_constraint(&path_constraint, &op_rs, &plan_rs)?;
+    Ok(())
+}