#![cfg(feature = "slow-query-capture")]
use std::time::{Duration, Instant};
use sz_orm_core::{Connection, DbError, Utc, Value};
use sz_orm_observability::{
query_logger::SerializablePhaseTiming, LogLevel, QueryLogEntry, QueryLogger,
};
pub struct SlowQueryGuard<C: Connection> {
inner: C,
logger: QueryLogger,
slow_threshold: Duration,
logged_jsons: Vec<String>,
logged_entries: Vec<QueryLogEntry>,
degraded: bool,
}
impl<C: Connection> SlowQueryGuard<C> {
pub fn new(conn: C) -> Self {
Self {
inner: conn,
logger: QueryLogger::new(),
slow_threshold: Duration::from_millis(200),
logged_jsons: Vec::new(),
logged_entries: Vec::new(),
degraded: false,
}
}
pub fn with_slow_threshold(mut self, threshold: Duration) -> Self {
self.slow_threshold = threshold;
self
}
pub fn with_sample_rate(mut self, rate: f64) -> Self {
if rate < 0.0 || rate > 1.0 {
self.degraded = true;
} else {
self.logger = self.logger.with_sample_rate(rate);
}
self
}
pub fn with_level(mut self, level: LogLevel) -> Self {
self.logger = self.logger.with_level(level);
self
}
pub async fn execute_with_logging(
&mut self,
sql: &str,
params: &[Value],
) -> Result<u64, DbError> {
let start = Instant::now();
let result = self.inner.execute_with_params(sql, params).await;
let elapsed = start.elapsed();
let is_slow = elapsed >= self.slow_threshold;
self.maybe_log(sql, params, elapsed, is_slow);
result
}
pub fn logged_jsons(&self) -> &[String] {
&self.logged_jsons
}
pub fn logged_entries(&self) -> &[QueryLogEntry] {
&self.logged_entries
}
pub fn is_degraded(&self) -> bool {
self.degraded
}
pub fn inner_mut(&mut self) -> &mut C {
&mut self.inner
}
fn maybe_log(&mut self, sql: &str, params: &[Value], elapsed: Duration, is_slow: bool) {
if self.degraded {
return;
}
let fingerprint = normalize_fingerprint(sql);
let param_strings: Vec<String> = params.iter().map(|v| v.to_string()).collect();
let entry = QueryLogEntry {
query_key: fingerprint,
sql: sql.to_string(),
params: param_strings,
total_elapsed_ms: elapsed.as_millis() as u64,
phase_breakdown: vec![SerializablePhaseTiming {
phase: "execute".to_string(),
start_ms: 0,
duration_ms: elapsed.as_millis() as u64,
}],
slow: is_slow,
from_cache: false,
timestamp: Utc::now().to_rfc3339(),
trace_id: None,
span_id: None,
};
if let Some(json) = self.logger.log(entry.clone()) {
self.logged_jsons.push(json);
self.logged_entries.push(entry);
}
}
}
pub fn normalize_fingerprint(sql: &str) -> String {
let mut result = String::with_capacity(sql.len());
let chars: Vec<char> = sql.chars().collect();
let mut i = 0;
while i < chars.len() {
let c = chars[i];
if c == '\'' {
result.push('?');
i += 1;
while i < chars.len() && chars[i] != '\'' {
i += 1;
}
if i < chars.len() {
i += 1;
}
} else if c == '$' && i + 1 < chars.len() && chars[i + 1].is_ascii_digit() {
result.push('?');
i += 2;
while i < chars.len() && chars[i].is_ascii_digit() {
i += 1;
}
} else if c.is_ascii_digit() {
let mut j = i;
while j < chars.len() && chars[j].is_ascii_digit() {
j += 1;
}
if j == chars.len() || !chars[j].is_ascii_alphabetic() {
result.push('?');
} else {
result.extend(&chars[i..j]);
}
i = j;
} else {
result.push(c);
i += 1;
}
}
result
}
#[cfg(test)]
mod tests {
use super::*;
use sz_orm_core::mock::MockConnection;
#[tokio::test]
async fn test_slow_query_above_threshold_logged() {
let mut mock = MockConnection::new();
mock.expect_any().with_rows_affected(1);
let mut guard = SlowQueryGuard::new(mock).with_slow_threshold(Duration::from_millis(0));
guard
.execute_with_logging("SELECT * FROM t WHERE id = ?", &[Value::I64(1)])
.await
.unwrap();
let entries = guard.logged_entries();
assert!(!entries.is_empty());
assert!(entries[0].slow);
}
#[tokio::test]
async fn test_non_slow_not_logged_without_sampling() {
let mut mock = MockConnection::new();
mock.expect_any().with_rows_affected(1);
let mut guard = SlowQueryGuard::new(mock)
.with_slow_threshold(Duration::from_secs(10))
.with_sample_rate(0.0);
guard.execute_with_logging("SELECT 1", &[]).await.unwrap();
assert!(guard.logged_entries().is_empty());
}
#[tokio::test]
async fn test_query_logger_init_fail_degrade() {
let mut mock = MockConnection::new();
mock.expect_any().with_rows_affected(1);
let mut guard = SlowQueryGuard::new(mock).with_sample_rate(2.0);
assert!(guard.is_degraded());
let result = guard.execute_with_logging("SELECT 1", &[]).await.unwrap();
assert_eq!(result, 1);
assert!(guard.logged_entries().is_empty());
}
#[tokio::test]
async fn test_param_masking() {
let mut mock = MockConnection::new();
mock.expect_any().with_rows_affected(1);
let mut guard = SlowQueryGuard::new(mock)
.with_slow_threshold(Duration::from_millis(0))
.with_level(LogLevel::Debug);
guard
.execute_with_logging(
"SELECT * FROM users WHERE phone = ?",
&[Value::String("13800138000".to_string())],
)
.await
.unwrap();
let jsons = guard.logged_jsons();
assert!(!jsons.is_empty());
assert!(!jsons[0].contains("13800138000"));
}
#[test]
fn test_fingerprint_normalization() {
let fp1 = normalize_fingerprint("SELECT * FROM users WHERE id = 1");
let fp2 = normalize_fingerprint("SELECT * FROM users WHERE id = 2");
assert_eq!(fp1, fp2);
assert!(fp1.contains('?'));
}
#[test]
fn test_fingerprint_pg_placeholders() {
let fp1 = normalize_fingerprint("SELECT * FROM users WHERE id = $1");
let fp2 = normalize_fingerprint("SELECT * FROM users WHERE id = $2");
assert_eq!(fp1, fp2);
}
#[test]
fn test_fingerprint_string_literals() {
let fp1 = normalize_fingerprint("SELECT * FROM users WHERE name = 'alice'");
let fp2 = normalize_fingerprint("SELECT * FROM users WHERE name = 'bob'");
assert_eq!(fp1, fp2);
}
#[tokio::test]
async fn test_no_full_logging() {
let mut mock = MockConnection::new();
mock.expect_any().with_rows_affected(1);
let mut guard = SlowQueryGuard::new(mock)
.with_slow_threshold(Duration::from_secs(10))
.with_sample_rate(0.0);
for _ in 0..100 {
guard.execute_with_logging("SELECT 1", &[]).await.unwrap();
}
assert!(guard.logged_entries().is_empty());
}
#[tokio::test]
async fn test_execute_result_returned_correctly() {
let mut mock = MockConnection::new();
mock.expect_any().with_rows_affected(42);
let mut guard = SlowQueryGuard::new(mock).with_slow_threshold(Duration::from_millis(0));
let result = guard
.execute_with_logging("INSERT INTO t VALUES (?)", &[Value::I64(1)])
.await
.unwrap();
assert_eq!(result, 42);
}
#[tokio::test]
async fn test_slow_query_error_still_logged() {
let mock = MockConnection::new().with_fallback(sz_orm_core::mock::FallbackBehavior::Error);
let mut guard = SlowQueryGuard::new(mock).with_slow_threshold(Duration::from_millis(0));
let _ = guard.execute_with_logging("SELECT bad_sql", &[]).await;
assert!(!guard.logged_entries().is_empty());
assert!(guard.logged_entries()[0].slow);
}
}