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