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: String,
23 #[pyo3(get)]
24 pub root_operation: String,
25 #[pyo3(get)]
26 pub start_time: DateTime<Utc>,
27 #[pyo3(get)]
28 pub end_time: Option<DateTime<Utc>>,
29 #[pyo3(get)]
30 pub duration_ms: Option<i64>,
31 #[pyo3(get)]
32 pub status_code: i32,
33 #[pyo3(get)]
34 pub status_message: Option<String>,
35 #[pyo3(get)]
36 pub span_count: i64,
37 #[pyo3(get)]
38 pub has_errors: bool,
39 #[pyo3(get)]
40 pub error_count: i64,
41 #[pyo3(get)]
42 pub resource_attributes: Vec<Attribute>,
43}
44
45#[cfg(feature = "server")]
46impl FromRow<'_, PgRow> for TraceListItem {
47 fn from_row(row: &PgRow) -> Result<Self, sqlx::Error> {
48 let resource_attributes: Vec<Attribute> =
49 serde_json::from_value(row.try_get("resource_attributes")?).unwrap_or_default();
50 Ok(TraceListItem {
51 trace_id: row.try_get("trace_id")?,
52 service_name: row.try_get("service_name")?,
53 scope: row.try_get("scope")?,
54 root_operation: row.try_get("root_operation")?,
55 start_time: row.try_get("start_time")?,
56 end_time: row.try_get("end_time")?,
57 duration_ms: row.try_get("duration_ms")?,
58 status_code: row.try_get("status_code")?,
59 status_message: row.try_get("status_message")?,
60 span_count: row.try_get("span_count")?,
61 has_errors: row.try_get("has_errors")?,
62 error_count: row.try_get("error_count")?,
63 resource_attributes,
64 })
65 }
66}
67
68#[pymethods]
69impl TraceListItem {
70 pub fn __str__(&self) -> String {
71 PyHelperFuncs::__str__(self)
72 }
73}
74
75#[derive(Debug, Serialize, Deserialize, Clone)]
76#[pyclass]
77pub struct TraceSpan {
78 #[pyo3(get)]
79 pub trace_id: String,
80 #[pyo3(get)]
81 pub span_id: String,
82 #[pyo3(get)]
83 pub parent_span_id: Option<String>,
84 #[pyo3(get)]
85 pub span_name: String,
86 #[pyo3(get)]
87 pub span_kind: Option<String>,
88 #[pyo3(get)]
89 pub start_time: DateTime<Utc>,
90 #[pyo3(get)]
91 pub end_time: Option<DateTime<Utc>>,
92 #[pyo3(get)]
93 pub duration_ms: Option<i64>,
94 #[pyo3(get)]
95 pub status_code: i32,
96 #[pyo3(get)]
97 pub status_message: Option<String>,
98 #[pyo3(get)]
99 pub attributes: Vec<Attribute>,
100 #[pyo3(get)]
101 pub events: Vec<SpanEvent>,
102 #[pyo3(get)]
103 pub links: Vec<SpanLink>,
104 #[pyo3(get)]
105 pub depth: i32,
106 #[pyo3(get)]
107 pub path: Vec<String>,
108 #[pyo3(get)]
109 pub root_span_id: String,
110 #[pyo3(get)]
111 pub service_name: String,
112 #[pyo3(get)]
113 pub span_order: i32,
114 pub input: Option<Value>,
115 pub output: Option<Value>,
116}
117
118#[pymethods]
119impl TraceSpan {
120 pub fn __str__(&self) -> String {
121 PyHelperFuncs::__str__(self)
122 }
123
124 #[getter]
125 pub fn input<'py>(&self, py: Python<'py>) -> Result<Bound<'py, PyAny>, TypeError> {
126 let input = match &self.input {
127 Some(v) => v,
128 None => &Value::Null,
129 };
130 Ok(json_to_pyobject_value(py, input)?
131 .into_bound_py_any(py)?
132 .clone())
133 }
134
135 #[getter]
136 pub fn output<'py>(&self, py: Python<'py>) -> Result<Bound<'py, PyAny>, TypeError> {
137 let output = match &self.output {
138 Some(v) => v,
139 None => &Value::Null,
140 };
141 Ok(json_to_pyobject_value(py, output)?
142 .into_bound_py_any(py)?
143 .clone())
144 }
145}
146
147#[cfg(feature = "server")]
148impl FromRow<'_, PgRow> for TraceSpan {
149 fn from_row(row: &PgRow) -> Result<Self, sqlx::Error> {
150 let attributes: Vec<Attribute> =
151 serde_json::from_value(row.try_get("attributes")?).unwrap_or_default();
152 let events: Vec<SpanEvent> =
153 serde_json::from_value(row.try_get("events")?).unwrap_or_default();
154 let links: Vec<SpanLink> =
155 serde_json::from_value(row.try_get("links")?).unwrap_or_default();
156 let input: Option<Value> = row.try_get("input")?;
157 let output: Option<Value> = row.try_get("output")?;
158
159 Ok(TraceSpan {
160 trace_id: row.try_get("trace_id")?,
161 span_id: row.try_get("span_id")?,
162 parent_span_id: row.try_get("parent_span_id")?,
163 span_name: row.try_get("span_name")?,
164 span_kind: row.try_get("span_kind")?,
165 start_time: row.try_get("start_time")?,
166 end_time: row.try_get("end_time")?,
167 duration_ms: row.try_get("duration_ms")?,
168 status_code: row.try_get("status_code")?,
169 status_message: row.try_get("status_message")?,
170 attributes,
171 events,
172 links,
173 depth: row.try_get("depth")?,
174 path: row.try_get("path")?,
175 root_span_id: row.try_get("root_span_id")?,
176 span_order: row.try_get("span_order")?,
177 input,
178 output,
179 service_name: row.try_get("service_name")?,
180 })
181 }
182}
183
184#[derive(Debug, Default, Clone, Serialize, Deserialize)]
185#[pyclass]
186pub struct TraceFilters {
187 #[pyo3(get, set)]
188 pub service_name: Option<String>,
189 #[pyo3(get, set)]
190 pub has_errors: Option<bool>,
191 #[pyo3(get, set)]
192 pub status_code: Option<i32>,
193 #[pyo3(get, set)]
194 pub start_time: Option<DateTime<Utc>>,
195 #[pyo3(get, set)]
196 pub end_time: Option<DateTime<Utc>>,
197 #[pyo3(get, set)]
198 pub limit: Option<i32>,
199 #[pyo3(get, set)]
200 pub cursor_start_time: Option<DateTime<Utc>>,
201 #[pyo3(get, set)]
202 pub cursor_trace_id: Option<String>,
203 #[pyo3(get, set)]
204 pub direction: Option<String>,
205 #[pyo3(get, set)]
206 pub attribute_filters: Option<Vec<String>>,
207 #[pyo3(get, set)]
208 pub trace_ids: Option<Vec<String>>,
209}
210
211#[pymethods]
212#[allow(clippy::too_many_arguments)]
213impl TraceFilters {
214 #[new]
215 #[pyo3(signature = (
216 service_name=None,
217 has_errors=None,
218 status_code=None,
219 start_time=None,
220 end_time=None,
221 limit=None,
222 cursor_start_time=None,
223 cursor_trace_id=None,
224 attribute_filters=None,
225 trace_ids=None
226 ))]
227 pub fn new(
228 service_name: Option<String>,
229 has_errors: Option<bool>,
230 status_code: Option<i32>,
231 start_time: Option<DateTime<Utc>>,
232 end_time: Option<DateTime<Utc>>,
233 limit: Option<i32>,
234 cursor_start_time: Option<DateTime<Utc>>,
235 cursor_trace_id: Option<String>,
236 attribute_filters: Option<Vec<String>>,
237 trace_ids: Option<Vec<String>>,
238 ) -> Self {
239 TraceFilters {
240 service_name,
241 has_errors,
242 status_code,
243 start_time,
244 end_time,
245 limit,
246 cursor_start_time,
247 cursor_trace_id,
248 direction: None,
249 attribute_filters,
250 trace_ids,
251 }
252 }
253}
254
255impl TraceFilters {
256 pub fn with_service(mut self, service: impl Into<String>) -> Self {
257 self.service_name = Some(service.into());
258 self
259 }
260
261 pub fn with_errors_only(mut self) -> Self {
262 self.has_errors = Some(true);
263 self
264 }
265
266 pub fn with_time_range(mut self, start: DateTime<Utc>, end: DateTime<Utc>) -> Self {
267 self.start_time = Some(start);
268 self.end_time = Some(end);
269 self
270 }
271
272 pub fn with_cursor(mut self, started_at: DateTime<Utc>, trace_id: impl Into<String>) -> Self {
273 self.cursor_start_time = Some(started_at);
274 self.cursor_trace_id = Some(trace_id.into());
275 self
276 }
277
278 pub fn limit(mut self, limit: i32) -> Self {
279 self.limit = Some(limit);
280 self
281 }
282
283 pub fn next_page(mut self, cursor: &TraceCursor) -> Self {
284 self.cursor_start_time = Some(cursor.start_time);
285 self.cursor_trace_id = Some(cursor.trace_id.clone());
286 self.direction = Some("next".to_string());
287 self
288 }
289
290 pub fn first_page(mut self) -> Self {
291 self.cursor_start_time = None;
292 self.cursor_trace_id = None;
293 self
294 }
295
296 pub fn previous_page(mut self, cursor: &TraceCursor) -> Self {
297 self.cursor_start_time = Some(cursor.start_time);
298 self.cursor_trace_id = Some(cursor.trace_id.clone());
299 self.direction = Some("previous".to_string());
300 self
301 }
302}
303
304#[cfg_attr(feature = "server", derive(sqlx::FromRow))]
305#[derive(Debug, Serialize, Deserialize, Clone)]
306#[pyclass]
307pub struct TraceMetricBucket {
308 #[pyo3(get)]
309 pub bucket_start: DateTime<Utc>,
310 #[pyo3(get)]
311 pub trace_count: i64,
312 #[pyo3(get)]
313 pub avg_duration_ms: f64,
314 #[pyo3(get)]
315 pub p50_duration_ms: Option<f64>,
316 #[pyo3(get)]
317 pub p95_duration_ms: Option<f64>,
318 #[pyo3(get)]
319 pub p99_duration_ms: Option<f64>,
320 #[pyo3(get)]
321 pub error_rate: f64,
322}
323
324#[pymethods]
325impl TraceMetricBucket {
326 pub fn __str__(&self) -> String {
327 PyHelperFuncs::__str__(self)
328 }
329}