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,158 @@ 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>,
+}
+
+impl UsageReportingOperationDetails {
+    fn operation_name_or_default(&self) -> String {
+        self.operation_name.as_deref().unwrap_or("").to_string()
+    }
+
+    fn operation_sig_or_default(&self) -> String {
+        self.operation_signature
+            .as_deref()
+            .unwrap_or("")
+            .to_string()
+    }
+
+    fn get_signature_and_operation(&self) -> String {
+        let op_name = self.operation_name.as_deref().unwrap_or("-").to_string();
+        let op_sig = self
+            .operation_signature
+            .as_deref()
+            .unwrap_or("")
+            .to_string();
+        format!("# {}\n{}", op_name, op_sig)
+    }
+}
+
 /// 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, persisted_query_id: String) -> UsageReporting {
+        match self {
+            UsageReporting::Operation(operation_details)
+            | UsageReporting::PersistedQuery {
+                operation_details, ..
+            } => UsageReporting::PersistedQuery {
+                operation_details: operation_details.clone(),
+                persisted_query_id,
+            },
+            // 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
-    /// 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>,
+    /// Metric data sent to Studio must be aggregated via grouped key of
+    /// (`client_name`, `client_version`, `stats_report_key`).
+    /// For errors, the report key is of the form "## <error name>\n".
+    /// For operations not requested by PQ, the report key is of the form "# <op name>\n<op sig>".
+    /// For operations requested by PQ, the report key is of the form "pq# <unique hash>", where the
+    /// unique hash is a string that is consistent for the same PQ and operation, but unique if either
+    /// is different. The actual PQ ID, operation name, and operation signature is passed as metadata.
+    /// We need to do this so that we can group stats for each combination of PQ and operation.
+    /// Note that the combination of signature and operation name is sometimes referred to in code as
+    /// the "operation signature".
+    pub(crate) fn get_stats_report_key(&self) -> String {
+        match self {
+            UsageReporting::Operation(operation_details) => {
+                operation_details.get_signature_and_operation()
+            }
+            UsageReporting::Error(error_key) => {
+                format!("## {}\n", error_key)
+            }
+            UsageReporting::PersistedQuery {
+                operation_details,
+                persisted_query_id,
+                ..
+            } => {
+                let string_to_hash = format!(
+                    "{}\n{}\n{}",
+                    persisted_query_id,
+                    operation_details.operation_name_or_default(),
+                    operation_details.operation_sig_or_default()
+                );
+                format!("pq# {}", Self::hash_string(&string_to_hash))
+            }
+        }
+    }
+
+    pub(crate) fn get_operation_id(&self) -> String {
+        let string_to_hash = match self {
+            UsageReporting::Operation(operation_details)
+            | UsageReporting::PersistedQuery {
+                operation_details, ..
+            } => operation_details.get_signature_and_operation(),
+            UsageReporting::Error(error_key) => {
+                format!("# # {}\n", error_key)
+            }
+        };
+        Self::hash_string(&string_to_hash)
+    }
+
+    pub(crate) fn get_operation_name(&self) -> String {
+        match self {
+            UsageReporting::Operation(operation_details)
+            | UsageReporting::PersistedQuery {
+                operation_details, ..
+            } => operation_details.operation_name_or_default(),
+            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 {
+                operation_details,
+                persisted_query_id,
+                ..
+            } => Some(QueryMetadata {
+                name: operation_details.operation_name_or_default(),
+                signature: operation_details.operation_sig_or_default(),
+                pq_id: persisted_query_id.clone(),
+            }),
+            // For now we only want to populate query metadata for PQ operations
+            UsageReporting::Operation { .. } | UsageReporting::Error { .. } => None,
+        }
+    }
+
+    fn hash_string(string_to_hash: &String) -> String {
+        let mut hasher = sha1::Sha1::new();
+        hasher.update(string_to_hash.as_bytes());
+        let result = hasher.finalize();
+        hex::encode(result)
+    }
 }
 
 /// A list of fields that will be resolved for a given type
@@ -186,9 +329,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 +505,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 +529,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 +563,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);