diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index a5bfc21957..c384f1f14a 100644 --- a/NEXT_CHANGELOG.md +++ b/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. diff --git a/apollo-router/src/query_planner/caching_query_planner.rs b/apollo-router/src/query_planner/caching_query_planner.rs index 9d7199c3e6..d44eed3dae 100644 --- a/apollo-router/src/query_planner/caching_query_planner.rs +++ b/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::() { 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()); + } + } }