use crate::server::{AppState, QueryCache};
use axum::{
extract::{Path, Query, State},
http::StatusCode,
response::IntoResponse,
Json,
};
use otelite_core::api::{ErrorResponse, SpanEntry, TraceDetail, TraceEntry, TracesResponse};
use otelite_core::storage::QueryParams;
use otelite_core::telemetry::Span;
use serde::{Deserialize, Serialize};
#[derive(Debug, Deserialize, Serialize, utoipa::IntoParams)]
pub struct TracesQuery {
#[serde(default)]
pub trace_id: Option<String>,
#[serde(default)]
pub service: Option<String>,
#[serde(default)]
pub resource: Option<String>,
#[serde(default)]
pub search: Option<String>,
#[serde(default)]
pub start_time: Option<i64>,
#[serde(default)]
pub end_time: Option<i64>,
#[serde(default = "default_limit")]
pub limit: usize,
#[serde(default)]
pub offset: usize,
}
fn default_limit() -> usize {
100
}
#[utoipa::path(
get,
path = "/api/traces",
params(TracesQuery),
responses(
(status = 200, description = "List of traces matching query", body = TracesResponse),
(status = 500, description = "Internal server error", body = ErrorResponse)
),
tag = "traces"
)]
pub async fn list_traces(
State(state): State<AppState>,
Query(params): Query<TracesQuery>,
) -> Result<impl IntoResponse, (StatusCode, Json<ErrorResponse>)> {
let cache_key = QueryCache::make_key(¶ms);
if let Some(cached) = state.cache.traces.get(&cache_key) {
return Ok((
StatusCode::OK,
[("content-type", "application/json")],
cached,
)
.into_response());
}
let limit = params.limit.min(1000);
let query = QueryParams {
start_time: params.start_time,
end_time: params.end_time,
limit: Some(limit * 10), trace_id: params.trace_id.clone(),
..Default::default()
};
let spans = state.storage.query_spans(&query).await.map_err(|e| {
(
StatusCode::INTERNAL_SERVER_ERROR,
Json(ErrorResponse::storage_error(format!("query spans: {}", e))),
)
})?;
let mut traces_map: std::collections::HashMap<String, Vec<Span>> =
std::collections::HashMap::new();
for span in spans {
traces_map
.entry(span.trace_id.clone())
.or_default()
.push(span);
}
let mut trace_entries: Vec<TraceEntry> = traces_map
.into_iter()
.map(|(trace_id, spans)| {
let start_time = spans.iter().map(|s| s.start_time).min().unwrap_or(0);
let end_time = spans.iter().map(|s| s.end_time).max().unwrap_or(0);
let duration = end_time - start_time;
let root_span = spans
.iter()
.find(|s| s.parent_span_id.is_none())
.or_else(|| spans.first());
let root_span_name = root_span
.map(|s| s.name.clone())
.unwrap_or_else(|| "Unknown".to_string());
let service_names: Vec<String> = {
let mut names: Vec<String> = spans
.iter()
.filter_map(|s| s.resource.as_ref())
.filter_map(|r| r.attributes.get("service.name"))
.cloned()
.collect::<std::collections::HashSet<_>>()
.into_iter()
.collect();
names.sort();
names
};
let has_errors = spans.iter().any(|s| {
matches!(
s.status.code,
otelite_core::telemetry::trace::StatusCode::Error
)
});
TraceEntry {
trace_id,
root_span_name,
start_time,
duration,
span_count: spans.len(),
service_names,
has_errors,
}
})
.collect();
trace_entries.sort_by_key(|b| std::cmp::Reverse(b.start_time));
let total = trace_entries.len();
let paginated_traces: Vec<TraceEntry> = trace_entries
.into_iter()
.skip(params.offset)
.take(limit)
.collect();
let response = TracesResponse {
traces: paginated_traces,
total,
limit,
offset: params.offset,
};
if let Ok(json) = serde_json::to_string(&response) {
state.cache.traces.insert(cache_key, json.clone());
Ok((StatusCode::OK, [("content-type", "application/json")], json).into_response())
} else {
Ok(Json(response).into_response())
}
}
#[utoipa::path(
get,
path = "/api/traces/{trace_id}",
params(
("trace_id" = String, Path, description = "Trace ID")
),
responses(
(status = 200, description = "Trace details with all spans", body = TraceDetail),
(status = 404, description = "Trace not found", body = ErrorResponse),
(status = 500, description = "Internal server error", body = ErrorResponse)
),
tag = "traces"
)]
pub async fn get_trace(
State(state): State<AppState>,
Path(trace_id): Path<String>,
) -> Result<impl IntoResponse, (StatusCode, Json<ErrorResponse>)> {
let query = QueryParams {
trace_id: Some(trace_id.clone()),
limit: Some(1000), ..Default::default()
};
let spans = state.storage.query_spans(&query).await.map_err(|e| {
(
StatusCode::INTERNAL_SERVER_ERROR,
Json(ErrorResponse::storage_error(format!(
"query trace spans: {}",
e
))),
)
})?;
if spans.is_empty() {
return Err((
StatusCode::NOT_FOUND,
Json(ErrorResponse::not_found(format!("Trace '{}'", trace_id))),
));
}
let start_time = spans.iter().map(|s| s.start_time).min().unwrap_or(0);
let end_time = spans.iter().map(|s| s.end_time).max().unwrap_or(0);
let duration = end_time - start_time;
let service_names: Vec<String> = {
let mut names: Vec<String> = spans
.iter()
.filter_map(|s| s.resource.as_ref())
.filter_map(|r| r.attributes.get("service.name"))
.cloned()
.collect::<std::collections::HashSet<_>>()
.into_iter()
.collect();
names.sort();
names
};
let span_entries: Vec<SpanEntry> = spans.into_iter().map(SpanEntry::from).collect();
let span_count = span_entries.len();
let trace_detail = TraceDetail {
trace_id,
spans: span_entries,
start_time,
end_time,
duration,
span_count,
service_names,
};
Ok(Json(trace_detail))
}
#[derive(Debug, Deserialize, utoipa::IntoParams)]
pub struct ExportQuery {
#[serde(default = "default_format")]
pub format: String,
#[serde(flatten)]
pub filters: TracesQuery,
}
fn default_format() -> String {
"json".to_string()
}
#[utoipa::path(
get,
path = "/api/traces/export",
params(ExportQuery),
responses(
(status = 200, description = "Exported traces in JSON format"),
(status = 400, description = "Invalid format parameter", body = ErrorResponse),
(status = 500, description = "Internal server error", body = ErrorResponse)
),
tag = "traces"
)]
pub async fn export_traces(
State(state): State<AppState>,
Query(params): Query<ExportQuery>,
) -> Result<impl IntoResponse, (StatusCode, Json<ErrorResponse>)> {
let query = QueryParams {
start_time: params.filters.start_time,
end_time: params.filters.end_time,
limit: Some(10000),
trace_id: params.filters.trace_id.clone().filter(|s| !s.is_empty()),
..Default::default()
};
let spans = state.storage.query_spans(&query).await.map_err(|e| {
(
StatusCode::INTERNAL_SERVER_ERROR,
Json(ErrorResponse::storage_error(format!(
"export traces: {}",
e
))),
)
})?;
match params.format.as_str() {
"json" => {
let span_entries: Vec<SpanEntry> = spans.into_iter().map(SpanEntry::from).collect();
Ok((
[
("Content-Type", "application/json"),
(
"Content-Disposition",
"attachment; filename=\"traces.json\"",
),
],
Json(span_entries),
)
.into_response())
},
_ => Err((
StatusCode::BAD_REQUEST,
Json(ErrorResponse::bad_request(
"Invalid format parameter. Use 'json'",
)),
)),
}
}