Add component_path to usage data sent from backend (#28031)

This PR adds `component_path` to usage events sent from backend.

GitOrigin-RevId: fa4332926215bd33df427ce91799b5e9f1405a1b

Author: emmaling27

crates/common/src/components/function_paths.rs

@@ -98,6 +98,10 @@ impl CanonicalizedComponentFunctionPath {
     pub fn is_system(&self) -> bool {
         self.udf_path.is_system()
     }
+
+    pub fn into_component_and_udf_path(self) -> (ComponentPath, CanonicalizedUdfPath) {
+        (self.component, self.udf_path)
+    }
 }
 
 impl From<CanonicalizedComponentFunctionPath> for ComponentFunctionPath {

crates/common/src/types/functions.rs

@@ -111,6 +111,18 @@ pub enum UdfIdentifier {
     Cli(String),
 }
 
+impl UdfIdentifier {
+    pub fn into_component_and_udf_path(self) -> (Option<String>, String) {
+        match self {
+            UdfIdentifier::Function(path) => {
+                let (component_path, udf_path) = path.clone().into_component_and_udf_path();
+                (component_path.serialize(), udf_path.to_string())
+            },
+            UdfIdentifier::Http(_) | UdfIdentifier::Cli(_) => (None, self.to_string()),
+        }
+    }
+}
+
 impl fmt::Display for UdfIdentifier {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match self {

crates/events/src/testing.rs

@@ -91,6 +91,7 @@ impl UsageCounterState {
     fn record_event(&mut self, event: UsageEvent) {
         match event {
             UsageEvent::FunctionCall {
+                component_path,
                 udf_id,
                 tag,
                 memory_megabytes,
@@ -100,7 +101,14 @@ impl UsageCounterState {
                 ..
             } => {
                 if is_tracked {
-                    *self.recent_calls.entry(udf_id.to_string()).or_default() += 1;
+                    let fn_name = if let Some(mut component) = component_path {
+                        component.push('/');
+                        component.push_str(&udf_id);
+                        component
+                    } else {
+                        udf_id.clone()
+                    };
+                    *self.recent_calls.entry(fn_name).or_default() += 1;
                     *self.recent_calls_by_tag.entry(tag).or_default() += 1;
 
                     // Convert into MB-milliseconds of compute time

crates/events/src/usage.rs

@@ -20,6 +20,9 @@ pub enum UsageEvent {
     FunctionCall {
         // The ExecutionId of a particular UDF
         id: String,
+        /// The path of a component. Uniquely identifies a component in a
+        /// project.
+        component_path: Option<String>,
         // The path / name of the UDF
         udf_id: String,
         // The type of the udf identifier (http, function, cli)
@@ -43,6 +46,7 @@ pub enum UsageEvent {
     /// invocation.
     FunctionStorageCalls {
         id: String,
+        component_path: Option<String>,
         udf_id: String,
         call: String,
         count: u64,
@@ -51,6 +55,7 @@ pub enum UsageEvent {
     /// function invocation.
     FunctionStorageBandwidth {
         id: String,
+        component_path: Option<String>,
         udf_id: String,
         ingress: u64,
         egress: u64,
@@ -70,13 +75,15 @@ pub enum UsageEvent {
     },
     DatabaseBandwidth {
         id: String,
+        component_path: Option<String>,
         udf_id: String,
         table_name: String,
         ingress: u64,
         egress: u64,
     },
     VectorBandwidth {
         id: String,
+        component_path: Option<String>,
         udf_id: String,
         table_name: String,
         ingress: u64,

crates/usage_tracking/src/lib.rs

@@ -118,9 +118,11 @@ impl UsageCounter {
             UdfIdentifier::Http(_) => (true, "http"),
             UdfIdentifier::Cli(_) => (false, "cli"),
         };
+        let (component_path, udf_id) = udf_path.clone().into_component_and_udf_path();
         usage_metrics.push(UsageEvent::FunctionCall {
             id: execution_id.to_string(),
-            udf_id: udf_path.to_string(),
+            component_path,
+            udf_id,
             udf_id_type: udf_id_type.to_string(),
             tag: call_type.tag().to_string(),
             memory_megabytes: call_type.memory_megabytes(),
@@ -158,25 +160,29 @@ impl UsageCounter {
         usage_metrics: &mut Vec<UsageEvent>,
     ) {
         // Merge the storage stats.
+        let (component_path, udf_id) = udf_path.clone().into_component_and_udf_path();
         for (storage_api, function_count) in stats.storage_calls {
             usage_metrics.push(UsageEvent::FunctionStorageCalls {
                 id: execution_id.to_string(),
-                udf_id: udf_path.to_string(),
+                component_path: component_path.clone(),
+                udf_id: udf_id.clone(),
                 call: storage_api,
                 count: function_count,
             });
         }
         usage_metrics.push(UsageEvent::FunctionStorageBandwidth {
             id: execution_id.to_string(),
-            udf_id: udf_path.to_string(),
+            component_path: component_path.clone(),
+            udf_id: udf_id.clone(),
             ingress: stats.storage_ingress_size,
             egress: stats.storage_egress_size,
         });
         // Merge "by table" bandwidth stats.
         for (table_name, ingress_size) in stats.database_ingress_size {
             usage_metrics.push(UsageEvent::DatabaseBandwidth {
                 id: execution_id.to_string(),
-                udf_id: udf_path.to_string(),
+                component_path: component_path.clone(),
+                udf_id: udf_id.clone(),
                 table_name,
                 ingress: ingress_size,
                 egress: 0,
@@ -185,7 +191,8 @@ impl UsageCounter {
         for (table_name, egress_size) in stats.database_egress_size {
             usage_metrics.push(UsageEvent::DatabaseBandwidth {
                 id: execution_id.to_string(),
-                udf_id: udf_path.to_string(),
+                component_path: component_path.clone(),
+                udf_id: udf_id.clone(),
                 table_name,
                 ingress: 0,
                 egress: egress_size,
@@ -194,7 +201,8 @@ impl UsageCounter {
         for (table_name, ingress_size) in stats.vector_ingress_size {
             usage_metrics.push(UsageEvent::VectorBandwidth {
                 id: execution_id.to_string(),
-                udf_id: udf_path.to_string(),
+                component_path: component_path.clone(),
+                udf_id: udf_id.clone(),
                 table_name,
                 ingress: ingress_size,
                 egress: 0,
@@ -203,7 +211,8 @@ impl UsageCounter {
         for (table_name, egress_size) in stats.vector_egress_size {
             usage_metrics.push(UsageEvent::VectorBandwidth {
                 id: execution_id.to_string(),
-                udf_id: udf_path.to_string(),
+                component_path: component_path.clone(),
+                udf_id: udf_id.clone(),
                 table_name,
                 ingress: 0,
                 egress: egress_size,