Skip to main content

juncture_telemetry/
trace_store.rs

1//! Storage trait and query types for telemetry data.
2//!
3//! Defines the `TraceStore` async trait that abstracts over different
4//! storage backends (`SQLite`, `PostgreSQL`, memory). All operations are
5//! async to support non-blocking batch writes.
6
7use crate::models::{EnrichedSession, Id, ModelStats, Observation, Session, SummaryStats, Trace};
8use chrono::{DateTime, Utc};
9use serde::{Deserialize, Serialize};
10
11/// Error type for trace store operations.
12#[derive(Debug, thiserror::Error)]
13pub enum StoreError {
14    /// Database connection or query error.
15    #[error("storage error: {0}")]
16    Storage(String),
17    /// Serialization/deserialization error.
18    #[error("serialization error: {0}")]
19    Serialization(#[from] serde_json::Error),
20    /// Record not found.
21    #[error("not found: {0}")]
22    NotFound(String),
23}
24
25/// Query parameters for filtering traces.
26#[derive(Clone, Debug, Default, Serialize, Deserialize)]
27#[serde(rename_all = "camelCase")]
28pub struct TraceQuery {
29    /// Filter by session (thread) identifier.
30    pub session_id: Option<String>,
31    /// Filter by user identifier.
32    pub user_id: Option<String>,
33    /// Filter by name (graph name).
34    pub name: Option<String>,
35    /// Filter by environment.
36    pub environment: Option<String>,
37    /// Filter by tags (AND logic).
38    pub tags: Vec<String>,
39    /// Start of time range (inclusive).
40    pub from_timestamp: Option<DateTime<Utc>>,
41    /// End of time range (inclusive).
42    pub to_timestamp: Option<DateTime<Utc>>,
43    /// Page number (0-indexed).
44    pub page: Option<u32>,
45    /// Page size (default 50).
46    pub page_size: Option<u32>,
47}
48
49/// Paginated response wrapper.
50#[derive(Clone, Debug, Serialize, Deserialize)]
51#[serde(rename_all = "camelCase")]
52pub struct PaginatedResponse<T> {
53    /// Items in this page.
54    pub data: Vec<T>,
55    /// Current page number (0-indexed).
56    pub page: u32,
57    /// Items per page.
58    pub page_size: u32,
59    /// Total number of items matching the query.
60    pub total_count: u64,
61}
62
63/// Daily aggregated statistics.
64#[derive(Clone, Debug, Serialize, Deserialize)]
65#[serde(rename_all = "camelCase")]
66pub struct DailyStats {
67    /// Date (YYYY-MM-DD).
68    pub date: String,
69    /// Number of traces.
70    pub trace_count: u64,
71    /// Total observations.
72    pub observation_count: u64,
73    /// Total tokens consumed.
74    pub total_tokens: u64,
75    /// Total cost in USD.
76    pub total_cost: f64,
77    /// Total duration in milliseconds.
78    pub total_duration_ms: u64,
79}
80
81/// Async trait for telemetry data storage.
82///
83/// Implementations must be `Send + Sync` to support concurrent
84/// access from the batch writer and web server.
85#[async_trait::async_trait]
86pub trait TraceStore: Send + Sync + 'static {
87    /// Insert or update a trace.
88    async fn upsert_trace(&self, trace: &Trace) -> Result<(), StoreError>;
89
90    /// Insert an observation.
91    async fn insert_observation(&self, observation: &Observation) -> Result<(), StoreError>;
92
93    /// Insert or update a session.
94    async fn upsert_session(&self, session: &Session) -> Result<(), StoreError>;
95
96    /// Get a trace by ID with its observations.
97    async fn get_trace(&self, id: Id) -> Result<Option<TraceWithObservations>, StoreError>;
98
99    /// Query traces with filtering and pagination.
100    async fn query_traces(
101        &self,
102        query: &TraceQuery,
103    ) -> Result<PaginatedResponse<Trace>, StoreError>;
104
105    /// Get a session by ID.
106    async fn get_session(&self, id: &str) -> Result<Option<Session>, StoreError>;
107
108    /// Query sessions with pagination.
109    async fn query_sessions(
110        &self,
111        page: u32,
112        page_size: u32,
113    ) -> Result<PaginatedResponse<Session>, StoreError>;
114
115    /// Get daily aggregated statistics.
116    async fn get_daily_stats(
117        &self,
118        from: DateTime<Utc>,
119        to: DateTime<Utc>,
120    ) -> Result<Vec<DailyStats>, StoreError>;
121
122    /// Get per-model aggregated statistics.
123    async fn get_model_stats(&self) -> Result<Vec<ModelStats>, StoreError>;
124
125    /// Get overall summary statistics with latency percentiles.
126    async fn get_summary_stats(&self) -> Result<SummaryStats, StoreError>;
127
128    /// Query enriched sessions with aggregated data.
129    async fn query_enriched_sessions(
130        &self,
131        page: u32,
132        page_size: u32,
133    ) -> Result<PaginatedResponse<EnrichedSession>, StoreError>;
134
135    /// Flush any pending writes. Called before process exit.
136    async fn flush(&self) -> Result<(), StoreError>;
137}
138
139/// A trace with its associated observations.
140#[derive(Clone, Debug, Serialize, Deserialize)]
141#[serde(rename_all = "camelCase")]
142pub struct TraceWithObservations {
143    /// The trace.
144    pub trace: Trace,
145    /// Observations belonging to this trace, ordered by `start_time`.
146    pub observations: Vec<Observation>,
147}