persisted queries have client names

.changesets/feat_glasser_pq_client_name.md

@@ -0,0 +1,13 @@
+### Client name support for Persisted Query Lists ([PR #6198](https://github.com/apollographql/router/pull/6198))
+
+The persisted query manifest fetched from Uplink can now contain a `clientName` field in each operation. Two operations with the same `id` but different `clientName` are considered to be distinct operations (and may have distinct bodies).
+
+Router resolves the client name by taking the first of these which exists:
+- Reading the `apollo_persisted_queries::client_name` context key (which may be set by a `router_service` plugin)
+- Reading the HTTP header named by `telemetry.apollo.client_name_header` (which defaults to `apollographql-client-name`)
+
+If a client name can be resolved for a request, Router first tries to find a persisted query with the specified ID and the resolved client name.
+
+If there is no operation with that ID and client name, or if a client name cannot be resolved, Router tries to find a persisted query with the specified ID and no client name specified.  (This means that existing PQ lists that do not contain client names will continue to work.)
+
+By [@glasser](https://github.com/glasser) in https://github.com/apollographql/router/pull/6198

apollo-router/src/configuration/snapshots/apollo_router__configuration__tests__schema_generation.snap

@@ -1747,7 +1747,7 @@ expression: "&schema"
         },
         "client_name_header": {
           "default": "apollographql-client-name",
-          "description": "The name of the header to extract from requests when populating 'client nane' for traces and metrics in Apollo Studio.",
+          "description": "The name of the header to extract from requests when populating 'client name' for traces and metrics in Apollo Studio.",
           "nullable": true,
           "type": "string"
         },

apollo-router/src/plugins/telemetry/apollo.rs

@@ -67,7 +67,7 @@ pub(crate) struct Config {
     #[schemars(skip)]
     pub(crate) apollo_graph_ref: Option<String>,
 
-    /// The name of the header to extract from requests when populating 'client nane' for traces and metrics in Apollo Studio.
+    /// The name of the header to extract from requests when populating 'client name' for traces and metrics in Apollo Studio.
     #[schemars(with = "Option<String>", default = "client_name_header_default_str")]
     #[serde(deserialize_with = "deserialize_header_name")]
     pub(crate) client_name_header: HeaderName,

apollo-router/src/plugins/telemetry/mod.rs

@@ -171,7 +171,7 @@ pub(crate) mod tracing;
 pub(crate) mod utils;
 
 // Tracing consts
-const CLIENT_NAME: &str = "apollo_telemetry::client_name";
+pub(crate) const CLIENT_NAME: &str = "apollo_telemetry::client_name";
 const CLIENT_VERSION: &str = "apollo_telemetry::client_version";
 const SUBGRAPH_FTV1: &str = "apollo_telemetry::subgraph_ftv1";
 pub(crate) const STUDIO_EXCLUDE: &str = "apollo_telemetry::studio::exclude";

apollo-router/src/services/layers/persisted_queries/manifest_poller.rs

@@ -21,8 +21,18 @@ use crate::uplink::stream_from_uplink_transforming_new_response;
 use crate::uplink::UplinkConfig;
 use crate::Configuration;
 
+/// The full identifier for an operation in a PQ list consists of an operation
+/// ID and an optional client name.
+#[derive(Debug, Clone, Eq, Hash, PartialEq)]
+pub struct FullPersistedQueryOperationId {
+    /// The operation ID (usually a hash).
+    pub operation_id: String,
+    /// The client name associated with the operation; if None, can be any client.
+    pub client_name: Option<String>,
+}
+
 /// An in memory cache of persisted queries.
-pub(crate) type PersistedQueryManifest = HashMap<String, String>;
+pub type PersistedQueryManifest = HashMap<FullPersistedQueryOperationId, String>;
 
 /// How the router should respond to requests that are not resolved as the IDs
 /// of an operation in the manifest. (For the most part this means "requests
@@ -212,7 +222,7 @@ impl PersistedQueryManifestPoller {
             if manifest_files.is_empty() {
                 return Err("no local persisted query list files specified".into());
             }
-            let mut manifest: HashMap<String, String> = PersistedQueryManifest::new();
+            let mut manifest = PersistedQueryManifest::new();
 
             for local_pq_list in manifest_files {
                 tracing::info!(
@@ -250,7 +260,13 @@ impl PersistedQueryManifestPoller {
                 }
 
                 for operation in manifest_file.operations {
-                    manifest.insert(operation.id, operation.body);
+                    manifest.insert(
+                        FullPersistedQueryOperationId {
+                            operation_id: operation.id,
+                            client_name: operation.client_name,
+                        },
+                        operation.body,
+                    );
                 }
             }
 
@@ -343,15 +359,35 @@ impl PersistedQueryManifestPoller {
         }
     }
 
-    pub(crate) fn get_operation_body(&self, persisted_query_id: &str) -> Option<String> {
+    pub(crate) fn get_operation_body(
+        &self,
+        persisted_query_id: &str,
+        client_name: Option<String>,
+    ) -> Option<String> {
         let state = self
             .state
             .read()
             .expect("could not acquire read lock on persisted query manifest state");
-        state
+        if let Some(body) = state
             .persisted_query_manifest
-            .get(persisted_query_id)
+            .get(&FullPersistedQueryOperationId {
+                operation_id: persisted_query_id.to_string(),
+                client_name: client_name.clone(),
+            })
             .cloned()
+        {
+            Some(body)
+        } else if client_name.is_some() {
+            state
+                .persisted_query_manifest
+                .get(&FullPersistedQueryOperationId {
+                    operation_id: persisted_query_id.to_string(),
+                    client_name: None,
+                })
+                .cloned()
+        } else {
+            None
+        }
     }
 
     pub(crate) fn get_all_operations(&self) -> Vec<String> {
@@ -588,7 +624,13 @@ async fn add_chunk_to_operations(
         match fetch_chunk(http_client.clone(), chunk_url).await {
             Ok(chunk) => {
                 for operation in chunk.operations {
-                    operations.insert(operation.id, operation.body);
+                    operations.insert(
+                        FullPersistedQueryOperationId {
+                            operation_id: operation.id,
+                            client_name: operation.client_name,
+                        },
+                        operation.body,
+                    );
                 }
                 return Ok(());
             }
@@ -674,9 +716,11 @@ pub(crate) struct SignedUrlChunk {
 
 /// A single operation containing an ID and a body,
 #[derive(Debug, Clone, Deserialize, Serialize)]
+#[serde(rename_all = "camelCase")]
 pub(crate) struct Operation {
     pub(crate) id: String,
     pub(crate) body: String,
+    pub(crate) client_name: Option<String>,
 }
 
 #[cfg(test)]
@@ -701,7 +745,7 @@ mod tests {
         )
         .await
         .unwrap();
-        assert_eq!(manifest_manager.get_operation_body(&id), Some(body))
+        assert_eq!(manifest_manager.get_operation_body(&id, None), Some(body))
     }
 
     #[tokio::test(flavor = "multi_thread")]
@@ -734,18 +778,26 @@ mod tests {
         )
         .await
         .unwrap();
-        assert_eq!(manifest_manager.get_operation_body(&id), Some(body))
+        assert_eq!(manifest_manager.get_operation_body(&id, None), Some(body))
     }
 
     #[test]
     fn safelist_body_normalization() {
-        let safelist = FreeformGraphQLSafelist::new(&PersistedQueryManifest::from([(
-            "valid-syntax".to_string(),
-            "fragment A on T { a }    query SomeOp { ...A ...B }    fragment,,, B on U{b c  } # yeah"
-                .to_string(),
-        ), (
-            "invalid-syntax".to_string(),
-            "}}}".to_string()),
+        let safelist = FreeformGraphQLSafelist::new(&PersistedQueryManifest::from([
+            (
+                FullPersistedQueryOperationId {
+                    operation_id: "valid-syntax".to_string(),
+                    client_name: None,
+                },
+                "fragment A on T { a }    query SomeOp { ...A ...B }    fragment,,, B on U{b c  } # yeah".to_string(),
+            ),
+            (
+                FullPersistedQueryOperationId {
+                    operation_id: "invalid-syntax".to_string(),
+                    client_name: None,
+                },
+                "}}}".to_string(),
+            ),
         ]));
 
         let is_allowed = |body: &str| -> bool {
@@ -795,6 +847,6 @@ mod tests {
         )
         .await
         .unwrap();
-        assert_eq!(manifest_manager.get_operation_body(&id), Some(body))
+        assert_eq!(manifest_manager.get_operation_body(&id, None), Some(body))
     }
 }

apollo-router/src/services/layers/persisted_queries/mod.rs

@@ -8,17 +8,21 @@ use http::header::CACHE_CONTROL;
 use http::HeaderValue;
 use http::StatusCode;
 use id_extractor::PersistedQueryIdExtractor;
+pub use manifest_poller::FullPersistedQueryOperationId;
+pub use manifest_poller::PersistedQueryManifest;
 pub(crate) use manifest_poller::PersistedQueryManifestPoller;
 use tower::BoxError;
 
 use self::manifest_poller::FreeformGraphQLAction;
 use super::query_analysis::ParsedDocument;
 use crate::graphql::Error as GraphQLError;
+use crate::plugins::telemetry::CLIENT_NAME;
 use crate::services::SupergraphRequest;
 use crate::services::SupergraphResponse;
 use crate::Configuration;
 
 const DONT_CACHE_RESPONSE_VALUE: &str = "private, no-cache, must-revalidate";
+const PERSISTED_QUERIES_CLIENT_NAME_CONTEXT_KEY: &str = "apollo_persisted_queries::client_name";
 
 struct UsedQueryIdFromManifest;
 
@@ -110,9 +114,21 @@ impl PersistedQueryLayer {
         } else {
             // if there is no query, look up the persisted query in the manifest
             // and put the body on the `supergraph_request`
-            if let Some(persisted_query_body) =
-                manifest_poller.get_operation_body(persisted_query_id)
-            {
+            if let Some(persisted_query_body) = manifest_poller.get_operation_body(
+                persisted_query_id,
+                // Use the first one of these that exists:
+                // - The PQL-specific context name entry
+                //   `apollo_persisted_queries::client_name` (which can be set
+                //   by router_service plugins)
+                // - The same name used by telemetry (ie, the value of the
+                //   header named by `telemetry.apollo.client_name_header`,
+                //   which defaults to `apollographql-client-name` by default)
+                request
+                    .context
+                    .get(PERSISTED_QUERIES_CLIENT_NAME_CONTEXT_KEY)
+                    .unwrap_or_default()
+                    .or_else(|| request.context.get(CLIENT_NAME).unwrap_or_default()),
+            ) {
                 let body = request.supergraph_request.body_mut();
                 body.query = Some(persisted_query_body);
                 body.extensions.remove("persistedQuery");
@@ -389,6 +405,7 @@ mod tests {
     use std::collections::HashMap;
     use std::time::Duration;
 
+    use maplit::hashmap;
     use serde_json::json;
 
     use super::*;
@@ -400,6 +417,7 @@ mod tests {
     use crate::services::layers::query_analysis::QueryAnalysisLayer;
     use crate::spec::Schema;
     use crate::test_harness::mocks::persisted_queries::*;
+    use crate::Context;
 
     #[tokio::test(flavor = "multi_thread")]
     async fn disabled_pq_layer_has_no_poller() {
@@ -479,6 +497,84 @@ mod tests {
         assert_eq!(request.supergraph_request.body().query, Some(body));
     }
 
+    #[tokio::test(flavor = "multi_thread")]
+    async fn enabled_pq_layer_with_client_names() {
+        let (_mock_guard, uplink_config) = mock_pq_uplink(&hashmap! {
+            FullPersistedQueryOperationId {
+                operation_id: "both-plain-and-cliented".to_string(),
+                client_name: None,
+            } => "query { bpac_no_client: __typename }".to_string(),
+            FullPersistedQueryOperationId {
+                operation_id: "both-plain-and-cliented".to_string(),
+                client_name: Some("web".to_string()),
+            } => "query { bpac_web_client: __typename }".to_string(),
+            FullPersistedQueryOperationId {
+                operation_id: "only-cliented".to_string(),
+                client_name: Some("web".to_string()),
+            } => "query { oc_web_client: __typename }".to_string(),
+        })
+        .await;
+
+        let pq_layer = PersistedQueryLayer::new(
+            &Configuration::fake_builder()
+                .persisted_query(PersistedQueries::builder().enabled(true).build())
+                .uplink(uplink_config)
+                .build()
+                .unwrap(),
+        )
+        .await
+        .unwrap();
+
+        let map_to_query = |operation_id: &str, client_name: Option<&str>| -> Option<String> {
+            let context = Context::new();
+            if let Some(client_name) = client_name {
+                context
+                    .insert(
+                        PERSISTED_QUERIES_CLIENT_NAME_CONTEXT_KEY,
+                        client_name.to_string(),
+                    )
+                    .unwrap();
+            }
+
+            let incoming_request = SupergraphRequest::fake_builder()
+                .extension(
+                    "persistedQuery",
+                    json!({"version": 1, "sha256Hash": operation_id.to_string()}),
+                )
+                .context(context)
+                .build()
+                .unwrap();
+
+            pq_layer
+                .supergraph_request(incoming_request)
+                .ok()
+                .expect("pq layer returned response instead of putting the query on the request")
+                .supergraph_request
+                .body()
+                .query
+                .clone()
+        };
+
+        assert_eq!(
+            map_to_query("both-plain-and-cliented", None),
+            Some("query { bpac_no_client: __typename }".to_string())
+        );
+        assert_eq!(
+            map_to_query("both-plain-and-cliented", Some("not-web")),
+            Some("query { bpac_no_client: __typename }".to_string())
+        );
+        assert_eq!(
+            map_to_query("both-plain-and-cliented", Some("web")),
+            Some("query { bpac_web_client: __typename }".to_string())
+        );
+        assert_eq!(
+            map_to_query("only-cliented", Some("web")),
+            Some("query { oc_web_client: __typename }".to_string())
+        );
+        assert_eq!(map_to_query("only-cliented", None), None);
+        assert_eq!(map_to_query("only-cliented", Some("not-web")), None);
+    }
+
     #[tokio::test(flavor = "multi_thread")]
     async fn pq_layer_passes_on_to_apq_layer_when_id_not_found() {
         let (_id, _body, manifest) = fake_manifest();
@@ -690,11 +786,17 @@ mod tests {
     #[tokio::test(flavor = "multi_thread")]
     async fn pq_layer_freeform_graphql_with_safelist() {
         let manifest = HashMap::from([(
-            "valid-syntax".to_string(),
+            FullPersistedQueryOperationId {
+              operation_id: "valid-syntax".to_string(),
+              client_name: None,
+            },
             "fragment A on Query { me { id } }    query SomeOp { ...A ...B }    fragment,,, B on Query{me{name,username}  } # yeah"
                 .to_string(),
         ), (
-            "invalid-syntax".to_string(),
+            FullPersistedQueryOperationId {
+              operation_id: "invalid-syntax".to_string(),
+              client_name: None,
+            },
             "}}}".to_string()),
         ]);