feat: add support of operationCountByType for apollo studio (#2979)

add support of operationCountByType for apollo studio

---------

Signed-off-by: Benjamin Coenen <[email protected]>

Author: bnjjj

.changesets/maint_climber_northeast_inspector_cried.md

@@ -0,0 +1,5 @@
+### Add support of operationCountByType for apollo studio ([PR #2979](https://github.com/apollographql/router/pull/2979))
+
+Add more details about the operation type for apollo studio usage reporting.
+
+By [@bnjjj](https://github.com/bnjjj) in https://github.com/apollographql/router/pull/2979

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

@@ -1,6 +1,7 @@
 //! Configuration for apollo telemetry.
 // This entire file is license key functionality
 use std::collections::HashMap;
+use std::fmt::Display;
 use std::num::NonZeroUsize;
 use std::ops::AddAssign;
 use std::time::SystemTime;
@@ -9,6 +10,7 @@ use derivative::Derivative;
 use http::header::HeaderName;
 use itertools::Itertools;
 use schemars::JsonSchema;
+use serde::ser::SerializeMap;
 use serde::Deserialize;
 use serde::Serialize;
 use url::Url;
@@ -25,6 +27,7 @@ use crate::plugins::telemetry::apollo_exporter::proto::reports::StatsContext;
 use crate::plugins::telemetry::apollo_exporter::proto::reports::Trace;
 use crate::plugins::telemetry::config::SamplerOption;
 use crate::plugins::telemetry::tracing::BatchProcessorConfig;
+use crate::query_planner::OperationKind;
 use crate::services::apollo_graph_reference;
 use crate::services::apollo_key;
 
@@ -209,9 +212,71 @@ pub(crate) enum SingleReport {
 #[derive(Default, Debug, Serialize)]
 pub(crate) struct Report {
     pub(crate) traces_per_query: HashMap<String, TracesAndStats>,
+    #[serde(serialize_with = "serialize_operation_count_by_type")]
+    pub(crate) operation_count_by_type:
+        HashMap<(OperationKind, Option<OperationSubType>), OperationCountByType>,
+}
+
+#[derive(Default, Debug, Serialize, PartialEq, Eq, Hash)]
+pub(crate) struct OperationCountByType {
+    pub(crate) r#type: OperationKind,
+    pub(crate) subtype: Option<OperationSubType>,
     pub(crate) operation_count: u64,
 }
 
+#[derive(Debug, Serialize, PartialEq, Eq, Hash, Clone, Copy)]
+#[serde(rename_all = "kebab-case")]
+pub(crate) enum OperationSubType {
+    // TODO
+}
+
+impl OperationSubType {
+    pub(crate) const fn as_str(&self) -> &'static str {
+        ""
+    }
+}
+
+impl Display for OperationSubType {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "")
+    }
+}
+
+impl From<OperationCountByType>
+    for crate::plugins::telemetry::apollo_exporter::proto::reports::report::OperationCountByType
+{
+    fn from(value: OperationCountByType) -> Self {
+        Self {
+            r#type: value.r#type.as_apollo_operation_type().to_string(),
+            subtype: value.subtype.map(|s| s.to_string()).unwrap_or_default(),
+            operation_count: value.operation_count,
+        }
+    }
+}
+
+fn serialize_operation_count_by_type<S>(
+    elt: &HashMap<(OperationKind, Option<OperationSubType>), OperationCountByType>,
+    serializer: S,
+) -> Result<S::Ok, S::Error>
+where
+    S: serde::Serializer,
+{
+    let mut map_ser = serializer.serialize_map(Some(elt.len()))?;
+    for ((op_type, op_subtype), v) in elt {
+        map_ser.serialize_entry(
+            &format!(
+                "{}{}",
+                op_type.as_apollo_operation_type(),
+                op_subtype
+                    .map(|o| "/".to_owned() + o.as_str())
+                    .unwrap_or_default()
+            ),
+            v,
+        )?;
+    }
+    map_ser.end()
+}
+
 impl Report {
     #[cfg(test)]
     pub(crate) fn new(reports: Vec<SingleStatsReport>) -> Report {
@@ -229,7 +294,11 @@ impl Report {
         let mut report = crate::plugins::telemetry::apollo_exporter::proto::reports::Report {
             header: Some(header),
             end_time: Some(SystemTime::now().into()),
-            operation_count: self.operation_count,
+            operation_count_by_type: self
+                .operation_count_by_type
+                .into_values()
+                .map(|op| op.into())
+                .collect(),
             traces_pre_aggregated: true,
             ..Default::default()
         };
@@ -269,7 +338,18 @@ impl AddAssign<SingleStatsReport> for Report {
             *self.traces_per_query.entry(k).or_default() += v;
         }
 
-        self.operation_count += report.operation_count;
+        if let Some(operation_count_by_type) = report.operation_count_by_type {
+            let key = (
+                operation_count_by_type.r#type,
+                operation_count_by_type.subtype,
+            );
+            self.operation_count_by_type
+                .entry(key)
+                .and_modify(|e| {
+                    e.operation_count += 1;
+                })
+                .or_insert(operation_count_by_type);
+        }
     }
 }
 

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

@@ -164,7 +164,7 @@ impl ApolloExporter {
 
     pub(crate) async fn submit_report(&self, report: Report) -> Result<(), ApolloExportError> {
         // We may be sending traces but with no operation count
-        if report.operation_count == 0 && report.traces_per_query.is_empty() {
+        if report.operation_count_by_type.is_empty() && report.traces_per_query.is_empty() {
             return Ok(());
         }
         tracing::debug!("submitting report: {:?}", report);

apollo-router/src/plugins/telemetry/metrics/apollo/snapshots/apollo_router__plugins__telemetry__metrics__apollo__studio__test__aggregation.snap

@@ -1302,5 +1302,11 @@ expression: aggregated_metrics
       }
     }
   },
-  "operation_count": 2
+  "operation_count_by_type": {
+    "query": {
+      "type": "query",
+      "subtype": null,
+      "operation_count": 2
+    }
+  }
 }

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

@@ -6,14 +6,15 @@ use serde::Serialize;
 use uuid::Uuid;
 
 use super::duration_histogram::DurationHistogram;
+use crate::plugins::telemetry::apollo::OperationCountByType;
 use crate::plugins::telemetry::apollo_exporter::proto::reports::ReferencedFieldsForType;
 use crate::plugins::telemetry::apollo_exporter::proto::reports::StatsContext;
 
 #[derive(Default, Debug, Serialize)]
 pub(crate) struct SingleStatsReport {
     pub(crate) request_id: Uuid,
     pub(crate) stats: HashMap<String, SingleStats>,
-    pub(crate) operation_count: u64,
+    pub(crate) operation_count_by_type: Option<OperationCountByType>,
 }
 
 #[derive(Default, Debug, Serialize)]
@@ -269,6 +270,7 @@ mod test {
 
     use super::*;
     use crate::plugins::telemetry::apollo::Report;
+    use crate::query_planner::OperationKind;
 
     #[test]
     fn test_aggregation() {
@@ -317,7 +319,12 @@ mod test {
 
         SingleStatsReport {
             request_id: Uuid::default(),
-            operation_count: count.inc_u64(),
+            operation_count_by_type: OperationCountByType {
+                r#type: OperationKind::Query,
+                subtype: None,
+                operation_count: count.inc_u64(),
+            }
+            .into(),
             stats: HashMap::from([(
                 stats_report_key.to_string(),
                 SingleStats {

apollo-router/src/plugins/telemetry/metrics/snapshots/apollo_router__plugins__telemetry__metrics__apollo__test__apollo_metrics_exclude.snap

@@ -6,6 +6,10 @@ expression: results
   {
     "request_id": "[REDACTED]",
     "stats": {},
-    "operation_count": 1
+    "operation_count_by_type": {
+      "type": "query",
+      "subtype": null,
+      "operation_count": 1
+    }
   }
 ]

apollo-router/src/plugins/telemetry/metrics/snapshots/apollo_router__plugins__telemetry__metrics__apollo__test__apollo_metrics_multiple_operations.snap

@@ -11,7 +11,7 @@ expression: results
           "context": {
             "client_name": "test_client",
             "client_version": "1.0-test",
-            "operation_type": "",
+            "operation_type": "query",
             "operation_subtype": ""
           },
           "query_latency_stats": {
@@ -39,6 +39,6 @@ expression: results
         "referenced_fields_by_type": {}
       }
     },
-    "operation_count": 0
+    "operation_count_by_type": null
   }
 ]

apollo-router/src/plugins/telemetry/metrics/snapshots/apollo_router__plugins__telemetry__metrics__apollo__test__apollo_metrics_parse_failure.snap

@@ -11,7 +11,7 @@ expression: results
           "context": {
             "client_name": "test_client",
             "client_version": "1.0-test",
-            "operation_type": "",
+            "operation_type": "query",
             "operation_subtype": ""
           },
           "query_latency_stats": {
@@ -39,6 +39,6 @@ expression: results
         "referenced_fields_by_type": {}
       }
     },
-    "operation_count": 0
+    "operation_count_by_type": null
   }
 ]

apollo-router/src/plugins/telemetry/metrics/snapshots/apollo_router__plugins__telemetry__metrics__apollo__test__apollo_metrics_single_operation.snap

@@ -11,7 +11,7 @@ expression: results
           "context": {
             "client_name": "test_client",
             "client_version": "1.0-test",
-            "operation_type": "",
+            "operation_type": "query",
             "operation_subtype": ""
           },
           "query_latency_stats": {
@@ -52,6 +52,10 @@ expression: results
         }
       }
     },
-    "operation_count": 1
+    "operation_count_by_type": {
+      "type": "query",
+      "subtype": null,
+      "operation_count": 1
+    }
   }
 ]

apollo-router/src/plugins/telemetry/metrics/snapshots/apollo_router__plugins__telemetry__metrics__apollo__test__apollo_metrics_unknown_operation.snap

@@ -11,7 +11,7 @@ expression: results
           "context": {
             "client_name": "test_client",
             "client_version": "1.0-test",
-            "operation_type": "",
+            "operation_type": "query",
             "operation_subtype": ""
           },
           "query_latency_stats": {
@@ -39,6 +39,6 @@ expression: results
         "referenced_fields_by_type": {}
       }
     },
-    "operation_count": 0
+    "operation_count_by_type": null
   }
 ]