Added expose_query_plan to federation source

libs/common/src/graphql.rs

@@ -322,6 +322,16 @@ pub struct GraphQLResponse {
 }
 
 impl GraphQLResponse {
+  pub fn append_extensions(&mut self, extensions: Map<String, Value>) {
+    if let Some(existing_extensions) = &mut self.extensions {
+      if let Value::Object(existing_extensions_map) = existing_extensions {
+        existing_extensions_map.extend(extensions);
+      }
+    } else {
+      self.extensions = Some(Value::Object(extensions));
+    }
+  }
+
   pub fn new_error(error: &str) -> Self {
     GraphQLResponse {
       data: None,

libs/config/conductor.schema.json

@@ -346,6 +346,14 @@
         "supergraph"
       ],
       "properties": {
+        "expose_query_plan": {
+          "description": "Exposes the query plan as JSON under \"extensions\"",
+          "default": null,
+          "type": [
+            "boolean",
+            "null"
+          ]
+        },
         "supergraph": {
           "description": "The endpoint URL for the GraphQL source.",
           "$ref": "#/definitions/SupergraphSourceConfig"

libs/config/src/lib.rs

@@ -383,6 +383,9 @@ fn graphql_source_definition_example() -> JsonSchemaExample<SourceDefinition> {
 pub struct FederationSourceConfig {
   /// The endpoint URL for the GraphQL source.
   pub supergraph: SupergraphSourceConfig,
+  /// Exposes the query plan as JSON under "extensions"
+  #[serde(default)]
+  pub expose_query_plan: Option<bool>,
 }
 
 #[derive(Deserialize, Serialize, Debug, Clone, JsonSchema)]

libs/engine/src/source/federation_source.rs

@@ -176,8 +176,18 @@ impl SourceRuntime for FederationSourceRuntime {
       let operation = downstream_request.parsed_operation;
 
       match execute_federation(&self.supergraph, operation).await {
-        Ok(response_data) => {
-          let response = serde_json::from_str::<GraphQLResponse>(&response_data).unwrap();
+        Ok((response_data, query_plan)) => {
+          let mut response = serde_json::from_str::<GraphQLResponse>(&response_data).unwrap();
+
+          if self.config.expose_query_plan.is_some_and(|v| v) {
+            let mut ext = serde_json::Map::new();
+            ext.insert(
+              "queryPlan".to_string(),
+              serde_json::value::to_value(query_plan).unwrap(),
+            );
+
+            response.append_extensions(ext);
+          }
 
           Ok(response)
         }

libs/federation_query_planner/src/lib.rs

@@ -2,6 +2,7 @@ use std::ops::Index;
 
 use anyhow::{Ok, Result};
 use graphql_parser::query::Document;
+use query_planner::QueryPlan;
 use serde_json::json;
 use supergraph::Supergraph;
 
@@ -20,7 +21,7 @@ pub mod user_query;
 pub async fn execute_federation(
   supergraph: &Supergraph,
   parsed_user_query: Document<'static, String>,
-) -> Result<String> {
+) -> Result<(String, QueryPlan)> {
   // println!("parsed_user_query: {:#?}", user_query);
   let mut user_query = parse_user_query(parsed_user_query)?;
   let query_plan = plan_for_user_query(supergraph, &mut user_query)?;
@@ -31,7 +32,10 @@ pub async fn execute_federation(
 
   // println!("response: {:#?}", json!(response_vec).to_string());
 
-  Ok(json!(response_vec.index(0).index(0).1).to_string())
+  Ok((
+    json!(response_vec.index(0).index(0).1).to_string(),
+    query_plan,
+  ))
 }
 
 #[cfg(test)]

test_config/config.yaml

@@ -13,6 +13,7 @@ sources:
   - id: fed
     type: federation
     config:
+      expose_query_plan: true
       supergraph:
         file: ./supergraph.graphql