ref(attributes): Introduce AttributeLike for shared normalizations

relay-event-normalization/src/eap/attribute_like.rs

@@ -0,0 +1,93 @@
+use relay_event_schema::protocol::{Attribute, Attributes, SpanData};
+use relay_protocol::{Annotated, Object, Value};
+
+/// An attribute collection which is keyed by an attribute key defined in [`relay_conventions`].
+///
+/// This exists as a common abstraction over [`SpanData`] and [`Attributes`]. Its purpose is to
+/// allow modern attribute based normalizations also to apply to transaction based spans,
+/// eliminating the need to duplicate code.
+///
+/// This largely mirrors the API of [`Attributes`] and provides it also for [`SpanData`].
+pub trait AttributesLike {
+    /// The attribute value.
+    type Value: AttributeLike;
+
+    /// Access this container through the underlying [`Object`].
+    fn as_object(&self) -> &Object<Self::Value>;
+    /// Access this container through the underlying [`Object`] mutably.
+    fn as_object_mut(&mut self) -> &mut Object<Self::Value>;
+
+    /// Checks whether this collection contains an attribute with the given `key`.
+    fn contains_key(&self, key: &str) -> bool {
+        self.as_object().contains_key(key)
+    }
+
+    /// Inserts an attribute with the given value into the collection.
+    fn insert(&mut self, key: String, value: Annotated<Self::Value>) {
+        self.as_object_mut().insert(key, value);
+    }
+
+    /// Removes a `key` from the attribute collection returning its value.
+    fn remove(&mut self, key: &str) -> Option<Annotated<Self::Value>> {
+        self.as_object_mut().remove(key)
+    }
+}
+
+impl AttributesLike for SpanData {
+    type Value = Value;
+
+    fn as_object(&self) -> &Object<Self::Value> {
+        &self.other
+    }
+
+    fn as_object_mut(&mut self) -> &mut Object<Self::Value> {
+        &mut self.other
+    }
+}
+
+impl AttributesLike for Attributes {
+    type Value = Attribute;
+
+    fn as_object(&self) -> &Object<Self::Value> {
+        &self.0
+    }
+
+    fn as_object_mut(&mut self) -> &mut Object<Self::Value> {
+        &mut self.0
+    }
+}
+
+/// An attribute value stored in [`AttributesLike`].
+///
+/// This only allows read only accessors as mutating the attribute value may not be allowed.
+/// For example [`Attribute`] has a `type` field which must match its `value field.
+pub trait AttributeLike: Clone + From<String> + From<i64> + From<f64> {
+    /// Returns a reference to the stored value in the attribute.
+    fn as_value(&self) -> Option<&Value>;
+
+    /// Returns the stored value as a `str`.
+    ///
+    /// Is `None` when the stored value is not a string.
+    fn as_str(&self) -> Option<&str> {
+        self.as_value().and_then(Value::as_str)
+    }
+
+    /// Returns the stored value as a `f64`.
+    ///
+    /// Is `None` when the stored value is not a float.
+    fn as_f64(&self) -> Option<f64> {
+        self.as_value().and_then(Value::as_f64)
+    }
+}
+
+impl AttributeLike for Value {
+    fn as_value(&self) -> Option<&Value> {
+        Some(self)
+    }
+}
+
+impl AttributeLike for Attribute {
+    fn as_value(&self) -> Option<&Value> {
+        self.value.value.value()
+    }
+}

relay-event-normalization/src/eap/mod.rs

@@ -3,7 +3,6 @@
 //! A central place for all modifications/normalizations for attributes.
 
 use std::borrow::Cow;
-use std::collections::BTreeMap;
 use std::fmt;
 use std::net::IpAddr;
 
@@ -30,13 +29,15 @@ use crate::{
 };
 
 mod ai;
+mod attribute_like;
 mod mobile;
 mod size;
 pub mod time;
 pub mod trace_metric;
 mod trimming;
 
 pub use self::ai::normalize_ai;
+pub use self::attribute_like::AttributesLike;
 pub use self::mobile::{normalize_mobile_attributes, normalize_mobile_measurements};
 pub use self::size::*;
 pub use self::trimming::TrimmingProcessor;
@@ -578,23 +579,21 @@ pub fn normalize_client_sample_rate(attributes: &mut Annotated<Attributes>) {
 ///
 /// Attributes with a status of `"backfill"` will be copied to their replacement name if the
 /// replacement name is not present. In any case, the original name is left alone.
-pub fn normalize_attribute_names(attributes: &mut Annotated<Attributes>) {
+pub fn normalize_attribute_names(attributes: &mut Annotated<impl AttributesLike>) {
     let Some(attributes) = attributes.value_mut() else {
         return;
     };
 
     normalize_attribute_names_inner(
-        &mut attributes.0,
+        attributes.as_object_mut(),
         relay_conventions::attribute_info_with_fragment,
     )
 }
 
 type AttributeInfoFn = fn(&str) -> Option<(&'static AttributeInfo, Option<&str>)>;
 
-fn normalize_attribute_names_inner<T>(
-    attributes: &mut BTreeMap<String, Annotated<T>>,
-    attribute_info: AttributeInfoFn,
-) where
+fn normalize_attribute_names_inner<T>(attributes: &mut Object<T>, attribute_info: AttributeInfoFn)
+where
     T: Clone,
 {
     let attribute_names: Vec<_> = attributes.keys().cloned().collect();
@@ -649,11 +648,6 @@ fn normalize_attribute_names_inner<T>(
     }
 }
 
-/// Like [`normalize_attribute_names`] but works on [`Object`] instead of [`Attributes`].
-pub fn normalize_attribute_names_obj(attributes: &mut Object<Value>) {
-    normalize_attribute_names_inner(attributes, relay_conventions::attribute_info_with_fragment)
-}
-
 /// Resolves the name of a replacement attribute for rewriting.
 ///
 /// There are two cases to consider:

relay-event-normalization/src/normalize/span/mod.rs

@@ -34,18 +34,18 @@ pub fn normalize_conventions(event: &mut Event) {
         .value_mut()
         .as_mut()
         .and_then(|c| c.get_mut::<TraceContext>())
-        .and_then(|c| c.data.value_mut().as_mut())
+        .map(|c| &mut c.data)
     {
-        crate::eap::normalize_attribute_names_obj(&mut data.other);
+        crate::eap::normalize_attribute_names(data);
     }
 
     if let Some(spans) = event.spans.value_mut() {
         for data in spans
             .iter_mut()
             .filter_map(|span| span.value_mut().as_mut())
-            .filter_map(|span| span.data.value_mut().as_mut())
+            .map(|span| &mut span.data)
         {
-            crate::eap::normalize_attribute_names_obj(&mut data.other);
+            crate::eap::normalize_attribute_names(data);
         }
     }
 }

relay-server/src/processing/legacy_spans/normalize.rs

@@ -119,12 +119,8 @@ pub fn normalize(
         dsc,
     } = config;
 
-    if let Some(data) = annotated_span
-        .value_mut()
-        .as_mut()
-        .and_then(|s| s.data.value_mut().as_mut())
-    {
-        relay_event_normalization::eap::normalize_attribute_names_obj(&mut data.other);
+    if let Some(data) = annotated_span.value_mut().as_mut().map(|s| &mut s.data) {
+        relay_event_normalization::eap::normalize_attribute_names(data);
     }
 
     set_segment_attributes(annotated_span);