use std::time::Instant;
use serde::{Deserialize, Serialize};
use tracing::{Level, debug, span, warn};
use crate::error::Result;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QueryPhaseSpan {
pub phase: String,
pub duration_us: u64,
pub success: bool,
pub error: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QueryExecutionTrace {
pub query_id: String,
pub query: String,
pub phases: Vec<QueryPhaseSpan>,
pub total_duration_us: u64,
pub success: bool,
pub error: Option<String>,
pub result_count: Option<usize>,
}
#[must_use = "call .finish() to construct the final value"]
pub struct QueryTraceBuilder {
pub(crate) query_id: String,
pub(crate) query: String,
pub(crate) phases: Vec<QueryPhaseSpan>,
start: Instant,
}
impl QueryTraceBuilder {
pub fn new(query_id: &str, query: &str) -> Self {
let query_str = crate::utils::text::truncate_for_display(query, 500);
Self {
query_id: query_id.to_string(),
query: query_str,
phases: Vec::new(),
start: Instant::now(),
}
}
pub fn record_phase_success(&mut self, phase_name: &str, duration_us: u64) {
self.phases.push(QueryPhaseSpan {
phase: phase_name.to_string(),
duration_us,
success: true,
error: None,
});
debug!(phase = phase_name, duration_us = duration_us, "Query phase completed");
}
pub fn record_phase_error(&mut self, phase_name: &str, duration_us: u64, error: &str) {
self.phases.push(QueryPhaseSpan {
phase: phase_name.to_string(),
duration_us,
success: false,
error: Some(error.to_string()),
});
warn!(
phase = phase_name,
duration_us = duration_us,
error = error,
"Query phase failed"
);
}
pub fn finish(
self,
success: bool,
error: Option<&str>,
result_count: Option<usize>,
) -> Result<QueryExecutionTrace> {
let total_duration_us = u64::try_from(self.start.elapsed().as_micros()).unwrap_or(u64::MAX);
Ok(QueryExecutionTrace {
query_id: self.query_id.clone(),
query: self.query,
phases: self.phases,
total_duration_us,
success,
error: error.map(|e| e.to_string()),
result_count,
})
}
#[must_use]
pub fn query_id(&self) -> &str {
&self.query_id
}
#[must_use]
pub fn elapsed_us(&self) -> u64 {
u64::try_from(self.start.elapsed().as_micros()).unwrap_or(u64::MAX)
}
}
impl QueryExecutionTrace {
#[must_use]
pub fn average_phase_duration_us(&self) -> u64 {
if self.phases.is_empty() {
0
} else {
self.phases.iter().map(|p| p.duration_us).sum::<u64>() / self.phases.len() as u64
}
}
#[must_use]
pub fn slowest_phase(&self) -> Option<&QueryPhaseSpan> {
self.phases.iter().max_by_key(|p| p.duration_us)
}
#[must_use]
pub fn to_log_string(&self) -> String {
let status = if self.success { "success" } else { "error" };
let phases_str = self
.phases
.iter()
.map(|p| format!("{}={}us", p.phase, p.duration_us))
.collect::<Vec<_>>()
.join(" ");
let error_str = self.error.as_ref().map(|e| format!(" error={}", e)).unwrap_or_default();
format!(
"query_id={} status={} total={}us phases=[{}]{}",
self.query_id, status, self.total_duration_us, phases_str, error_str
)
}
}
pub fn create_query_span(query_id: &str, query_text: &str) -> tracing::Span {
span!(
Level::DEBUG,
"query_execute",
query_id = query_id,
query = truncate_query(query_text, 100),
)
}
pub fn create_phase_span(phase_name: &str, query_id: &str) -> tracing::Span {
span!(Level::DEBUG, "query_phase", phase = phase_name, query_id = query_id)
}
pub(crate) fn truncate_query(query: &str, max_len: usize) -> String {
crate::utils::text::truncate_for_display(query, max_len)
}