Pass parameters explicitely in QueryPlannerRequest instead of context

apollo-router/src/error.rs

@@ -1,4 +1,5 @@
 //! Router errors.
+use std::collections::HashMap;
 use std::sync::Arc;
 
 use apollo_compiler::validation::DiagnosticList;
@@ -406,6 +407,19 @@ impl IntoGraphQLErrors for QueryPlannerError {
     }
 }
 
+impl QueryPlannerError {
+    pub(crate) fn usage_reporting(&self) -> Option<UsageReporting> {
+        match self {
+            QueryPlannerError::PlanningErrors(pe) => Some(pe.usage_reporting.clone()),
+            QueryPlannerError::SpecError(e) => Some(UsageReporting {
+                stats_report_key: e.get_error_key().to_string(),
+                referenced_fields_by_type: HashMap::new(),
+            }),
+            _ => None,
+        }
+    }
+}
+
 #[derive(Clone, Debug, Error, Serialize, Deserialize)]
 /// Container for planner setup errors
 pub(crate) struct PlannerErrors(Arc<Vec<PlannerError>>);

apollo-router/src/plugins/demand_control/cost_calculator/static_cost.rs

@@ -578,11 +578,13 @@ mod tests {
     use ahash::HashMapExt;
     use apollo_federation::query_plan::query_planner::QueryPlanner;
     use bytes::Bytes;
+    use router_bridge::planner::PlanOptions;
     use test_log::test;
     use tower::Service;
 
     use super::*;
     use crate::introspection::IntrospectionCache;
+    use crate::plugins::authorization::CacheKeyMetadata;
     use crate::query_planner::BridgeQueryPlanner;
     use crate::services::layers::query_analysis::ParsedDocument;
     use crate::services::QueryPlannerContent;
@@ -681,10 +683,16 @@ mod tests {
 
         let ctx = Context::new();
         ctx.extensions()
-            .with_lock(|mut lock| lock.insert::<ParsedDocument>(query));
+            .with_lock(|mut lock| lock.insert::<ParsedDocument>(query.clone()));
 
         let planner_res = planner
-            .call(QueryPlannerRequest::new(query_str.to_string(), None, ctx))
+            .call(QueryPlannerRequest::new(
+                query_str.to_string(),
+                None,
+                query,
+                CacheKeyMetadata::default(),
+                PlanOptions::default(),
+            ))
             .await
             .unwrap();
         let query_plan = match planner_res.content.unwrap() {

apollo-router/src/query_planner/bridge_query_planner.rs

@@ -43,7 +43,6 @@ use crate::metrics::meter_provider;
 use crate::plugins::authorization::AuthorizationPlugin;
 use crate::plugins::authorization::CacheKeyMetadata;
 use crate::plugins::authorization::UnauthorizedPaths;
-use crate::plugins::progressive_override::LABELS_TO_OVERRIDE_KEY;
 use crate::plugins::telemetry::config::ApolloSignatureNormalizationAlgorithm;
 use crate::plugins::telemetry::config::Conf as TelemetryConfig;
 use crate::query_planner::convert::convert_root_query_plan_node;
@@ -594,21 +593,14 @@ impl Service<QueryPlannerRequest> for BridgeQueryPlanner {
         let QueryPlannerRequest {
             query: original_query,
             operation_name,
-            context,
+            document,
+            metadata,
+            plan_options,
         } = req;
 
-        let metadata = context
-            .extensions()
-            .with_lock(|lock| lock.get::<CacheKeyMetadata>().cloned().unwrap_or_default());
         let this = self.clone();
         let fut = async move {
-            let mut doc = match context
-                .extensions()
-                .with_lock(|lock| lock.get::<ParsedDocument>().cloned())
-            {
-                None => return Err(QueryPlannerError::SpecError(SpecError::UnknownFileId)),
-                Some(d) => d,
-            };
+            let mut doc = document;
 
             let api_schema = this.schema.api_schema();
             match add_defer_labels(api_schema, &doc.ast) {
@@ -635,19 +627,9 @@ impl Service<QueryPlannerRequest> for BridgeQueryPlanner {
                         operation_name.as_deref(),
                         Arc::new(QueryHash(hash)),
                     )?;
-                    context
-                        .extensions()
-                        .with_lock(|mut lock| lock.insert::<ParsedDocument>(doc.clone()));
                 }
             }
 
-            let plan_options = PlanOptions {
-                override_conditions: context
-                    .get(LABELS_TO_OVERRIDE_KEY)
-                    .unwrap_or_default()
-                    .unwrap_or_default(),
-            };
-
             let res = this
                 .get(
                     QueryKey {
@@ -664,27 +646,8 @@ impl Service<QueryPlannerRequest> for BridgeQueryPlanner {
             match res {
                 Ok(query_planner_content) => Ok(QueryPlannerResponse::builder()
                     .content(query_planner_content)
-                    .context(context)
                     .build()),
-                Err(e) => {
-                    match &e {
-                        QueryPlannerError::PlanningErrors(pe) => {
-                            context.extensions().with_lock(|mut lock| {
-                                lock.insert(Arc::new(pe.usage_reporting.clone()))
-                            });
-                        }
-                        QueryPlannerError::SpecError(e) => {
-                            context.extensions().with_lock(|mut lock| {
-                                lock.insert(Arc::new(UsageReporting {
-                                    stats_report_key: e.get_error_key().to_string(),
-                                    referenced_fields_by_type: HashMap::new(),
-                                }))
-                            });
-                        }
-                        _ => (),
-                    }
-                    Err(e)
-                }
+                Err(e) => Err(e),
             }
         };
 

apollo-router/src/query_planner/bridge_query_planner_pool.rs

@@ -300,9 +300,11 @@ impl tower::Service<QueryPlannerRequest> for BridgeQueryPlannerPool {
 
 mod tests {
     use opentelemetry_sdk::metrics::data::Gauge;
+    use router_bridge::planner::PlanOptions;
 
     use super::*;
     use crate::metrics::FutureMetricsExt;
+    use crate::plugins::authorization::CacheKeyMetadata;
     use crate::spec::Query;
     use crate::Context;
 
@@ -326,11 +328,19 @@ mod tests {
 
             let doc = Query::parse_document(&query, None, &schema, &config).unwrap();
             let context = Context::new();
-            context.extensions().with_lock(|mut lock| lock.insert(doc));
-
-            pool.call(QueryPlannerRequest::new(query, None, context))
-                .await
-                .unwrap();
+            context
+                .extensions()
+                .with_lock(|mut lock| lock.insert(doc.clone()));
+
+            pool.call(QueryPlannerRequest::new(
+                query,
+                None,
+                doc,
+                CacheKeyMetadata::default(),
+                PlanOptions::default(),
+            ))
+            .await
+            .unwrap();
 
             let metrics = crate::metrics::collect_metrics();
             let heap_used = metrics.find("apollo.router.v8.heap.used").unwrap();