Update Apollo usage reporting to include persisted query metrics

.changesets/feat_persisted_query_reporting.md

@@ -0,0 +1,5 @@
+### Enables reporting of persisted query usage by PQ ID to Apollo ([PR #7166](https://github.com/apollographql/router/pull/7166))
+
+This change allows the router to report usage metrics by persisted query ID to Apollo, so that we can show usage stats for PQs.
+
+By [@bonnici](https://github.com/bonnici) in https://github.com/apollographql/router/pull/7166

apollo-router/src/apollo_studio_interop/mod.rs

@@ -28,9 +28,11 @@ use apollo_compiler::schema::ExtendedType;
 use apollo_compiler::validation::Valid;
 use serde::Deserialize;
 use serde::Serialize;
+use sha2::Digest;
 
 use crate::json_ext::Object;
 use crate::json_ext::Value as JsonValue;
+use crate::plugins::telemetry::apollo_exporter::proto::reports::QueryMetadata;
 use crate::plugins::telemetry::config::ApolloSignatureNormalizationAlgorithm;
 use crate::spec::Fragments;
 use crate::spec::Query;
@@ -161,17 +163,143 @@ impl AddAssign<ReferencedEnums> for AggregatedExtendedReferenceStats {
     }
 }
 
+#[derive(Deserialize, Serialize, Debug, PartialEq, Eq, Clone, Default)]
+#[serde(rename_all = "camelCase")]
+pub(crate) struct UsageReportingOperationDetails {
+    /// The operation name, or None if there is no operation name
+    operation_name: Option<String>,
+    /// The normalized operation signature, or None if there is no valid signature
+    operation_signature: Option<String>,
+    /// a list of all types and fields referenced in the query
+    #[serde(default)]
+    referenced_fields_by_type: HashMap<String, ReferencedFieldsForType>,
+}
+
 /// UsageReporting fields, that will be used to send stats to uplink/studio
 #[derive(Deserialize, Serialize, Debug, PartialEq, Eq, Clone)]
 #[serde(rename_all = "camelCase")]
-pub(crate) struct UsageReporting {
+pub(crate) enum UsageReporting {
+    Operation(UsageReportingOperationDetails),
+    PersistedQuery {
+        operation_details: UsageReportingOperationDetails,
+        persisted_query_id: String,
+    },
+    Error(String),
+}
+
+impl UsageReporting {
+    pub(crate) fn with_pq_id(&self, maybe_pq_id: Option<String>) -> UsageReporting {
+        match self {
+            UsageReporting::Operation(op_details)
+            | UsageReporting::PersistedQuery {
+                operation_details: op_details,
+                ..
+            } => match maybe_pq_id {
+                Some(pq_id) => UsageReporting::PersistedQuery {
+                    operation_details: op_details.clone(),
+                    persisted_query_id: pq_id,
+                },
+                None => UsageReporting::Operation(op_details.clone()),
+            },
+            // PQ ID has no effect on errors
+            UsageReporting::Error { .. } => self.clone(),
+        }
+    }
+
     /// The `stats_report_key` is a unique identifier derived from schema and query.
-    /// Metric data  sent to Studio must be aggregated
+    /// Metric data sent to Studio must be aggregated
     /// via grouped key of (`client_name`, `client_version`, `stats_report_key`).
-    pub(crate) stats_report_key: String,
-    /// a list of all types and fields referenced in the query
-    #[serde(default)]
-    pub(crate) referenced_fields_by_type: HashMap<String, ReferencedFieldsForType>,
+    /// For errors, the report key is of the form "## <error name>\n".
+    /// For operations requested by PQ, the report key is of the form "pq# <pq id>".
+    /// For operations not requested by PQ, the report key is of the form "# <op name>\n<op sig>".
+    /// Note that this combination of operation name and signature is sometimes referred to in code as
+    /// "operation signature" even though it also contains the operation name as the first line.
+    pub(crate) fn get_stats_report_key(&self) -> String {
+        match self {
+            UsageReporting::Operation(operation_details) => {
+                let stats_report_key_op_name = operation_details
+                    .operation_name
+                    .as_deref()
+                    .unwrap_or("-")
+                    .to_string();
+                format!(
+                    "# {}\n{}",
+                    stats_report_key_op_name,
+                    self.get_operation_signature(),
+                )
+            }
+            UsageReporting::PersistedQuery {
+                persisted_query_id, ..
+            } => {
+                format!("pq# {}", persisted_query_id)
+            }
+            UsageReporting::Error(error_key) => format!("## {}\n", error_key),
+        }
+    }
+
+    pub(crate) fn get_operation_signature(&self) -> String {
+        match self {
+            UsageReporting::Operation(operation_details)
+            | UsageReporting::PersistedQuery {
+                operation_details, ..
+            } => operation_details
+                .operation_signature
+                .clone()
+                .unwrap_or("".to_string()),
+            UsageReporting::Error { .. } => "".to_string(),
+        }
+    }
+
+    pub(crate) fn get_operation_id(&self) -> String {
+        let string_to_hash = match self {
+            UsageReporting::Operation { .. } | UsageReporting::PersistedQuery { .. } => {
+                self.get_stats_report_key()
+            }
+            UsageReporting::Error(error_key) => format!("# # {}\n", error_key),
+        };
+
+        let mut hasher = sha1::Sha1::new();
+        hasher.update(string_to_hash.as_bytes());
+        let result = hasher.finalize();
+        hex::encode(result)
+    }
+
+    pub(crate) fn get_operation_name(&self) -> String {
+        match self {
+            UsageReporting::Operation(operation_details)
+            | UsageReporting::PersistedQuery {
+                operation_details, ..
+            } => operation_details
+                .operation_name
+                .clone()
+                .unwrap_or("".to_string()),
+            UsageReporting::Error(error_key) => format!("# {}", error_key),
+        }
+    }
+
+    pub(crate) fn get_referenced_fields(&self) -> HashMap<String, ReferencedFieldsForType> {
+        match self {
+            UsageReporting::Operation(operation_details)
+            | UsageReporting::PersistedQuery {
+                operation_details, ..
+            } => operation_details.referenced_fields_by_type.clone(),
+            UsageReporting::Error { .. } => HashMap::default(),
+        }
+    }
+
+    pub(crate) fn get_query_metadata(&self) -> Option<QueryMetadata> {
+        match self {
+            UsageReporting::PersistedQuery {
+                persisted_query_id, ..
+            } => Some(QueryMetadata {
+                name: self.get_operation_name(),
+                signature: self.get_operation_signature(),
+                pq_id: persisted_query_id.clone(),
+            }),
+            // For now we only want to populate query metadata for PQ operations
+            _ => None,
+        }
+    }
 }
 
 /// A list of fields that will be resolved for a given type
@@ -186,9 +314,10 @@ pub(crate) struct ReferencedFieldsForType {
     pub(crate) is_interface: bool,
 }
 
-/// Generate a UsageReporting containing the stats_report_key (a normalized version of the operation signature)
-/// and referenced fields of an operation. The document used to generate the signature and for the references can be
-/// different to handle cases where the operation has been filtered, but we want to keep the same signature.
+/// Generate a UsageReporting containing the data required to generate a stats_report_key (either a normalized version of
+/// the operation signature or an error key or a PQ ID) and referenced fields of an operation. The document used to
+/// generate the signature and for the references can be different to handle cases where the operation has been filtered,
+/// but we want to keep the same signature.
 pub(crate) fn generate_usage_reporting(
     signature_doc: &ExecutableDocument,
     references_doc: &ExecutableDocument,
@@ -361,13 +490,22 @@ struct UsageGenerator<'a> {
 
 impl UsageGenerator<'_> {
     fn generate_usage_reporting(&mut self) -> UsageReporting {
-        UsageReporting {
-            stats_report_key: self.generate_stats_report_key(),
+        UsageReporting::Operation(UsageReportingOperationDetails {
+            operation_name: self.get_operation_name(),
+            operation_signature: self.generate_normalized_signature(),
             referenced_fields_by_type: self.generate_apollo_reporting_refs(),
-        }
+        })
     }
 
-    fn generate_stats_report_key(&mut self) -> String {
+    fn get_operation_name(&self) -> Option<String> {
+        self.signature_doc
+            .operations
+            .get(self.operation_name.as_deref())
+            .ok()
+            .and_then(|operation| operation.name.as_ref().map(|node| node.to_string()))
+    }
+
+    fn generate_normalized_signature(&mut self) -> Option<String> {
         self.fragments_map.clear();
 
         match self
@@ -376,10 +514,10 @@ impl UsageGenerator<'_> {
             .get(self.operation_name.as_deref())
             .ok()
         {
-            None => "".to_string(),
+            None => None,
             Some(operation) => {
                 self.extract_signature_fragments(&operation.selection_set);
-                self.format_operation_for_report(operation)
+                Some(self.format_operation_signature_for_report(operation))
             }
         }
     }
@@ -410,15 +548,10 @@ impl UsageGenerator<'_> {
         }
     }
 
-    fn format_operation_for_report(&self, operation: &Node<Operation>) -> String {
-        // The result in the name of the operation
-        let op_name = match &operation.name {
-            None => "-".into(),
-            Some(node) => node.to_string(),
-        };
-        let mut result = format!("# {}\n", op_name);
+    fn format_operation_signature_for_report(&self, operation: &Node<Operation>) -> String {
+        let mut result = String::new();
 
-        // Followed by a sorted list of fragments
+        // The signature starts with a sorted list of fragments
         let mut sorted_fragments: Vec<_> = self.fragments_map.iter().collect();
         sorted_fragments.sort_by_key(|&(k, _)| k);
 

apollo-router/src/apollo_studio_interop/tests.rs

@@ -5,7 +5,7 @@ use super::*;
 use crate::Configuration;
 
 fn assert_expected_signature(actual: &UsageReporting, expected_sig: &str) {
-    assert_eq!(actual.stats_report_key, expected_sig);
+    assert_eq!(actual.get_stats_report_key(), expected_sig);
 }
 
 macro_rules! assert_extended_references {
@@ -300,3 +300,110 @@ async fn test_enums_from_response_fragments() {
     let generated = enums_from_response(query_str, op_name, schema_str, response_str);
     assert_enums_from_response!(&generated);
 }
+
+#[test]
+fn apollo_operation_id_hash() {
+    let usage_reporting = UsageReporting::Operation(UsageReportingOperationDetails {
+        operation_name: Some("IgnitionMeQuery".to_string()),
+        operation_signature: Some("query IgnitionMeQuery{me{id}}".to_string()),
+        referenced_fields_by_type: HashMap::new(),
+    });
+
+    assert_eq!(
+        "d1554552698157b05c2a462827fb4367a4548ee5",
+        usage_reporting.get_operation_id()
+    );
+}
+
+// The Apollo operation ID hash for these errors is based on a slightly different string. E.g. instead of hashing
+// "## GraphQLValidationFailure\n" we should hash "# # GraphQLValidationFailure".
+#[test]
+fn apollo_error_operation_id_hash() {
+    assert_eq!(
+        "ea4f152696abedca148b016d72df48842b713697",
+        UsageReporting::Error("GraphQLValidationFailure".into()).get_operation_id()
+    );
+    assert_eq!(
+        "3f410834f13153f401ffe73f7e454aa500d10bf7",
+        UsageReporting::Error("GraphQLParseFailure".into()).get_operation_id()
+    );
+    assert_eq!(
+        "7486043da2085fed407d942508a572ef88dc8120",
+        UsageReporting::Error("GraphQLUnknownOperationName".into()).get_operation_id()
+    );
+}
+
+#[test]
+fn test_get_stats_report_key() {
+    let usage_reporting_for_errors = UsageReporting::Error("GraphQLParseFailure".into());
+    assert_eq!(
+        "## GraphQLParseFailure\n",
+        usage_reporting_for_errors.get_stats_report_key()
+    );
+
+    let usage_reporting_for_pq = UsageReporting::PersistedQuery {
+        operation_details: UsageReportingOperationDetails {
+            operation_name: Some("SomeQuery".into()),
+            operation_signature: Some("query SomeQuery{thing{id}}".into()),
+            referenced_fields_by_type: HashMap::new(),
+        },
+        persisted_query_id: "SomePqId".into(),
+    };
+    assert_eq!(
+        "pq# SomePqId",
+        usage_reporting_for_pq.get_stats_report_key()
+    );
+
+    let usage_reporting_for_named_operation =
+        UsageReporting::Operation(UsageReportingOperationDetails {
+            operation_name: Some("SomeQuery".into()),
+            operation_signature: Some("query SomeQuery{thing{id}}".into()),
+            referenced_fields_by_type: HashMap::new(),
+        });
+    assert_eq!(
+        "# SomeQuery\nquery SomeQuery{thing{id}}",
+        usage_reporting_for_named_operation.get_stats_report_key()
+    );
+
+    let usage_reporting_for_unnamed_operation =
+        UsageReporting::Operation(UsageReportingOperationDetails {
+            operation_name: None,
+            operation_signature: Some("query{thing{id}}".into()),
+            referenced_fields_by_type: HashMap::new(),
+        });
+    assert_eq!(
+        "# -\nquery{thing{id}}",
+        usage_reporting_for_unnamed_operation.get_stats_report_key()
+    );
+}
+
+#[test]
+fn test_get_operation_name() {
+    let usage_reporting_for_errors = UsageReporting::Error("GraphQLParseFailure".into());
+    assert_eq!(
+        "# GraphQLParseFailure",
+        usage_reporting_for_errors.get_operation_name()
+    );
+
+    let usage_reporting_for_named_operation =
+        UsageReporting::Operation(UsageReportingOperationDetails {
+            operation_name: Some("SomeQuery".into()),
+            operation_signature: Some("query SomeQuery{thing{id}}".into()),
+            referenced_fields_by_type: HashMap::new(),
+        });
+    assert_eq!(
+        "SomeQuery",
+        usage_reporting_for_named_operation.get_operation_name()
+    );
+
+    let usage_reporting_for_unnamed_operation =
+        UsageReporting::Operation(UsageReportingOperationDetails {
+            operation_name: None,
+            operation_signature: Some("query{thing{id}}".into()),
+            referenced_fields_by_type: HashMap::new(),
+        });
+    assert_eq!(
+        "",
+        usage_reporting_for_unnamed_operation.get_operation_name()
+    );
+}

apollo-router/src/error.rs

@@ -1,5 +1,4 @@
 //! Router errors.
-use std::collections::HashMap;
 use std::ops::Deref;
 use std::sync::Arc;
 
@@ -271,7 +270,7 @@ pub(crate) enum QueryPlannerError {
     JoinError(String),
 
     /// empty query plan. This behavior is unexpected and we suggest opening an issue to apollographql/router with a reproduction.
-    EmptyPlan(UsageReporting), // usage_reporting_signature
+    EmptyPlan(String), // usage_reporting stats_report_key
 
     /// unhandled planner result
     UnhandledPlannerResult,
@@ -434,10 +433,9 @@ impl IntoGraphQLErrors for QueryPlannerError {
 impl QueryPlannerError {
     pub(crate) fn usage_reporting(&self) -> Option<UsageReporting> {
         match self {
-            QueryPlannerError::SpecError(e) => Some(UsageReporting {
-                stats_report_key: e.get_error_key().to_string(),
-                referenced_fields_by_type: HashMap::new(),
-            }),
+            QueryPlannerError::SpecError(e) => {
+                Some(UsageReporting::Error(e.get_error_key().to_string()))
+            }
             _ => None,
         }
     }