From a4d17bc038ddca6ed8c31f26a01012824142707d Mon Sep 17 00:00:00 2001 From: Benjamin Coenen <5719034+bnjjj@users.noreply.github.com> Date: Fri, 19 Aug 2022 17:55:10 +0200 Subject: [PATCH 1/2] fix(query_planner): put the USAGE_REPORTING back in the context if the query plan has been cached Signed-off-by: Benjamin Coenen <5719034+bnjjj@users.noreply.github.com> --- .../src/query_planner/caching_query_planner.rs | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/apollo-router/src/query_planner/caching_query_planner.rs b/apollo-router/src/query_planner/caching_query_planner.rs index 9d7199c3e6..a64dcf9e24 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 { From 8dc815070ec701949ed0080901dd1ea171691165 Mon Sep 17 00:00:00 2001 From: Benjamin Coenen <5719034+bnjjj@users.noreply.github.com> Date: Mon, 22 Aug 2022 10:07:01 +0200 Subject: [PATCH 2/2] add test + changelog Signed-off-by: Benjamin Coenen <5719034+bnjjj@users.noreply.github.com> --- NEXT_CHANGELOG.md | 6 ++ .../query_planner/caching_query_planner.rs | 57 ++++++++++++++++++- 2 files changed, 60 insertions(+), 3 deletions(-) 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 a64dcf9e24..d44eed3dae 100644 --- a/apollo-router/src/query_planner/caching_query_planner.rs +++ b/apollo-router/src/query_planner/caching_query_planner.rs @@ -187,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)] @@ -232,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 { @@ -270,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()); + } + } }