diff --git a/src/android/chunk.rs b/src/android/chunk.rs index 3dc336a..c3162e2 100644 --- a/src/android/chunk.rs +++ b/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) {} + fn get_received(&self) -> f64 { self.received } diff --git a/src/lib.rs b/src/lib.rs index 477f7b1..7728f66 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -146,6 +146,7 @@ fn decompress_profile(profile: &[u8]) -> PyResult { fn vroomrs(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; + m.add_class::()?; 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)?)?; diff --git a/src/profile_chunk.rs b/src/profile_chunk.rs index e38f9fd..357c1c8 100644 --- a/src/profile_chunk.rs +++ b/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}, }; @@ -45,12 +45,10 @@ impl ProfileChunk { platform: &str, ) -> Result { match platform { - "android" => { - let android: AndroidChunk = serde_json::from_slice(profile)?; - Ok(ProfileChunk { - profile: Box::new(android), - }) - } + // Only profiles without a version use the legacy android format: + // android profiles with a version (e.g. "2") use the sample format, + // so we fall back to version-based detection. + "android" => Self::from_json_vec(profile), _ => { let sample: SampleChunk = serde_json::from_slice(profile)?; Ok(ProfileChunk { @@ -129,6 +127,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 { + 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) { + self.profile.set_attachments(attachments); + } + /// Returns the received timestamp. /// /// Returns: @@ -365,6 +385,60 @@ 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::() + .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::() + .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> { diff --git a/src/sample/v2.rs b/src/sample/v2.rs index 91ae500..7e17e7f 100644 --- a/src/sample/v2.rs +++ b/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(default, skip_serializing_if = "Vec::is_empty")] + pub attachments: Vec, } #[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) { + 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; diff --git a/src/types.rs b/src/types.rs index 3cbdd25..0fc99c1 100644 --- a/src/types.rs +++ b/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); 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 name of the attachment (e.g. `raw_profile`). + #[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, + /// 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, 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 { diff --git a/vroomrs.pyi b/vroomrs.pyi index 7f19a8a..bd103f5 100644 --- a/vroomrs.pyi +++ b/vroomrs.pyi @@ -316,7 +316,30 @@ class ProfileChunk: int: The project ID to which the profile belongs. """ ... - + + def get_attachments(self) -> List["Attachment"]: + """ + 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. + """ + ... + + def set_attachments(self, attachments: List["Attachment"]) -> None: + """ + 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. + """ + ... + def get_received(self) -> float: """ Returns the received timestamp. @@ -325,7 +348,7 @@ class ProfileChunk: float: The received timestamp. """ ... - + def get_release(self) -> Optional[str]: """ Returns the release. @@ -460,6 +483,22 @@ class ProfileChunk: """ ... +class Attachment: + """ + A file related to the chunk (e.g. a raw profile), stored in the object store. + """ + + name: str + """The attachment kind, e.g. `raw_profile`.""" + + content_type: Optional[str] + """The content type of the attachment, or None if not available.""" + + stored_id: str + """The object store ID of the attachment.""" + + def __init__(self, name: str, content_type: Optional[str], stored_id: str) -> None: ... + class CallTreeFunction: """ Represents function metrics from a call tree