From ade0bd59585748176d917a000f9b3d3e19db2e3f Mon Sep 17 00:00:00 2001 From: Ankur Goyal Date: Tue, 31 Mar 2026 22:49:36 -0700 Subject: [PATCH] Support project-name full-span parents --- src/log_queue/queue.rs | 62 +++++++++++++++++++++++++++++++------- src/span.rs | 68 +++++++++++++++++++++++++----------------- src/span_components.rs | 36 ++++++++++++++++------ src/types.rs | 11 +++++-- 4 files changed, 127 insertions(+), 50 deletions(-) diff --git a/src/log_queue/queue.rs b/src/log_queue/queue.rs index 56d0d49..23e874d 100644 --- a/src/log_queue/queue.rs +++ b/src/log_queue/queue.rs @@ -467,15 +467,52 @@ impl LogQueueCore { Some(ParentSpanInfo::FullSpan { object_type, object_id, + compute_object_metadata_args, span_id: parent_span_id, root_span_id: parent_root_span_id, propagated_event: _, }) => { let span_parents = Some(vec![parent_span_id]); let dest = match object_type { - SpanObjectType::Experiment => LogDestination::experiment(object_id), - SpanObjectType::ProjectLogs => LogDestination::project_logs(object_id), - SpanObjectType::PlaygroundLogs => LogDestination::playground_logs(object_id), + SpanObjectType::Experiment => { + LogDestination::experiment(object_id.ok_or_else(|| { + anyhow::anyhow!("experiment parent span is missing object_id") + })?) + } + SpanObjectType::ProjectLogs => { + if let Some(object_id) = object_id { + LogDestination::project_logs(object_id) + } else { + let args = compute_object_metadata_args.as_ref().ok_or_else(|| { + anyhow::anyhow!( + "project-log parent span is missing compute_object_metadata_args" + ) + })?; + if let Some(project_id) = args.get("project_id").and_then(Value::as_str) + { + LogDestination::project_logs(project_id.to_string()) + } else { + let project_name = args + .get("project_name") + .and_then(Value::as_str) + .ok_or_else(|| anyhow::anyhow!("missing project_name"))?; + let project_id = self + .ensure_project_id( + token, + &org_id, + org_name.as_deref(), + project_name, + ) + .await?; + LogDestination::project_logs(project_id) + } + } + } + SpanObjectType::PlaygroundLogs => { + LogDestination::playground_logs(object_id.ok_or_else(|| { + anyhow::anyhow!("playground parent span is missing object_id") + })?) + } }; (parent_root_span_id, span_parents, dest) } @@ -740,13 +777,18 @@ impl LogQueueCore { object_id, .. }) => match object_type { - SpanObjectType::Experiment => LogDestination::experiment(object_id.clone()), - SpanObjectType::ProjectLogs => { - LogDestination::project_logs(object_id.clone()) - } - SpanObjectType::PlaygroundLogs => { - LogDestination::playground_logs(object_id.clone()) - } + SpanObjectType::Experiment => match object_id.clone() { + Some(object_id) => LogDestination::experiment(object_id), + None => return, + }, + SpanObjectType::ProjectLogs => match object_id.clone() { + Some(object_id) => LogDestination::project_logs(object_id), + None => return, + }, + SpanObjectType::PlaygroundLogs => match object_id.clone() { + Some(object_id) => LogDestination::playground_logs(object_id), + None => return, + }, }, None => return, }, diff --git a/src/span.rs b/src/span.rs index 440c6b4..45bc5e8 100644 --- a/src/span.rs +++ b/src/span.rs @@ -483,24 +483,32 @@ impl SpanHandle { // Determine object_type and object_id from parent_info if available, // otherwise default to ProjectLogs - let (object_type, object_id) = match &self.parent_info { - Some(ParentSpanInfo::Experiment { object_id }) => { - (SpanObjectType::Experiment, Some(object_id.clone())) - } - Some(ParentSpanInfo::ProjectLogs { object_id }) => { - (SpanObjectType::ProjectLogs, Some(object_id.clone())) - } - Some(ParentSpanInfo::PlaygroundLogs { object_id }) => { - (SpanObjectType::PlaygroundLogs, Some(object_id.clone())) - } - Some(ParentSpanInfo::FullSpan { - object_type, - object_id, - .. - }) => (*object_type, Some(object_id.clone())), - // Default to ProjectLogs if no parent - _ => (SpanObjectType::ProjectLogs, None), - }; + let (object_type, object_id, inherited_compute_object_metadata_args) = + match &self.parent_info { + Some(ParentSpanInfo::Experiment { object_id }) => { + (SpanObjectType::Experiment, Some(object_id.clone()), None) + } + Some(ParentSpanInfo::ProjectLogs { object_id }) => { + (SpanObjectType::ProjectLogs, Some(object_id.clone()), None) + } + Some(ParentSpanInfo::PlaygroundLogs { object_id }) => ( + SpanObjectType::PlaygroundLogs, + Some(object_id.clone()), + None, + ), + Some(ParentSpanInfo::FullSpan { + object_type, + object_id, + compute_object_metadata_args, + .. + }) => ( + *object_type, + object_id.clone(), + compute_object_metadata_args.clone(), + ), + // Default to ProjectLogs if no parent + _ => (SpanObjectType::ProjectLogs, None, None), + }; // Use root_span_id from parent_info (FullSpan) if available, otherwise // fall back to this span's own span_id (it is the root). @@ -510,16 +518,18 @@ impl SpanHandle { }; let compute_object_metadata_args = if object_id.is_none() { - inner.project_name.as_ref().map(|project_name| { - let mut args = Map::new(); - args.insert( - "project_name".to_string(), - Value::String(project_name.clone()), - ); - args + inherited_compute_object_metadata_args.or_else(|| { + inner.project_name.as_ref().map(|project_name| { + let mut args = Map::new(); + args.insert( + "project_name".to_string(), + Value::String(project_name.clone()), + ); + args + }) }) } else { - None + inherited_compute_object_metadata_args }; Ok(SpanComponents { @@ -865,9 +875,10 @@ mod tests { let parent_info = ParentSpanInfo::FullSpan { object_type: SpanObjectType::ProjectLogs, - object_id: "project-123".to_string(), + object_id: Some("project-123".to_string()), span_id: "parent-span-id".to_string(), root_span_id: "root-span-id".to_string(), + compute_object_metadata_args: None, propagated_event: Some(parent_propagated), }; @@ -913,9 +924,10 @@ mod tests { let parent_info = ParentSpanInfo::FullSpan { object_type: SpanObjectType::Experiment, - object_id: "exp-123".to_string(), + object_id: Some("exp-123".to_string()), span_id: "span-456".to_string(), root_span_id: "root-789".to_string(), + compute_object_metadata_args: None, propagated_event: Some(propagated), }; diff --git a/src/span_components.rs b/src/span_components.rs index 719ae89..6b764ec 100644 --- a/src/span_components.rs +++ b/src/span_components.rs @@ -373,10 +373,6 @@ impl SpanComponents { /// Convert SpanComponents to ParentSpanInfo for creating child spans pub fn to_parent_span_info(&self) -> Result { - // For FullSpan variant, we need object_id, span_id, and root_span_id - let object_id = self.object_id.clone().ok_or_else(|| { - BraintrustError::InvalidConfig("object_id required for parent span".to_string()) - })?; let span_id = self.span_id.clone().ok_or_else(|| { BraintrustError::InvalidConfig("span_id required for parent span".to_string()) })?; @@ -384,9 +380,28 @@ impl SpanComponents { BraintrustError::InvalidConfig("root_span_id required for parent span".to_string()) })?; + match self.object_type { + SpanObjectType::ProjectLogs => { + if self.object_id.is_none() && self.compute_object_metadata_args.is_none() { + return Err(BraintrustError::InvalidConfig( + "project-log parent span requires object_id or compute_object_metadata_args" + .to_string(), + )); + } + } + SpanObjectType::Experiment | SpanObjectType::PlaygroundLogs => { + if self.object_id.is_none() { + return Err(BraintrustError::InvalidConfig( + "object_id required for parent span".to_string(), + )); + } + } + } + Ok(ParentSpanInfo::FullSpan { object_type: self.object_type, - object_id, + object_id: self.object_id.clone(), + compute_object_metadata_args: self.compute_object_metadata_args.clone(), span_id, root_span_id, propagated_event: self.propagated_event.clone(), @@ -399,13 +414,14 @@ impl SpanComponents { ParentSpanInfo::FullSpan { object_type, object_id, + compute_object_metadata_args, span_id, root_span_id, propagated_event, } => Some(Self { object_type: *object_type, - object_id: Some(object_id.clone()), - compute_object_metadata_args: None, + object_id: object_id.clone(), + compute_object_metadata_args: compute_object_metadata_args.clone(), row_id: None, span_id: Some(span_id.clone()), root_span_id: Some(root_span_id.clone()), @@ -585,9 +601,10 @@ mod tests { span_id, root_span_id, propagated_event, + .. } => { assert_eq!(object_type, SpanObjectType::ProjectLogs); - assert_eq!(object_id, "project-123"); + assert_eq!(object_id, Some("project-123".to_string())); assert_eq!(span_id, "span-456"); assert_eq!(root_span_id, "root-789"); assert!(propagated_event.is_some()); @@ -608,9 +625,10 @@ mod tests { let parent = ParentSpanInfo::FullSpan { object_type: SpanObjectType::Experiment, - object_id: "exp-123".to_string(), + object_id: Some("exp-123".to_string()), span_id: "span-456".to_string(), root_span_id: "root-789".to_string(), + compute_object_metadata_args: None, propagated_event: Some(propagated), }; diff --git a/src/types.rs b/src/types.rs index d9cb829..956a873 100644 --- a/src/types.rs +++ b/src/types.rs @@ -357,7 +357,10 @@ pub enum ParentSpanInfo { }, FullSpan { object_type: SpanObjectType, - object_id: String, + #[serde(skip_serializing_if = "Option::is_none")] + object_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + compute_object_metadata_args: Option>, span_id: String, root_span_id: String, #[serde(skip_serializing_if = "Option::is_none")] @@ -869,9 +872,10 @@ mod tests { fn parent_span_info_full_span_serializes_object_type_as_u8() { let parent = ParentSpanInfo::FullSpan { object_type: SpanObjectType::Experiment, - object_id: "exp-123".to_string(), + object_id: Some("exp-123".to_string()), span_id: "span-1".to_string(), root_span_id: "root-1".to_string(), + compute_object_metadata_args: None, propagated_event: None, }; @@ -908,9 +912,10 @@ mod tests { fn parent_span_info_serializes_propagated_event() { let parent = ParentSpanInfo::FullSpan { object_type: SpanObjectType::ProjectLogs, - object_id: "proj-123".to_string(), + object_id: Some("proj-123".to_string()), span_id: "span-1".to_string(), root_span_id: "root-1".to_string(), + compute_object_metadata_args: None, propagated_event: Some(Map::from_iter([( "metrics".to_string(), json!({ "foo": 0.1 }),