feat(chunk): Add attachments to profile chunks

src/android/chunk.rs

@@ -4,7 +4,7 @@ use serde::{Deserialize, Serialize};
 
 use crate::{
     nodetree::Node,
-    types::{CallTreeError, CallTreesStr, ChunkInterface, ClientSDK, DebugMeta},
+    types::{Attachment, CallTreeError, CallTreesStr, ChunkInterface, ClientSDK, DebugMeta},
 };
 
 use super::Android;
@@ -79,6 +79,14 @@ impl ChunkInterface for AndroidChunk {
         self.project_id
     }
 
+    // Attachments are only supported for sample chunks:
+    // the getter always returns an empty list and the setter is a no-op.
+    fn get_attachments(&self) -> &[Attachment] {
+        &[]
+    }
+
+    fn set_attachments(&mut self, _attachments: Vec<Attachment>) {}
+
     fn get_received(&self) -> f64 {
         self.received
     }

src/lib.rs

@@ -146,6 +146,7 @@ fn decompress_profile(profile: &[u8]) -> PyResult<Profile> {
 fn vroomrs(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<ProfileChunk>()?;
     m.add_class::<CallTreeFunction>()?;
+    m.add_class::<types::Attachment>()?;
     m.add_function(wrap_pyfunction!(profile_chunk_from_json_str, m)?)?;
     m.add_function(wrap_pyfunction!(decompress_profile_chunk, m)?)?;
     m.add_function(wrap_pyfunction!(profile_from_json_str, m)?)?;

src/profile_chunk.rs

@@ -6,7 +6,7 @@ use crate::{
     android::chunk::AndroidChunk,
     nodetree::CallTreeFunction,
     sample::v2::SampleChunk,
-    types::{CallTreesStr, ChunkInterface},
+    types::{Attachment, CallTreesStr, ChunkInterface},
     utils::{compress_lz4, decompress_lz4},
 };
 
@@ -18,45 +18,50 @@ pub struct ProfileChunk {
 
 #[derive(serde::Deserialize)]
 struct MinimumProfile {
+    #[serde(default)]
+    platform: String,
     version: Option<String>,
+    // Present only on legacy android profiles: the raw, unparsed profile.
+    // Ignored here, we only care whether it's set.
+    sampled_profile: Option<serde::de::IgnoredAny>,
 }
 
 impl ProfileChunk {
     pub(crate) fn from_json_vec(profile: &[u8]) -> Result<Self, serde_json::Error> {
         let min_prof: MinimumProfile = serde_json::from_slice(profile)?;
-        match min_prof.version {
-            None => {
-                let android: AndroidChunk = serde_json::from_slice(profile)?;
-                Ok(ProfileChunk {
-                    profile: Box::new(android),
-                })
-            }
-            Some(_) => {
-                let sample: SampleChunk = serde_json::from_slice(profile)?;
-                Ok(ProfileChunk {
-                    profile: Box::new(sample),
-                })
-            }
-        }
+        Self::from_json_vec_with_profile(profile, &min_prof.platform, &min_prof)
     }
 
     pub(crate) fn from_json_vec_and_platform(
         profile: &[u8],
         platform: &str,
     ) -> Result<Self, serde_json::Error> {
-        match platform {
-            "android" => {
-                let android: AndroidChunk = serde_json::from_slice(profile)?;
-                Ok(ProfileChunk {
-                    profile: Box::new(android),
-                })
-            }
-            _ => {
-                let sample: SampleChunk = serde_json::from_slice(profile)?;
-                Ok(ProfileChunk {
-                    profile: Box::new(sample),
-                })
-            }
+        let min_prof: MinimumProfile = serde_json::from_slice(profile)?;
+        Self::from_json_vec_with_profile(profile, platform, &min_prof)
+    }
+
+    fn from_json_vec_with_profile(
+        profile: &[u8],
+        platform: &str,
+        min_prof: &MinimumProfile,
+    ) -> Result<Self, serde_json::Error> {
+        // A chunk is a legacy android chunk when its platform is android and it
+        // either carries no version (the original legacy format) or still contains
+        // a `sampled_profile`. The latter is an interim fix: some legacy android
+        // profiles are sent with version "2" even though they don't follow the
+        // sample v2 format.
+        let is_legacy_android = platform == "android"
+            && (min_prof.version.is_none() || min_prof.sampled_profile.is_some());
+        if is_legacy_android {
+            let android: AndroidChunk = serde_json::from_slice(profile)?;
+            Ok(ProfileChunk {
+                profile: Box::new(android),
+            })
+        } else {
+            let sample: SampleChunk = serde_json::from_slice(profile)?;
+            Ok(ProfileChunk {
+                profile: Box::new(sample),
+            })
         }
     }
 
@@ -129,6 +134,28 @@ impl ProfileChunk {
         self.profile.get_project_id()
     }
 
+    /// Returns the attachments related to this chunk.
+    ///
+    /// Returns:
+    ///     list[Attachment]
+    ///         The attachments (e.g. a raw profile) related to this chunk.
+    ///         Empty if no attachments are available.
+    pub fn get_attachments(&self) -> Vec<Attachment> {
+        self.profile.get_attachments().to_vec()
+    }
+
+    /// Sets the attachments related to this chunk.
+    ///
+    /// Attachments are only supported for sample chunks:
+    /// this is a no-op for Android chunks.
+    ///
+    /// Args:
+    ///     attachments (list[Attachment]): The attachments related to this
+    ///         chunk, replacing any existing ones. An empty list clears them.
+    pub fn set_attachments(&mut self, attachments: Vec<Attachment>) {
+        self.profile.set_attachments(attachments);
+    }
+
     /// Returns the received timestamp.
     ///
     /// Returns:
@@ -365,6 +392,83 @@ mod tests {
         }
     }
 
+    #[test]
+    fn test_android_platform_with_version_is_sample_chunk() {
+        // A chunk with platform=android but a version set uses the
+        // sample v2 format and must not be treated as a legacy
+        // android chunk, no matter how it's deserialized.
+        let payload = include_bytes!("../tests/fixtures/sample/v2/valid_cocoa.json");
+        let mut value: serde_json::Value = serde_json::from_slice(payload).unwrap();
+        value["platform"] = "android".into();
+        let json = serde_json::to_vec(&value).unwrap();
+
+        for chunk in [
+            ProfileChunk::from_json_vec(&json).unwrap(),
+            ProfileChunk::from_json_vec_and_platform(&json, "android").unwrap(),
+        ] {
+            assert_eq!(chunk.get_platform(), "android");
+            assert!(chunk
+                .profile
+                .as_any()
+                .downcast_ref::<SampleChunk>()
+                .is_some());
+        }
+
+        // Legacy android chunks (no version) still deserialize as such.
+        let payload = include_bytes!("../tests/fixtures/android/chunk/valid.json");
+        let chunk = ProfileChunk::from_json_vec_and_platform(payload, "android").unwrap();
+        assert!(chunk
+            .profile
+            .as_any()
+            .downcast_ref::<AndroidChunk>()
+            .is_some());
+    }
+
+    #[test]
+    fn test_legacy_android_with_sampled_profile_is_android_chunk() {
+        // Interim fix: some legacy android profiles are sent with a version
+        // (e.g. "2") even though they carry a `sampled_profile` and don't
+        // follow the sample v2 format. They must be treated as android chunks.
+        let payload = include_bytes!("../tests/fixtures/android/chunk/valid.json");
+        let mut value: serde_json::Value = serde_json::from_slice(payload).unwrap();
+        value["version"] = "2".into();
+        value["sampled_profile"] = "AAAAA".into();
+        let json = serde_json::to_vec(&value).unwrap();
+
+        for chunk in [
+            ProfileChunk::from_json_vec(&json).unwrap(),
+            ProfileChunk::from_json_vec_and_platform(&json, "android").unwrap(),
+        ] {
+            assert!(chunk
+                .profile
+                .as_any()
+                .downcast_ref::<AndroidChunk>()
+                .is_some());
+        }
+    }
+
+    #[test]
+    fn test_attachments_survive_compression() {
+        use crate::types::Attachment;
+
+        // The sentry writer flow: deserialize the chunk, stamp the
+        // attachments, compress and store. The attachments must survive
+        // into the stored chunk representation.
+        let payload = include_bytes!("../tests/fixtures/sample/v2/valid_cocoa.json");
+        let mut chunk = ProfileChunk::from_json_vec(payload).unwrap();
+        let attachments = vec![Attachment {
+            name: "raw_profile".to_string(),
+            content_type: Some("application/x-perfetto".to_string()),
+            stored_id: "aef123345".to_string(),
+        }];
+        chunk.set_attachments(attachments.clone());
+        assert_eq!(chunk.get_attachments(), attachments);
+
+        let compressed = chunk.compress().unwrap();
+        let decompressed = ProfileChunk::decompress(compressed.as_slice()).unwrap();
+        assert_eq!(decompressed.get_attachments(), attachments);
+    }
+
     #[test]
     fn test_from_json_vec_and_platform() {
         struct TestStruct<'a> {

src/sample/v2.rs

@@ -10,7 +10,7 @@ use std::rc::Rc;
 use super::{SampleError, ThreadMetadata};
 use crate::frame::Frame;
 use crate::nodetree::Node;
-use crate::types::{CallTreeError, CallTreesStr, ChunkInterface};
+use crate::types::{Attachment, CallTreeError, CallTreesStr, ChunkInterface};
 use crate::types::{ClientSDK, DebugMeta};
 
 #[derive(Serialize, Deserialize, Debug, Default, PartialEq)]
@@ -46,6 +46,9 @@ pub struct SampleChunk {
     // `measurements` contains CPU/memory measurements we do during the capture of the chunk.
     #[serde(skip_serializing_if = "Option::is_none")]
     pub measurements: Option<serde_json::Value>,
+
+    #[serde(default, skip_serializing_if = "Vec::is_empty")]
+    pub attachments: Vec<Attachment>,
 }
 
 #[derive(Serialize, Deserialize, Debug, Default, PartialEq)]
@@ -243,6 +246,14 @@ impl ChunkInterface for SampleChunk {
         self.project_id
     }
 
+    fn get_attachments(&self) -> &[Attachment] {
+        &self.attachments
+    }
+
+    fn set_attachments(&mut self, attachments: Vec<Attachment>) {
+        self.attachments = attachments;
+    }
+
     fn get_received(&self) -> f64 {
         self.received
     }
@@ -331,6 +342,53 @@ mod tests {
         assert!(r.is_ok(), "{r:#?}")
     }
 
+    #[test]
+    fn test_attachments() {
+        use crate::types::Attachment;
+
+        let payload = include_bytes!("../../tests/fixtures/sample/v2/valid_cocoa.json");
+        let mut value: serde_json::Value = serde_json::from_slice(payload).unwrap();
+
+        // An absent field deserializes to an empty list,
+        // which is skipped during serialization.
+        let chunk: SampleChunk = serde_json::from_value(value.clone()).unwrap();
+        assert!(chunk.attachments.is_empty());
+        let serialized = serde_json::to_value(&chunk).unwrap();
+        assert!(serialized.get("attachments").is_none());
+
+        // A present field round-trips.
+        let attachments_json = serde_json::json!([{
+            "name": "raw_profile",
+            "content_type": "application/x-perfetto",
+            "stored_id": "aef123345"
+        }]);
+        value["attachments"] = attachments_json.clone();
+        let mut chunk: SampleChunk = serde_json::from_value(value).unwrap();
+        assert_eq!(
+            chunk.get_attachments(),
+            &[Attachment {
+                name: "raw_profile".to_string(),
+                content_type: Some("application/x-perfetto".to_string()),
+                stored_id: "aef123345".to_string(),
+            }]
+        );
+        let serialized = serde_json::to_value(&chunk).unwrap();
+        assert_eq!(serialized["attachments"], attachments_json);
+
+        // The setter overwrites and clears the list.
+        chunk.set_attachments(vec![Attachment {
+            name: "raw_profile".to_string(),
+            content_type: None,
+            stored_id: "fff999".to_string(),
+        }]);
+        assert_eq!(chunk.get_attachments()[0].stored_id, "fff999");
+        assert_eq!(chunk.get_attachments()[0].content_type, None);
+        chunk.set_attachments(Vec::new());
+        assert!(chunk.get_attachments().is_empty());
+        let serialized = serde_json::to_value(&chunk).unwrap();
+        assert!(serialized.get("attachments").is_none());
+    }
+
     #[test]
     fn test_call_trees() {
         use crate::nodetree::Node;

src/types.rs

@@ -1,7 +1,7 @@
 use chrono::{DateTime, Utc};
 
 use pyo3::exceptions::PyValueError;
-use pyo3::{pyclass, PyErr};
+use pyo3::{pyclass, pymethods, PyErr};
 use serde::{Deserialize, Serialize};
 use std::any::Any;
 use std::borrow::Cow;
@@ -83,6 +83,8 @@ pub trait ChunkInterface {
     fn get_platform(&self) -> String;
     fn get_profiler_id(&self) -> &str;
     fn get_project_id(&self) -> u64;
+    fn get_attachments(&self) -> &[Attachment];
+    fn set_attachments(&mut self, attachments: Vec<Attachment>);
     fn get_received(&self) -> f64;
     fn get_release(&self) -> Option<&str>;
     fn get_retention_days(&self) -> i32;
@@ -107,6 +109,50 @@ pub trait ChunkInterface {
     fn as_any(&self) -> &dyn Any;
 }
 
+/// A file related to the chunk (e.g. a raw profile), stored in the object store.
+#[pyclass(eq)]
+#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
+pub struct Attachment {
+    /// The attachment file name, e.g. `raw_profile.pftrace`.
+    #[pyo3(get)]
+    pub name: String,
+    /// The MIME content type of the attachment (e.g. `application/x-perfetto`), if known.
+    ///
+    /// Intentionally kept as a free-form string for now to stay open to new
+    /// content types; it is passed through and not interpreted within this repo.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    #[pyo3(get)]
+    pub content_type: Option<String>,
+    /// The identifier of the attachment in the object store.
+    #[pyo3(get)]
+    pub stored_id: String,
+}
+
+#[pymethods]
+impl Attachment {
+    #[new]
+    #[pyo3(signature = (name, content_type, stored_id))]
+    fn new(name: String, content_type: Option<String>, stored_id: String) -> Self {
+        Attachment {
+            name,
+            content_type,
+            stored_id,
+        }
+    }
+
+    fn __repr__(&self) -> String {
+        format!(
+            "Attachment(name={:?}, content_type={}, stored_id={:?})",
+            self.name,
+            match &self.content_type {
+                Some(content_type) => format!("{content_type:?}"),
+                None => "None".to_string(),
+            },
+            self.stored_id
+        )
+    }
+}
+
 #[pyclass]
 #[derive(Debug, Serialize, Deserialize, Clone, Default, PartialEq)]
 pub struct Transaction {