1use crate::error::TypeError;
2use crate::json_to_pyobject_value;
3use crate::trace::{Attribute, SpanEvent, SpanLink};
4use crate::PyHelperFuncs;
5use crate::TraceCursor;
6use chrono::{DateTime, Utc};
7use pyo3::prelude::*;
8use pyo3::IntoPyObjectExt;
9use serde::{Deserialize, Serialize};
10use serde_json::Value;
11#[cfg(feature = "server")]
12use sqlx::{postgres::PgRow, FromRow, Row};
13
14#[derive(Debug, Serialize, Deserialize, Clone)]
15#[pyclass]
16pub struct TraceListItem {
17 #[pyo3(get)]
18 pub trace_id: String,
19 #[pyo3(get)]
20 pub service_name: String,
21 #[pyo3(get)]
22 pub scope_name: String,
23 #[pyo3(get)]
24 pub scope_version: String,
25 #[pyo3(get)]
26 pub root_operation: String,
27 #[pyo3(get)]
28 pub start_time: DateTime<Utc>,
29 #[pyo3(get)]
30 pub end_time: Option<DateTime<Utc>>,
31 #[pyo3(get)]
32 pub duration_ms: Option<i64>,
33 #[pyo3(get)]
34 pub status_code: i32,
35 #[pyo3(get)]
36 pub status_message: Option<String>,
37 #[pyo3(get)]
38 pub span_count: i64,
39 #[pyo3(get)]
40 pub has_errors: bool,
41 #[pyo3(get)]
42 pub error_count: i64,
43 #[pyo3(get)]
44 pub resource_attributes: Vec<Attribute>,
45}
46
47#[cfg(feature = "server")]
48impl FromRow<'_, PgRow> for TraceListItem {
49 fn from_row(row: &PgRow) -> Result<Self, sqlx::Error> {
50 let resource_attributes: Vec<Attribute> =
51 serde_json::from_value(row.try_get("resource_attributes")?).unwrap_or_default();
52 Ok(TraceListItem {
53 trace_id: row.try_get("trace_id")?,
54 service_name: row.try_get("service_name")?,
55 scope_name: row.try_get("scope_name")?,
56 scope_version: row.try_get("scope_version")?,
57 root_operation: row.try_get("root_operation")?,
58 start_time: row.try_get("start_time")?,
59 end_time: row.try_get("end_time")?,
60 duration_ms: row.try_get("duration_ms")?,
61 status_code: row.try_get("status_code")?,
62 status_message: row.try_get("status_message")?,
63 span_count: row.try_get("span_count")?,
64 has_errors: row.try_get("has_errors")?,
65 error_count: row.try_get("error_count")?,
66 resource_attributes,
67 })
68 }
69}
70
71#[pymethods]
72impl TraceListItem {
73 pub fn __str__(&self) -> String {
74 PyHelperFuncs::__str__(self)
75 }
76}
77
78#[derive(Debug, Serialize, Deserialize, Clone)]
79#[pyclass]
80pub struct TraceSpan {
81 #[pyo3(get)]
82 pub trace_id: String,
83 #[pyo3(get)]
84 pub span_id: String,
85 #[pyo3(get)]
86 pub parent_span_id: Option<String>,
87 #[pyo3(get)]
88 pub span_name: String,
89 #[pyo3(get)]
90 pub span_kind: Option<String>,
91 #[pyo3(get)]
92 pub start_time: DateTime<Utc>,
93 #[pyo3(get)]
94 pub end_time: DateTime<Utc>,
95 #[pyo3(get)]
96 pub duration_ms: i64,
97 #[pyo3(get)]
98 pub status_code: i32,
99 #[pyo3(get)]
100 pub status_message: Option<String>,
101 #[pyo3(get)]
102 pub attributes: Vec<Attribute>,
103 #[pyo3(get)]
104 pub events: Vec<SpanEvent>,
105 #[pyo3(get)]
106 pub links: Vec<SpanLink>,
107 #[pyo3(get)]
108 pub depth: i32,
109 #[pyo3(get)]
110 pub path: Vec<String>,
111 #[pyo3(get)]
112 pub root_span_id: String,
113 #[pyo3(get)]
114 pub service_name: String,
115 #[pyo3(get)]
116 pub span_order: i32,
117 pub input: Option<Value>,
118 pub output: Option<Value>,
119}
120
121#[pymethods]
122impl TraceSpan {
123 pub fn __str__(&self) -> String {
124 PyHelperFuncs::__str__(self)
125 }
126
127 #[getter]
128 pub fn input<'py>(&self, py: Python<'py>) -> Result<Bound<'py, PyAny>, TypeError> {
129 let input = match &self.input {
130 Some(v) => v,
131 None => &Value::Null,
132 };
133 Ok(json_to_pyobject_value(py, input)?
134 .into_bound_py_any(py)?
135 .clone())
136 }
137
138 #[getter]
139 pub fn output<'py>(&self, py: Python<'py>) -> Result<Bound<'py, PyAny>, TypeError> {
140 let output = match &self.output {
141 Some(v) => v,
142 None => &Value::Null,
143 };
144 Ok(json_to_pyobject_value(py, output)?
145 .into_bound_py_any(py)?
146 .clone())
147 }
148}
149
150#[cfg(feature = "server")]
151impl FromRow<'_, PgRow> for TraceSpan {
152 fn from_row(row: &PgRow) -> Result<Self, sqlx::Error> {
153 let attributes: Vec<Attribute> =
154 serde_json::from_value(row.try_get("attributes")?).unwrap_or_default();
155 let events: Vec<SpanEvent> =
156 serde_json::from_value(row.try_get("events")?).unwrap_or_default();
157 let links: Vec<SpanLink> =
158 serde_json::from_value(row.try_get("links")?).unwrap_or_default();
159 let input: Option<Value> = row.try_get("input")?;
160 let output: Option<Value> = row.try_get("output")?;
161
162 Ok(TraceSpan {
163 trace_id: row.try_get("trace_id")?,
164 span_id: row.try_get("span_id")?,
165 parent_span_id: row.try_get("parent_span_id")?,
166 span_name: row.try_get("span_name")?,
167 span_kind: row.try_get("span_kind")?,
168 start_time: row.try_get("start_time")?,
169 end_time: row.try_get("end_time")?,
170 duration_ms: row.try_get("duration_ms")?,
171 status_code: row.try_get("status_code")?,
172 status_message: row.try_get("status_message")?,
173 attributes,
174 events,
175 links,
176 depth: row.try_get("depth")?,
177 path: row.try_get("path")?,
178 root_span_id: row.try_get("root_span_id")?,
179 span_order: row.try_get("span_order")?,
180 input,
181 output,
182 service_name: row.try_get("service_name")?,
183 })
184 }
185}
186
187#[derive(Debug, Default, Clone, Serialize, Deserialize)]
188#[pyclass]
189pub struct TraceFilters {
190 #[pyo3(get, set)]
191 pub service_name: Option<String>,
192 #[pyo3(get, set)]
193 pub has_errors: Option<bool>,
194 #[pyo3(get, set)]
195 pub status_code: Option<i32>,
196 #[pyo3(get, set)]
197 pub start_time: Option<DateTime<Utc>>,
198 #[pyo3(get, set)]
199 pub end_time: Option<DateTime<Utc>>,
200 #[pyo3(get, set)]
201 pub limit: Option<i32>,
202 #[pyo3(get, set)]
203 pub cursor_start_time: Option<DateTime<Utc>>,
204 #[pyo3(get, set)]
205 pub cursor_trace_id: Option<String>,
206 #[pyo3(get, set)]
207 pub direction: Option<String>,
208 #[pyo3(get, set)]
209 pub attribute_filters: Option<Vec<String>>,
210 #[pyo3(get, set)]
211 pub trace_ids: Option<Vec<String>>,
212 #[pyo3(get, set)]
213 pub entity_uid: Option<String>,
214}
215
216#[pymethods]
217#[allow(clippy::too_many_arguments)]
218impl TraceFilters {
219 #[new]
220 #[pyo3(signature = (
221 service_name=None,
222 has_errors=None,
223 status_code=None,
224 start_time=None,
225 end_time=None,
226 limit=None,
227 cursor_start_time=None,
228 cursor_trace_id=None,
229 attribute_filters=None,
230 trace_ids=None,
231 entity_uid=None
232 ))]
233 pub fn new(
234 service_name: Option<String>,
235 has_errors: Option<bool>,
236 status_code: Option<i32>,
237 start_time: Option<DateTime<Utc>>,
238 end_time: Option<DateTime<Utc>>,
239 limit: Option<i32>,
240 cursor_start_time: Option<DateTime<Utc>>,
241 cursor_trace_id: Option<String>,
242 attribute_filters: Option<Vec<String>>,
243 trace_ids: Option<Vec<String>>,
244 entity_uid: Option<String>,
245 ) -> Self {
246 TraceFilters {
247 service_name,
248 has_errors,
249 status_code,
250 start_time,
251 end_time,
252 limit,
253 cursor_start_time,
254 cursor_trace_id,
255 direction: None,
256 attribute_filters,
257 trace_ids,
258 entity_uid,
259 }
260 }
261}
262
263impl TraceFilters {
264 pub fn with_service(mut self, service: impl Into<String>) -> Self {
265 self.service_name = Some(service.into());
266 self
267 }
268
269 pub fn with_errors_only(mut self) -> Self {
270 self.has_errors = Some(true);
271 self
272 }
273
274 pub fn with_time_range(mut self, start: DateTime<Utc>, end: DateTime<Utc>) -> Self {
275 self.start_time = Some(start);
276 self.end_time = Some(end);
277 self
278 }
279
280 pub fn with_cursor(mut self, started_at: DateTime<Utc>, trace_id: impl Into<String>) -> Self {
281 self.cursor_start_time = Some(started_at);
282 self.cursor_trace_id = Some(trace_id.into());
283 self
284 }
285
286 pub fn limit(mut self, limit: i32) -> Self {
287 self.limit = Some(limit);
288 self
289 }
290
291 pub fn next_page(mut self, cursor: &TraceCursor) -> Self {
292 self.cursor_start_time = Some(cursor.start_time);
293 self.cursor_trace_id = Some(cursor.trace_id.clone());
294 self.direction = Some("next".to_string());
295 self
296 }
297
298 pub fn first_page(mut self) -> Self {
299 self.cursor_start_time = None;
300 self.cursor_trace_id = None;
301 self
302 }
303
304 pub fn previous_page(mut self, cursor: &TraceCursor) -> Self {
305 self.cursor_start_time = Some(cursor.start_time);
306 self.cursor_trace_id = Some(cursor.trace_id.clone());
307 self.direction = Some("previous".to_string());
308 self
309 }
310
311 pub fn with_entity_uid(mut self, entity_uid: impl Into<String>) -> Self {
312 self.entity_uid = Some(entity_uid.into());
313 self
314 }
315
316 pub fn parsed_trace_ids(&self) -> Result<Option<Vec<Vec<u8>>>, TypeError> {
320 match &self.trace_ids {
321 None => Ok(None),
322 Some(ids) => {
323 let parsed: Result<Vec<Vec<u8>>, _> = ids
324 .iter()
325 .map(|hex| {
326 let bytes = hex::decode(hex)?;
327
328 if bytes.len() != 16 {
329 return Err(TypeError::InvalidLength(format!(
330 "Invalid trace_id length: expected 16 bytes, got {} for '{}'",
331 bytes.len(),
332 hex
333 )));
334 }
335
336 Ok(bytes)
337 })
338 .collect();
339
340 parsed.map(Some)
341 }
342 }
343 }
344
345 pub fn parsed_cursor_trace_id(&self) -> Result<Option<Vec<u8>>, TypeError> {
347 match &self.cursor_trace_id {
348 None => Ok(None),
349 Some(hex) => {
350 let bytes = hex::decode(hex)?;
351
352 if bytes.len() != 16 {
353 return Err(TypeError::InvalidLength(format!(
354 "Invalid cursor_trace_id length: expected 16 bytes, got {}",
355 bytes.len()
356 )));
357 }
358
359 Ok(Some(bytes))
360 }
361 }
362 }
363}
364
365#[cfg_attr(feature = "server", derive(sqlx::FromRow))]
366#[derive(Debug, Serialize, Deserialize, Clone)]
367#[pyclass]
368pub struct TraceMetricBucket {
369 #[pyo3(get)]
370 pub bucket_start: DateTime<Utc>,
371 #[pyo3(get)]
372 pub trace_count: i64,
373 #[pyo3(get)]
374 pub avg_duration_ms: f64,
375 #[pyo3(get)]
376 pub p50_duration_ms: Option<f64>,
377 #[pyo3(get)]
378 pub p95_duration_ms: Option<f64>,
379 #[pyo3(get)]
380 pub p99_duration_ms: Option<f64>,
381 #[pyo3(get)]
382 pub error_rate: f64,
383}
384
385#[pymethods]
386impl TraceMetricBucket {
387 pub fn __str__(&self) -> String {
388 PyHelperFuncs::__str__(self)
389 }
390}