Get client name for more places; tests

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,
@@ -176,7 +176,7 @@ fn otlp_endpoint_default() -> Url {
     Url::parse(OTLP_ENDPOINT_DEFAULT).expect("must be valid url")
 }
 
-const fn client_name_header_default_str() -> &'static str {
+pub(crate) const fn client_name_header_default_str() -> &'static str {
     "apollographql-client-name"
 }
 

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

@@ -21,14 +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(crate) struct FullPersistedQueryOperationId {
-    operation_id: String,
-    client_name: Option<String>,
+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<FullPersistedQueryOperationId, 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

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

@@ -8,18 +8,22 @@ 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::apollo::client_name_header_default_str;
 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;
 
@@ -113,7 +117,29 @@ impl PersistedQueryLayer {
             // and put the body on the `supergraph_request`
             if let Some(persisted_query_body) = manifest_poller.get_operation_body(
                 persisted_query_id,
-                request.context.get(CLIENT_NAME).unwrap_or_default(),
+                // 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 if telemetry is enabled
+                //   (ie, the value of the header named by
+                //   `telemetry.apollo.client_name_header`, which defaults to
+                //   `apollographql-client-name` by default)
+                // - The value in the `apollographql-client-name` header
+                //   (whether or not telemetry is enabled)
+                request
+                    .context
+                    .get(PERSISTED_QUERIES_CLIENT_NAME_CONTEXT_KEY)
+                    .unwrap_or_default()
+                    .or_else(|| request.context.get(CLIENT_NAME).unwrap_or_default())
+                    .or_else(|| {
+                        request
+                            .supergraph_request
+                            .headers()
+                            .get(client_name_header_default_str())
+                            .map(|hv| hv.to_str().unwrap_or_default())
+                            .map(str::to_string)
+                    }),
             ) {
                 let body = request.supergraph_request.body_mut();
                 body.query = Some(persisted_query_body);

apollo-router/src/test_harness/mocks/persisted_queries.rs

@@ -14,6 +14,8 @@ use wiremock::Mock;
 use wiremock::MockServer;
 use wiremock::ResponseTemplate;
 
+pub use crate::services::layers::persisted_queries::FullPersistedQueryOperationId;
+pub use crate::services::layers::persisted_queries::PersistedQueryManifest;
 use crate::uplink::Endpoints;
 use crate::uplink::UplinkConfig;
 
@@ -32,7 +34,7 @@ pub async fn mock_empty_pq_uplink() -> (UplinkMockGuard, UplinkConfig) {
 
 /// Mocks an uplink server with a persisted query list with a delay.
 pub async fn mock_pq_uplink_with_delay(
-    manifest: &HashMap<String, String>,
+    manifest: &PersistedQueryManifest,
     delay: Duration,
 ) -> (UplinkMockGuard, UplinkConfig) {
     let (guard, url) = mock_pq_uplink_one_endpoint(manifest, Some(delay)).await;
@@ -43,7 +45,7 @@ pub async fn mock_pq_uplink_with_delay(
 }
 
 /// Mocks an uplink server with a persisted query list containing operations passed to this function.
-pub async fn mock_pq_uplink(manifest: &HashMap<String, String>) -> (UplinkMockGuard, UplinkConfig) {
+pub async fn mock_pq_uplink(manifest: &PersistedQueryManifest) -> (UplinkMockGuard, UplinkConfig) {
     let (guard, url) = mock_pq_uplink_one_endpoint(manifest, None).await;
     (
         guard,
@@ -58,22 +60,29 @@ pub struct UplinkMockGuard {
 }
 
 #[derive(Deserialize, Serialize)]
+#[serde(rename_all = "camelCase")]
 struct Operation {
     id: String,
     body: String,
+    #[serde(skip_serializing_if = "Option::is_none", default)]
+    client_name: Option<String>,
 }
 
 /// Mocks an uplink server; returns a single Url rather than a full UplinkConfig, so you
 /// can combine it with another one to test failover.
 pub async fn mock_pq_uplink_one_endpoint(
-    manifest: &HashMap<String, String>,
+    manifest: &PersistedQueryManifest,
     delay: Option<Duration>,
 ) -> (UplinkMockGuard, Url) {
     let operations: Vec<Operation> = manifest
         // clone the manifest so the caller can still make assertions about it
         .clone()
         .drain()
-        .map(|(id, body)| Operation { id, body })
+        .map(|(full_id, body)| Operation {
+            id: full_id.operation_id,
+            body,
+            client_name: full_id.client_name,
+        })
         .collect();
 
     let mock_gcs_server = MockServer::start().await;