fix(query_planner): put the USAGE_REPORTING back in the context if the query plan has been cached

NEXT_CHANGELOG.md

@@ -452,6 +452,12 @@ There isn't much use for QueryPlanner plugins. Most of the logic done there can
 
 By [@o0Ignition0o](https://github.com/o0Ignition0o)
 
+### Include usage reporting data in the context even when the query plan has been cached ([#1559](https://github.com/apollographql/router/issues/1559))
+
+Include usage reporting data in the context even when the query plan has been cached when calling `CachingQueryPlanner`.
+
+By [@bnjjj](https://github.com/bnjjj) in https://github.com/apollographql/router/pull/1559
+
 ### 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/query_planner/caching_query_planner.rs

@@ -126,7 +126,23 @@ where
                     .map_err(|_| QueryPlannerError::UnhandledPlannerResult)?;
 
                 match res {
-                    Ok(content) => Ok(QueryPlannerResponse { content, context }),
+                    Ok(content) => {
+                        if let QueryPlannerContent::Plan { plan, .. } = &content {
+                            match (&plan.usage_reporting).serialize(Serializer) {
+                                Ok(v) => {
+                                    context.insert_json_value(USAGE_REPORTING, v);
+                                }
+                                Err(e) => {
+                                    tracing::error!(
+                                        "usage reporting was not serializable to context, {}",
+                                        e
+                                    );
+                                }
+                            }
+                        }
+
+                        Ok(QueryPlannerResponse { content, context })
+                    }
                     Err(error) => {
                         if let Some(error) = error.downcast_ref::<QueryPlannerError>() {
                             if let QueryPlannerError::PlanningErrors(pe) = &error {
@@ -171,12 +187,14 @@ where
 mod tests {
     use mockall::mock;
     use mockall::predicate::*;
+    use query_planner::QueryPlan;
     use router_bridge::planner::PlanErrors;
     use router_bridge::planner::UsageReporting;
     use test_log::test;
     use tower::Service;
 
     use super::*;
+    use crate::query_planner::QueryPlanOptions;
 
     mock! {
         #[derive(Debug)]
@@ -216,11 +234,8 @@ mod tests {
     async fn test_plan() {
         let mut delegate = MockMyQueryPlanner::new();
         delegate.expect_clone().returning(|| {
-            println!("cloning query planner");
             let mut planner = MockMyQueryPlanner::new();
             planner.expect_sync_call().times(0..2).returning(|_| {
-                println!("calling query planner");
-
                 Err(QueryPlannerError::from(PlanErrors {
                     errors: Default::default(),
                     usage_reporting: UsageReporting {
@@ -254,4 +269,56 @@ mod tests {
             .await
             .is_err());
     }
+
+    macro_rules! test_query_plan {
+        () => {
+            include_str!("testdata/query_plan.json")
+        };
+    }
+
+    #[test(tokio::test)]
+    async fn test_usage_reporting() {
+        let mut delegate = MockMyQueryPlanner::new();
+        delegate.expect_clone().returning(|| {
+            let mut planner = MockMyQueryPlanner::new();
+            planner.expect_sync_call().times(0..2).returning(|_| {
+                let query_plan: QueryPlan = QueryPlan {
+                    root: serde_json::from_str(test_query_plan!()).unwrap(),
+                    options: QueryPlanOptions::default(),
+                    usage_reporting: UsageReporting {
+                        stats_report_key: "this is a test report key".to_string(),
+                        referenced_fields_by_type: Default::default(),
+                    },
+                };
+                let qp_content = QueryPlannerContent::Plan {
+                    query: Arc::new(Query::default()),
+                    plan: Arc::new(query_plan),
+                };
+
+                Ok(QueryPlannerResponse::builder()
+                    .content(qp_content)
+                    .context(Context::new())
+                    .build())
+            });
+            planner
+        });
+
+        let mut planner = CachingQueryPlanner::new(delegate, 10).await;
+
+        for _ in 0..5 {
+            assert!(planner
+                .call(QueryPlannerRequest::new(
+                    "".into(),
+                    Some("".into()),
+                    Context::new()
+                ))
+                .await
+                .unwrap()
+                .context
+                .get::<_, UsageReporting>(USAGE_REPORTING)
+                .ok()
+                .flatten()
+                .is_some());
+        }
+    }
 }