Pass parameters explicitely in QueryPlannerRequest instead of context

apollographql · Oct 29, 2024 · 3d68bbb · 3d68bbb
1 parent 39c6c03
commit 3d68bbb
Show file tree

Hide file tree

Showing 5 changed files with 55 additions and 20 deletions.
diff --git a/apollo-router/src/plugins/demand_control/cost_calculator/static_cost.rs b/apollo-router/src/plugins/demand_control/cost_calculator/static_cost.rs
@@ -583,6 +583,7 @@ mod tests {
 
     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 +682,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,
+                ctx,
+                query,
+                CacheKeyMetadata::default(),
+            ))
             .await
             .unwrap();
         let query_plan = match planner_res.content.unwrap() {

diff --git a/apollo-router/src/query_planner/bridge_query_planner.rs b/apollo-router/src/query_planner/bridge_query_planner.rs
@@ -594,21 +594,23 @@ impl Service<QueryPlannerRequest> for BridgeQueryPlanner {
         let QueryPlannerRequest {
             query: original_query,
             operation_name,
+            document,
+            metadata,
             context,
         } = req;
 
-        let metadata = context
-            .extensions()
-            .with_lock(|lock| lock.get::<CacheKeyMetadata>().cloned().unwrap_or_default());
+        // 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; /*match context
+                                        .extensions()
+                                        .with_lock(|lock| lock.get::<ParsedDocument>().cloned())
+                                    {
+                                        None => return Err(QueryPlannerError::SpecError(SpecError::UnknownFileId)),
+                                        Some(d) => d,
+                                    };*/
 
             let api_schema = this.schema.api_schema();
             match add_defer_labels(api_schema, &doc.ast) {

diff --git a/apollo-router/src/query_planner/bridge_query_planner_pool.rs b/apollo-router/src/query_planner/bridge_query_planner_pool.rs
@@ -303,6 +303,7 @@ mod tests {
 
     use super::*;
     use crate::metrics::FutureMetricsExt;
+    use crate::plugins::authorization::CacheKeyMetadata;
     use crate::spec::Query;
     use crate::Context;
 
@@ -326,11 +327,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,
+                context,
+                doc,
+                CacheKeyMetadata::default(),
+            ))
+            .await
+            .unwrap();
 
             let metrics = crate::metrics::collect_metrics();
             let heap_used = metrics.find("apollo.router.v8.heap.used").unwrap();

diff --git a/apollo-router/src/query_planner/caching_query_planner.rs b/apollo-router/src/query_planner/caching_query_planner.rs
@@ -332,13 +332,15 @@ where
                 }
 
                 context.extensions().with_lock(|mut lock| {
-                    lock.insert::<ParsedDocument>(doc);
-                    lock.insert(caching_key.metadata)
+                    lock.insert::<ParsedDocument>(doc.clone());
+                    lock.insert(caching_key.metadata.clone())
                 });
 
                 let request = QueryPlannerRequest {
                     query,
                     operation_name,
+                    document: doc,
+                    metadata: caching_key.metadata,
                     context: context.clone(),
                 };
 
@@ -503,6 +505,7 @@ where
             } = request;
 
             let schema = self.schema.api_schema();
+            // FIXME: does this work with authorization?
             if let Ok(modified_query) = add_defer_labels(schema, &doc.ast) {
                 query = modified_query.to_string();
             }
@@ -511,6 +514,8 @@ where
                 .query(query)
                 .and_operation_name(operation_name)
                 .context(context)
+                .document(doc)
+                .metadata(caching_key.metadata)
                 .build();
 
             // some clients might timeout and cancel the request before query planning is finished,

diff --git a/apollo-router/src/services/query_planner.rs b/apollo-router/src/services/query_planner.rs
@@ -13,6 +13,8 @@ use crate::graphql;
 use crate::query_planner::QueryPlan;
 use crate::Context;
 
+use super::layers::query_analysis::ParsedDocument;
+
 assert_impl_all!(Request: Send);
 /// [`Context`] for the request.
 #[derive(Derivative)]
@@ -21,6 +23,8 @@ pub(crate) struct Request {
     pub(crate) query: String,
     pub(crate) operation_name: Option<String>,
     pub(crate) context: Context,
+    pub(crate) document: ParsedDocument,
+    pub(crate) metadata: crate::plugins::authorization::CacheKeyMetadata,
 }
 
 #[buildstructor::buildstructor]
@@ -29,11 +33,19 @@ impl Request {
     ///
     /// Required parameters are required in non-testing code to create a QueryPlannerRequest.
     #[builder]
-    pub(crate) fn new(query: String, operation_name: Option<String>, context: Context) -> Request {
+    pub(crate) fn new(
+        query: String,
+        operation_name: Option<String>,
+        context: Context,
+        document: ParsedDocument,
+        metadata: crate::plugins::authorization::CacheKeyMetadata,
+    ) -> Request {
         Self {
             query,
             operation_name,
             context,
+            document,
+            metadata,
         }
     }
 }