Skip to main content

platform_core/
telemetry_query.rs

1use crate::error::AppResult;
2use async_trait::async_trait;
3use chrono::{DateTime, Utc};
4use serde::{Deserialize, Serialize};
5use serde_json::Value;
6use std::fmt::Debug;
7use std::sync::Arc;
8
9#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
10pub struct TelemetrySpan {
11    pub id: String,
12    pub name: String,
13    pub status: Option<String>,
14    pub started_at: DateTime<Utc>,
15    pub ended_at: DateTime<Utc>,
16    pub attributes: Value,
17}
18
19#[derive(Debug, Clone, Default, PartialEq, Eq)]
20pub struct TelemetrySpanQuery {
21    pub correlation_id: Option<String>,
22    pub story_id: Option<String>,
23    pub function_run_id: Option<String>,
24    pub outbox_event_id: Option<String>,
25}
26
27impl TelemetrySpanQuery {
28    pub fn by_correlation_id(correlation_id: impl Into<String>) -> Self {
29        Self {
30            correlation_id: Some(correlation_id.into()),
31            ..Self::default()
32        }
33    }
34
35    pub fn by_function_run_id(function_run_id: impl Into<String>) -> Self {
36        Self {
37            function_run_id: Some(function_run_id.into()),
38            ..Self::default()
39        }
40    }
41
42    pub fn by_outbox_event_id(outbox_event_id: impl Into<String>) -> Self {
43        Self {
44            outbox_event_id: Some(outbox_event_id.into()),
45            ..Self::default()
46        }
47    }
48}
49
50#[async_trait]
51pub trait TelemetrySpanProvider: Debug + Send + Sync {
52    async fn query_spans(&self, query: TelemetrySpanQuery) -> AppResult<Vec<TelemetrySpan>>;
53}
54
55#[derive(Debug, Default)]
56pub struct NoopTelemetrySpanProvider;
57
58#[async_trait]
59impl TelemetrySpanProvider for NoopTelemetrySpanProvider {
60    async fn query_spans(&self, _query: TelemetrySpanQuery) -> AppResult<Vec<TelemetrySpan>> {
61        Ok(Vec::new())
62    }
63}
64
65#[derive(Debug, Clone, Default)]
66pub struct InMemoryTelemetrySpanProvider {
67    spans: Arc<Vec<TelemetrySpan>>,
68}
69
70impl InMemoryTelemetrySpanProvider {
71    pub fn new(spans: impl Into<Vec<TelemetrySpan>>) -> Self {
72        Self {
73            spans: Arc::new(spans.into()),
74        }
75    }
76}
77
78#[async_trait]
79impl TelemetrySpanProvider for InMemoryTelemetrySpanProvider {
80    async fn query_spans(&self, query: TelemetrySpanQuery) -> AppResult<Vec<TelemetrySpan>> {
81        Ok(self
82            .spans
83            .iter()
84            .filter(|span| span_matches_query(span, &query))
85            .cloned()
86            .collect())
87    }
88}
89
90fn span_matches_query(span: &TelemetrySpan, query: &TelemetrySpanQuery) -> bool {
91    let selectors = [
92        query
93            .correlation_id
94            .as_deref()
95            .map(|value| ("lenso.correlation_id", value)),
96        query
97            .story_id
98            .as_deref()
99            .map(|value| ("lenso.story_id", value)),
100        query
101            .function_run_id
102            .as_deref()
103            .map(|value| ("lenso.function_run_id", value)),
104        query
105            .outbox_event_id
106            .as_deref()
107            .map(|value| ("lenso.outbox_event_id", value)),
108    ];
109
110    let selected = selectors.into_iter().flatten().collect::<Vec<_>>();
111    if selected.is_empty() {
112        return false;
113    }
114
115    selected
116        .iter()
117        .all(|(key, expected)| span_attribute(span, key) == Some(*expected))
118}
119
120fn span_attribute<'a>(span: &'a TelemetrySpan, key: &str) -> Option<&'a str> {
121    span.attributes.get(key).and_then(Value::as_str)
122}
123
124#[cfg(test)]
125mod tests {
126    use super::*;
127
128    #[tokio::test]
129    async fn in_memory_provider_queries_by_correlation_id() {
130        let provider = InMemoryTelemetrySpanProvider::new([
131            test_span(
132                "span_a",
133                serde_json::json!({ "lenso.correlation_id": "corr_a" }),
134            ),
135            test_span(
136                "span_b",
137                serde_json::json!({ "lenso.correlation_id": "corr_b" }),
138            ),
139        ]);
140
141        let spans = provider
142            .query_spans(TelemetrySpanQuery::by_correlation_id("corr_a"))
143            .await
144            .expect("query should succeed");
145
146        assert_eq!(spans.len(), 1);
147        assert_eq!(spans[0].id, "span_a");
148    }
149
150    #[tokio::test]
151    async fn in_memory_provider_queries_by_function_run_id() {
152        let provider = InMemoryTelemetrySpanProvider::new([
153            test_span(
154                "span_a",
155                serde_json::json!({ "lenso.function_run_id": "fnrun_a" }),
156            ),
157            test_span(
158                "span_b",
159                serde_json::json!({ "lenso.outbox_event_id": "evt_b" }),
160            ),
161        ]);
162
163        let spans = provider
164            .query_spans(TelemetrySpanQuery::by_function_run_id("fnrun_a"))
165            .await
166            .expect("query should succeed");
167
168        assert_eq!(spans.len(), 1);
169        assert_eq!(spans[0].id, "span_a");
170    }
171
172    fn test_span(id: &str, attributes: Value) -> TelemetrySpan {
173        TelemetrySpan {
174            attributes,
175            ended_at: "2026-05-31T00:00:01Z"
176                .parse()
177                .expect("timestamp should parse"),
178            id: id.to_owned(),
179            name: id.to_owned(),
180            started_at: "2026-05-31T00:00:00Z"
181                .parse()
182                .expect("timestamp should parse"),
183            status: Some("ok".to_owned()),
184        }
185    }
186}