fix: ExposeQueryPlan: Move the plugin code to execution_service (#1543)

fixes #1541, part of #1545

This commit moves the ExposeQueryPlan behavior to the execution_service, while we're progressively removing query_planner_service from the api.

NEXT_CHANGELOG.md

@@ -437,6 +437,12 @@ By [@SimonSapin](https://github.com/SimonSapin)
 
 ## 🐛 Fixes
 
+### Expose query plan: move the behavior to the execution_service ([#1541](https://github.com/apollographql/router/issues/1541))
+
+There isn't much use for QueryPlanner plugins. Most of the logic done there can be done in `execution_service`. Moreover users could get inconsistent plugin behavior because it depends on whether the QueryPlanner cache hits or not.
+
+By [@o0Ignition0o](https://github.com/o0Ignition0o)
+
 ### Accept SIGTERM as shutdown signal ([PR #1497](https://github.com/apollographql/router/pull/1497))
 
 This will make containers stop faster as they will not have to wait until a SIGKILL to stop the router.

apollo-router/src/plugins/expose_query_plan.rs

@@ -3,19 +3,15 @@ use futures::stream::once;
 use futures::StreamExt;
 use http::HeaderValue;
 use serde_json_bytes::json;
-use tower::util::BoxService;
 use tower::BoxError;
 use tower::ServiceExt as TowerServiceExt;
 
 use crate::layers::ServiceExt;
 use crate::plugin::Plugin;
 use crate::plugin::PluginInit;
 use crate::register_plugin;
-use crate::services::QueryPlannerContent;
-use crate::services::QueryPlannerRequest;
-use crate::services::QueryPlannerResponse;
-use crate::services::SupergraphRequest;
-use crate::services::SupergraphResponse;
+use crate::stages::execution;
+use crate::stages::supergraph;
 
 const EXPOSE_QUERY_PLAN_HEADER_NAME: &str = "Apollo-Expose-Query-Plan";
 const ENABLE_EXPOSE_QUERY_PLAN_ENV: &str = "APOLLO_EXPOSE_QUERY_PLAN";
@@ -38,45 +34,37 @@ impl Plugin for ExposeQueryPlan {
         })
     }
 
-    fn query_planner_service(
-        &self,
-        service: BoxService<QueryPlannerRequest, QueryPlannerResponse, BoxError>,
-    ) -> BoxService<QueryPlannerRequest, QueryPlannerResponse, BoxError> {
+    fn execution_service(&self, service: execution::BoxService) -> execution::BoxService {
         service
-            .map_response(move |res| {
-                if res
+            .map_request(move |req: execution::Request| {
+                if req
                     .context
                     .get::<_, bool>(ENABLED_CONTEXT_KEY)
                     .ok()
                     .flatten()
                     .is_some()
                 {
-                    if let QueryPlannerContent::Plan { plan, .. } = &res.content {
-                        res.context
-                            .insert(QUERY_PLAN_CONTEXT_KEY, plan.root.clone())
-                            .unwrap();
-                    }
+                    req.context
+                        .insert(QUERY_PLAN_CONTEXT_KEY, req.query_plan.root.clone())
+                        .unwrap();
                 }
 
-                res
+                req
             })
             .boxed()
     }
 
-    fn supergraph_service(
-        &self,
-        service: BoxService<SupergraphRequest, SupergraphResponse, BoxError>,
-    ) -> BoxService<SupergraphRequest, SupergraphResponse, BoxError> {
+    fn supergraph_service(&self, service: supergraph::BoxService) -> supergraph::BoxService {
         let conf_enabled = self.enabled;
         service
-            .map_future_with_context(move |req: &SupergraphRequest| {
+            .map_future_with_context(move |req: &supergraph::Request| {
                 let is_enabled = conf_enabled && req.originating_request.headers().get(EXPOSE_QUERY_PLAN_HEADER_NAME) == Some(&HeaderValue::from_static("true"));
                 if is_enabled {
                     req.context.insert(ENABLED_CONTEXT_KEY, true).unwrap();
                 }
                 (req.originating_request.body().query.clone(), is_enabled)
             }, move |(query, is_enabled): (Option<String>, bool), f| async move {
-                let mut res: Result<SupergraphResponse, BoxError>  = f.await;
+                let mut res: supergraph::Result = f.await;
                 res = match res {
                     Ok(mut res) => {
                         if is_enabled {
@@ -120,7 +108,6 @@ mod tests {
     use serde_json::Value as jValue;
     use serde_json_bytes::ByteString;
     use serde_json_bytes::Value;
-    use tower::util::BoxCloneService;
     use tower::Service;
 
     use super::*;
@@ -146,9 +133,7 @@ mod tests {
 
     static VALID_QUERY: &str = r#"query TopProducts($first: Int) { topProducts(first: $first) { upc name reviews { id product { name } author { id name } } } }"#;
 
-    async fn build_mock_router(
-        plugin: Box<dyn DynPlugin>,
-    ) -> BoxCloneService<SupergraphRequest, SupergraphResponse, BoxError> {
+    async fn build_mock_supergraph(plugin: Box<dyn DynPlugin>) -> supergraph::BoxCloneService {
         let mut extensions = Object::new();
         extensions.insert("test", Value::String(ByteString::from("value")));
 
@@ -206,19 +191,19 @@ mod tests {
             .expect("Plugin not created")
     }
 
-    async fn execute_router_test(
+    async fn execute_supergraph_test(
         query: &str,
         body: &Response,
-        mut router_service: BoxCloneService<SupergraphRequest, SupergraphResponse, BoxError>,
+        mut supergraph_service: supergraph::BoxCloneService,
     ) {
-        let request = SupergraphRequest::fake_builder()
+        let request = supergraph::Request::fake_builder()
             .query(query.to_string())
             .variable("first", 2usize)
             .header(EXPOSE_QUERY_PLAN_HEADER_NAME, "true")
             .build()
             .expect("expecting valid request");
 
-        let response = router_service
+        let response = supergraph_service
             .ready()
             .await
             .unwrap()
@@ -235,14 +220,26 @@ mod tests {
     #[tokio::test]
     async fn it_expose_query_plan() {
         let plugin = get_plugin(&serde_json::json!(true)).await;
-        let router = build_mock_router(plugin).await;
-        execute_router_test(VALID_QUERY, &*EXPECTED_RESPONSE_WITH_QUERY_PLAN, router).await;
+        let supergraph = build_mock_supergraph(plugin).await;
+        execute_supergraph_test(
+            VALID_QUERY,
+            &*EXPECTED_RESPONSE_WITH_QUERY_PLAN,
+            supergraph.clone(),
+        )
+        .await;
+        // let's try that again
+        execute_supergraph_test(VALID_QUERY, &*EXPECTED_RESPONSE_WITH_QUERY_PLAN, supergraph).await;
     }
 
     #[tokio::test]
     async fn it_doesnt_expose_query_plan() {
         let plugin = get_plugin(&serde_json::json!(false)).await;
-        let router = build_mock_router(plugin).await;
-        execute_router_test(VALID_QUERY, &*EXPECTED_RESPONSE_WITHOUT_QUERY_PLAN, router).await;
+        let supergraph = build_mock_supergraph(plugin).await;
+        execute_supergraph_test(
+            VALID_QUERY,
+            &*EXPECTED_RESPONSE_WITHOUT_QUERY_PLAN,
+            supergraph,
+        )
+        .await;
     }
 }