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 {
query_id: String,
query: String,
phases: Vec<QueryPhaseSpan>,
start: Instant,
}
impl QueryTraceBuilder {
pub fn new(query_id: &str, query: &str) -> Self {
let query_str = if query.len() > 500 {
format!("{}...", &query[..500])
} else {
query.to_string()
};
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,
})
}
pub fn query_id(&self) -> &str {
&self.query_id
}
pub fn elapsed_us(&self) -> u64 {
u64::try_from(self.start.elapsed().as_micros()).unwrap_or(u64::MAX)
}
}
impl QueryExecutionTrace {
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
}
}
pub fn slowest_phase(&self) -> Option<&QueryPhaseSpan> {
self.phases.iter().max_by_key(|p| p.duration_us)
}
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)
}
fn truncate_query(query: &str, max_len: usize) -> String {
if query.len() > max_len {
format!("{}...", &query[..max_len])
} else {
query.to_string()
}
}
#[cfg(test)]
mod tests {
#![allow(clippy::unwrap_used)]
use super::*;
#[test]
fn test_trace_builder_new() {
let builder = QueryTraceBuilder::new("query_1", "{ user { id } }");
assert_eq!(builder.query_id, "query_1");
assert_eq!(builder.query, "{ user { id } }");
assert!(builder.phases.is_empty());
}
#[test]
fn test_trace_builder_truncate_long_query() {
let long_query = "a".repeat(600);
let builder = QueryTraceBuilder::new("query_1", &long_query);
assert!(builder.query.len() < 600);
assert!(builder.query.ends_with("..."));
}
#[test]
fn test_record_phase_success() {
let mut builder = QueryTraceBuilder::new("query_1", "{ user { id } }");
builder.record_phase_success("parse", 100);
builder.record_phase_success("validate", 50);
assert_eq!(builder.phases.len(), 2);
assert_eq!(builder.phases[0].phase, "parse");
assert_eq!(builder.phases[0].duration_us, 100);
assert!(builder.phases[0].success);
}
#[test]
fn test_record_phase_error() {
let mut builder = QueryTraceBuilder::new("query_1", "{ user { id } }");
builder.record_phase_error("parse", 100, "Invalid syntax");
assert_eq!(builder.phases.len(), 1);
assert_eq!(builder.phases[0].phase, "parse");
assert!(!builder.phases[0].success);
assert_eq!(builder.phases[0].error, Some("Invalid syntax".to_string()));
}
#[test]
fn test_finish_success() {
let mut builder = QueryTraceBuilder::new("query_1", "{ user { id } }");
builder.record_phase_success("parse", 100);
builder.record_phase_success("execute", 500);
let trace = builder.finish(true, None, Some(10)).unwrap();
assert!(trace.success);
assert_eq!(trace.query_id, "query_1");
assert_eq!(trace.phases.len(), 2);
assert_eq!(trace.result_count, Some(10));
}
#[test]
fn test_finish_error() {
let mut builder = QueryTraceBuilder::new("query_1", "{ user { id } }");
builder.record_phase_success("parse", 100);
builder.record_phase_error("execute", 50, "Database connection failed");
let trace = builder.finish(false, Some("Database connection failed"), None).unwrap();
assert!(!trace.success);
assert_eq!(trace.error, Some("Database connection failed".to_string()));
}
#[test]
fn test_average_phase_duration() {
let mut builder = QueryTraceBuilder::new("query_1", "{ user { id } }");
builder.record_phase_success("parse", 100);
builder.record_phase_success("validate", 200);
builder.record_phase_success("execute", 300);
let trace = builder.finish(true, None, None).unwrap();
assert_eq!(trace.average_phase_duration_us(), 200);
}
#[test]
fn test_slowest_phase() {
let mut builder = QueryTraceBuilder::new("query_1", "{ user { id } }");
builder.record_phase_success("parse", 100);
builder.record_phase_success("execute", 500);
builder.record_phase_success("cache_check", 50);
let trace = builder.finish(true, None, None).unwrap();
let slowest = trace.slowest_phase().unwrap();
assert_eq!(slowest.phase, "execute");
assert_eq!(slowest.duration_us, 500);
}
#[test]
fn test_to_log_string_success() {
let mut builder = QueryTraceBuilder::new("query_1", "{ user { id } }");
builder.record_phase_success("parse", 100);
builder.record_phase_success("execute", 500);
let trace = builder.finish(true, None, Some(5)).unwrap();
let log_str = trace.to_log_string();
assert!(log_str.contains("query_id=query_1"));
assert!(log_str.contains("status=success"));
assert!(log_str.contains("parse=100us"));
assert!(log_str.contains("execute=500us"));
}
#[test]
fn test_to_log_string_error() {
let mut builder = QueryTraceBuilder::new("query_1", "{ user { id } }");
builder.record_phase_success("parse", 100);
builder.record_phase_error("validate", 50, "Type mismatch");
let trace = builder.finish(false, Some("Type mismatch"), None).unwrap();
let log_str = trace.to_log_string();
assert!(log_str.contains("status=error"));
assert!(log_str.contains("error=Type mismatch"));
}
#[test]
fn test_average_phase_duration_empty() {
let builder = QueryTraceBuilder::new("query_1", "{ user { id } }");
let trace = builder.finish(true, None, None).unwrap();
assert_eq!(trace.average_phase_duration_us(), 0);
}
#[test]
fn test_elapsed_us() {
let builder = QueryTraceBuilder::new("query_1", "{ user { id } }");
let elapsed = builder.elapsed_us();
let _ = elapsed;
}
#[test]
fn test_trace_serialization() {
let mut builder = QueryTraceBuilder::new("query_1", "{ user { id } }");
builder.record_phase_success("parse", 100);
let trace = builder.finish(true, None, Some(5)).unwrap();
let json = serde_json::to_string(&trace).expect("serialize should work");
let restored: QueryExecutionTrace =
serde_json::from_str(&json).expect("deserialize should work");
assert_eq!(restored.query_id, trace.query_id);
assert_eq!(restored.phases.len(), trace.phases.len());
}
#[test]
fn test_query_phase_span_serialize() {
let span = QueryPhaseSpan {
phase: "parse".to_string(),
duration_us: 100,
success: true,
error: None,
};
let json = serde_json::to_string(&span).expect("serialize should work");
let restored: QueryPhaseSpan =
serde_json::from_str(&json).expect("deserialize should work");
assert_eq!(restored.phase, span.phase);
assert_eq!(restored.duration_us, span.duration_us);
}
#[test]
fn test_truncate_query_helper() {
assert_eq!(truncate_query("hello", 100), "hello");
assert!(truncate_query(&"a".repeat(200), 50).ends_with("..."));
}
}