Skip to main content

scouter_types/trace/
sql.rs

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    /// Convert hex trace_ids directly to byte vectors for PostgreSQL BYTEA binding
317    /// Returns None if no trace_ids are set
318    /// Returns error if any trace_id is invalid hex or wrong length
319    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    /// Convert hex cursor_trace_id directly to bytes for PostgreSQL BYTEA binding
346    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}