expose query plan in extensions in GraphQL Response

NEXT_CHANGELOG.md

@@ -29,6 +29,12 @@ By [@USERNAME](https://github.com/USERNAME) in https://github.com/apollographql/
 
 ## 🚀 Features
 
+### Expose query plan in extensions for GraphQL response (experimental) ([PR #1470](https://github.com/apollographql/router/pull/1470))
+
+Expose query plan in extensions for GraphQL response. Only experimental for now, no documentation available.
+
+By [@bnjjj](https://github.com/bnjjj) in https://github.com/apollographql/router/pull/1470
+
 ### Add support of global rate limit and timeout. [PR #1347](https://github.com/apollographql/router/pull/1347)
 
 Additions to the traffic shaping plugin:

apollo-router/src/configuration/snapshots/apollo_router__configuration__tests__schema_generation.snap

@@ -1,6 +1,6 @@
 ---
 source: apollo-router/src/configuration/mod.rs
-assertion_line: 914
+assertion_line: 909
 expression: "&schema"
 ---
 {
@@ -283,6 +283,9 @@ expression: "&schema"
       "description": "Plugin configuration",
       "default": null,
       "properties": {
+        "experimental.expose_query_plan": {
+          "type": "boolean"
+        },
         "experimental.include_subgraph_errors": {
           "type": "object",
           "properties": {

apollo-router/src/context.rs

@@ -90,6 +90,14 @@ impl Context {
         self.entries.insert(key.into(), value)
     }
 
+    /// Get a json value from the context using the provided key.
+    pub fn get_json_value<K>(&self, key: K) -> Option<Value>
+    where
+        K: Into<String>,
+    {
+        self.entries.get(&key.into()).map(|v| v.value().clone())
+    }
+
     /// Upsert a value in the context using the provided key and resolving
     /// function.
     ///

apollo-router/src/plugins/expose_query_plan.rs

@@ -0,0 +1,248 @@
+use futures::future::ready;
+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::RouterRequest;
+use crate::services::RouterResponse;
+
+const EXPOSE_QUERY_PLAN_HEADER_NAME: &str = "Apollo-Expose-Query-Plan";
+const ENABLE_EXPOSE_QUERY_PLAN_ENV: &str = "APOLLO_EXPOSE_QUERY_PLAN";
+const QUERY_PLAN_CONTEXT_KEY: &str = "experimental::expose_query_plan.plan";
+const ENABLED_CONTEXT_KEY: &str = "experimental::expose_query_plan.enabled";
+
+#[derive(Debug, Clone)]
+struct ExposeQueryPlan {
+    enabled: bool,
+}
+
+#[async_trait::async_trait]
+impl Plugin for ExposeQueryPlan {
+    type Config = bool;
+
+    async fn new(init: PluginInit<Self::Config>) -> Result<Self, BoxError> {
+        Ok(ExposeQueryPlan {
+            enabled: init.config,
+        })
+    }
+
+    fn query_planning_service(
+        &self,
+        service: BoxService<QueryPlannerRequest, QueryPlannerResponse, BoxError>,
+    ) -> BoxService<QueryPlannerRequest, QueryPlannerResponse, BoxError> {
+        service
+            .map_response(move |res| {
+                if res
+                    .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();
+                    }
+                }
+
+                res
+            })
+            .boxed()
+    }
+
+    fn router_service(
+        &self,
+        service: BoxService<RouterRequest, RouterResponse, BoxError>,
+    ) -> BoxService<RouterRequest, RouterResponse, BoxError> {
+        let conf_enabled =
+            self.enabled || std::env::var(ENABLE_EXPOSE_QUERY_PLAN_ENV).as_deref() == Ok("true");
+        service
+            .map_future_with_context(move |req: &RouterRequest| {
+                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<RouterResponse, BoxError>  = f.await;
+                res = match res {
+                    Ok(mut res) => {
+                        if is_enabled {
+                            let (parts, stream) = http::Response::from(res.response).into_parts();
+                            let (mut first, rest) = stream.into_future().await;
+
+                            if let Some(first) = &mut first {
+                                if let Some(plan) =
+                                    res.context.get_json_value(QUERY_PLAN_CONTEXT_KEY)
+                                {
+                                    first
+                                        .extensions
+                                        .insert("apolloQueryPlan", json!({ "object": { "kind": "QueryPlan", "node": plan, "text": query } }));
+                                }
+                            }
+                            res.response = http::Response::from_parts(
+                                parts,
+                                once(ready(first.unwrap_or_default())).chain(rest).boxed(),
+                            )
+                            .into();
+                        }
+
+                        Ok(res)
+                    }
+                    Err(err) => Err(err),
+                };
+
+                res
+            })
+            .boxed()
+    }
+}
+
+register_plugin!("experimental", "expose_query_plan", ExposeQueryPlan);
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use once_cell::sync::Lazy;
+    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::*;
+    use crate::graphql::Response;
+    use crate::json_ext::Object;
+    use crate::plugin::test::MockSubgraph;
+    use crate::plugin::DynPlugin;
+    use crate::services::PluggableRouterServiceBuilder;
+    use crate::Schema;
+
+    static EXPECTED_RESPONSE_WITH_QUERY_PLAN: Lazy<Response> = Lazy::new(|| {
+        serde_json::from_str(include_str!(
+            "../../tests/fixtures/expected_response_with_queryplan.json"
+        ))
+        .unwrap()
+    });
+    static EXPECTED_RESPONSE_WITHOUT_QUERY_PLAN: Lazy<Response> = Lazy::new(|| {
+        serde_json::from_str(include_str!(
+            "../../tests/fixtures/expected_response_without_queryplan.json"
+        ))
+        .unwrap()
+    });
+
+    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<RouterRequest, RouterResponse, BoxError> {
+        let mut extensions = Object::new();
+        extensions.insert("test", Value::String(ByteString::from("value")));
+
+        let account_mocks = vec![
+            (
+                r#"{"query":"query TopProducts__accounts__3($representations:[_Any!]!){_entities(representations:$representations){...on User{name}}}","operationName":"TopProducts__accounts__3","variables":{"representations":[{"__typename":"User","id":"1"},{"__typename":"User","id":"2"},{"__typename":"User","id":"1"}]}}"#,
+                r#"{"data":{"_entities":[{"name":"Ada Lovelace"},{"name":"Alan Turing"},{"name":"Ada Lovelace"}]}}"#
+            )
+        ].into_iter().map(|(query, response)| (serde_json::from_str(query).unwrap(), serde_json::from_str(response).unwrap())).collect();
+        let account_service = MockSubgraph::new(account_mocks);
+
+        let review_mocks = vec![
+            (
+                r#"{"query":"query TopProducts__reviews__1($representations:[_Any!]!){_entities(representations:$representations){...on Product{reviews{id product{__typename upc}author{__typename id}}}}}","operationName":"TopProducts__reviews__1","variables":{"representations":[{"__typename":"Product","upc":"1"},{"__typename":"Product","upc":"2"}]}}"#,
+                r#"{"data":{"_entities":[{"reviews":[{"id":"1","product":{"__typename":"Product","upc":"1"},"author":{"__typename":"User","id":"1"}},{"id":"4","product":{"__typename":"Product","upc":"1"},"author":{"__typename":"User","id":"2"}}]},{"reviews":[{"id":"2","product":{"__typename":"Product","upc":"2"},"author":{"__typename":"User","id":"1"}}]}]}}"#
+            )
+            ].into_iter().map(|(query, response)| (serde_json::from_str(query).unwrap(), serde_json::from_str(response).unwrap())).collect();
+        let review_service = MockSubgraph::new(review_mocks);
+
+        let product_mocks = vec![
+            (
+                r#"{"query":"query TopProducts__products__0($first:Int){topProducts(first:$first){__typename upc name}}","operationName":"TopProducts__products__0","variables":{"first":2}}"#,
+                r#"{"data":{"topProducts":[{"__typename":"Product","upc":"1","name":"Table"},{"__typename":"Product","upc":"2","name":"Couch"}]}}"#
+            ),
+            (
+                r#"{"query":"query TopProducts__products__2($representations:[_Any!]!){_entities(representations:$representations){...on Product{name}}}","operationName":"TopProducts__products__2","variables":{"representations":[{"__typename":"Product","upc":"1"},{"__typename":"Product","upc":"1"},{"__typename":"Product","upc":"2"}]}}"#,
+                r#"{"data":{"_entities":[{"name":"Table"},{"name":"Table"},{"name":"Couch"}]}}"#
+            )
+            ].into_iter().map(|(query, response)| (serde_json::from_str(query).unwrap(), serde_json::from_str(response).unwrap())).collect();
+
+        let product_service = MockSubgraph::new(product_mocks).with_extensions(extensions);
+
+        let schema =
+            include_str!("../../../apollo-router-benchmarks/benches/fixtures/supergraph.graphql");
+        let schema = Arc::new(Schema::parse(schema, &Default::default()).unwrap());
+
+        let builder = PluggableRouterServiceBuilder::new(schema.clone());
+        let builder = builder
+            .with_dyn_plugin("experimental.expose_query_plan".to_string(), plugin)
+            .with_subgraph_service("accounts", account_service.clone())
+            .with_subgraph_service("reviews", review_service.clone())
+            .with_subgraph_service("products", product_service.clone());
+
+        let router = builder.build().await.expect("should build").test_service();
+
+        router
+    }
+
+    async fn get_plugin(config: &jValue) -> Box<dyn DynPlugin> {
+        crate::plugin::plugins()
+            .get("experimental.expose_query_plan")
+            .expect("Plugin not found")
+            .create_instance_without_schema(config)
+            .await
+            .expect("Plugin not created")
+    }
+
+    async fn execute_router_test(
+        query: &str,
+        body: &Response,
+        mut router_service: BoxCloneService<RouterRequest, RouterResponse, BoxError>,
+    ) {
+        let request = RouterRequest::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
+            .ready()
+            .await
+            .unwrap()
+            .call(request)
+            .await
+            .unwrap()
+            .next_response()
+            .await
+            .unwrap();
+
+        assert_eq!(response, *body);
+    }
+
+    #[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;
+    }
+
+    #[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;
+    }
+}

apollo-router/src/plugins/mod.rs

@@ -3,6 +3,7 @@
 //! These plugins are compiled into the router and configured via YAML configuration.
 
 pub mod csrf;
+mod expose_query_plan;
 mod forbid_mutations;
 mod headers;
 mod include_subgraph_errors;

apollo-router/src/query_planner/mod.rs

@@ -11,6 +11,7 @@ use futures::prelude::*;
 use opentelemetry::trace::SpanKind;
 use router_bridge::planner::UsageReporting;
 use serde::Deserialize;
+use serde::Serialize;
 use tokio::sync::broadcast::Sender;
 use tokio_stream::wrappers::BroadcastStream;
 use tracing::Instrument;
@@ -64,7 +65,7 @@ impl QueryPlan {
 }
 
 /// Query plans are composed of a set of nodes.
-#[derive(Debug, PartialEq, Deserialize)]
+#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)]
 #[serde(rename_all = "PascalCase", tag = "kind")]
 pub(crate) enum PlanNode {
     /// These nodes must be executed in order.
@@ -725,6 +726,7 @@ pub(crate) mod fetch {
 
     use indexmap::IndexSet;
     use serde::Deserialize;
+    use serde::Serialize;
     use tower::ServiceExt;
     use tracing::instrument;
     use tracing::Instrument;
@@ -743,7 +745,7 @@ pub(crate) mod fetch {
     use crate::*;
 
     /// GraphQL operation type.
-    #[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, Deserialize)]
+    #[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, Deserialize, Serialize)]
     #[serde(rename_all = "camelCase")]
     pub enum OperationKind {
         Query,
@@ -768,7 +770,7 @@ pub(crate) mod fetch {
     }
 
     /// A fetch node.
-    #[derive(Debug, PartialEq, Deserialize)]
+    #[derive(Debug, Clone, PartialEq, Deserialize, Serialize)]
     #[serde(rename_all = "camelCase")]
     pub(crate) struct FetchNode {
         /// The name of the service or subgraph that the fetch is querying.
@@ -1084,7 +1086,7 @@ pub(crate) mod fetch {
 }
 
 /// A flatten node.
-#[derive(Debug, PartialEq, Deserialize)]
+#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)]
 #[serde(rename_all = "camelCase")]
 pub(crate) struct FlattenNode {
     /// The path when result should be merged.
@@ -1095,7 +1097,7 @@ pub(crate) struct FlattenNode {
 }
 
 /// A primary query for a Defer node, the non deferred part
-#[derive(Debug, PartialEq, Deserialize)]
+#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)]
 #[serde(rename_all = "camelCase")]
 pub(crate) struct Primary {
     /// Optional path, set if and only if the defer node is a
@@ -1113,7 +1115,7 @@ pub(crate) struct Primary {
 /// The "deferred" parts of the defer (note that it's an array). Each
 /// of those deferred elements will correspond to a different chunk of
 /// the response to the client (after the initial non-deferred one that is).
-#[derive(Debug, Clone, PartialEq, Deserialize)]
+#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)]
 #[serde(rename_all = "camelCase")]
 pub(crate) struct DeferredNode {
     /// References one or more fetch node(s) (by `id`) within
@@ -1144,7 +1146,7 @@ impl DeferredNode {
     }
 }
 /// A deferred node.
-#[derive(Debug, Clone, PartialEq, Deserialize)]
+#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)]
 #[serde(rename_all = "camelCase")]
 pub(crate) struct Depends {
     id: String,