use std::fs::OpenOptions;
use std::io::{BufWriter, Write};
use std::path::{Path, PathBuf};
use std::time::Instant;
use chrono::{DateTime, Utc};
use parking_lot::Mutex;
use serde::Serialize;
use super::guard::QueryCompletionStatus;
#[derive(Debug, Serialize)]
pub struct QueryAuditEntry {
pub query: String,
pub timestamp: DateTime<Utc>,
pub duration_ms: u64,
pub result_count: usize,
pub memory_usage: usize,
pub predicates: Vec<String>,
pub outcome: QueryOutcome,
pub applied_limits: AppliedLimits,
pub user: String,
}
#[derive(Debug, Clone, Serialize)]
pub struct AppliedLimits {
pub timeout_ms: u64,
pub result_cap: usize,
pub memory_limit: usize,
}
#[derive(Debug, Clone, Serialize)]
pub enum QueryOutcome {
Success,
Timeout {
partial_results_count: usize,
},
ResultCapExceeded {
returned: usize,
limit: usize,
},
MemoryLimitExceeded {
returned: usize,
memory_bytes: usize,
},
CostLimitExceeded {
estimated: usize,
limit: usize,
},
Error(String),
}
impl QueryOutcome {
#[must_use]
pub fn from_completion_status(status: &QueryCompletionStatus, result_count: usize) -> Self {
match status {
QueryCompletionStatus::Complete => Self::Success,
QueryCompletionStatus::TimedOut { .. } => Self::Timeout {
partial_results_count: result_count,
},
QueryCompletionStatus::ResultCapReached { count, limit } => Self::ResultCapExceeded {
returned: *count,
limit: *limit,
},
QueryCompletionStatus::MemoryLimitReached { usage_bytes, .. } => {
Self::MemoryLimitExceeded {
returned: result_count,
memory_bytes: *usage_bytes,
}
}
}
}
#[must_use]
pub fn has_partial_results(&self) -> bool {
match self {
Self::Timeout {
partial_results_count,
} if *partial_results_count > 0 => true,
Self::ResultCapExceeded { .. } | Self::MemoryLimitExceeded { .. } => true,
_ => false,
}
}
#[must_use]
pub fn is_success(&self) -> bool {
matches!(self, Self::Success)
}
}
#[derive(Debug, Clone)]
pub struct AuditLogConfig {
pub log_path: PathBuf,
pub buffer_size: usize,
pub flush_interval_secs: u64,
pub max_file_size: usize,
pub rotation_count: usize,
pub user_context: String,
}
impl Default for AuditLogConfig {
fn default() -> Self {
Self {
log_path: PathBuf::from(".sqry-audit.jsonl"),
buffer_size: 100,
flush_interval_secs: 60,
max_file_size: 10 * 1024 * 1024, rotation_count: 5,
user_context: Self::detect_user_context(),
}
}
}
impl AuditLogConfig {
#[must_use]
pub fn detect_user_context() -> String {
std::env::var("SQRY_USER")
.or_else(|_| std::env::var("USER"))
.or_else(|_| std::env::var("USERNAME"))
.unwrap_or_else(|_| "unknown".to_string())
}
#[must_use]
pub fn with_log_path(self, path: PathBuf) -> Self {
Self {
log_path: path,
..self
}
}
#[must_use]
pub fn with_buffer_size(self, size: usize) -> Self {
Self {
buffer_size: size,
..self
}
}
pub fn validate_path(&self) -> Result<(), PathValidationError> {
let parent = self.log_path.parent().unwrap_or(Path::new("."));
if parent.exists() {
let test_path = parent.join(".sqry-audit-test");
match std::fs::write(&test_path, b"test") {
Ok(()) => {
let _ = std::fs::remove_file(&test_path);
Ok(())
}
Err(e) if e.kind() == std::io::ErrorKind::PermissionDenied => {
Err(PathValidationError::PermissionDenied {
path: self.log_path.clone(),
suggestion:
"Use --audit-log to specify a writable path (e.g., ~/.sqry/audit.jsonl)"
.into(),
})
}
Err(e) => Err(PathValidationError::IoError(e)),
}
} else {
match std::fs::create_dir_all(parent) {
Ok(()) => {
let _ = std::fs::remove_dir(parent);
Ok(())
}
Err(e) if e.kind() == std::io::ErrorKind::PermissionDenied => {
Err(PathValidationError::PermissionDenied {
path: self.log_path.clone(),
suggestion:
"Use --audit-log to specify a writable path (e.g., ~/.sqry/audit.jsonl)"
.into(),
})
}
Err(e) => Err(PathValidationError::IoError(e)),
}
}
}
}
#[derive(Debug, thiserror::Error)]
pub enum PathValidationError {
#[error("Permission denied for audit log path {path:?}: {suggestion}")]
PermissionDenied {
path: PathBuf,
suggestion: String,
},
#[error("I/O error validating audit log path: {0}")]
IoError(#[from] std::io::Error),
}
pub struct QueryAuditLogger {
config: AuditLogConfig,
buffer: Mutex<Vec<QueryAuditEntry>>,
last_flush: Mutex<Instant>,
}
impl QueryAuditLogger {
pub fn new(config: AuditLogConfig) -> Result<Self, std::io::Error> {
if let Some(parent) = config.log_path.parent()
&& !parent.as_os_str().is_empty()
&& !parent.exists()
{
std::fs::create_dir_all(parent).map_err(|e| {
std::io::Error::new(
e.kind(),
format!(
"Failed to create audit log directory {}: {e}",
parent.display()
),
)
})?;
}
Ok(Self {
config,
buffer: Mutex::new(Vec::new()),
last_flush: Mutex::new(Instant::now()),
})
}
pub fn log(&self, entry: QueryAuditEntry) {
let mut buffer = self.buffer.lock();
buffer.push(entry);
let should_flush = buffer.len() >= self.config.buffer_size
|| self.last_flush.lock().elapsed().as_secs() >= self.config.flush_interval_secs;
if should_flush {
let entries: Vec<_> = buffer.drain(..).collect();
drop(buffer);
self.flush_entries(&entries);
*self.last_flush.lock() = Instant::now();
}
}
pub fn flush(&self) {
let entries: Vec<_> = self.buffer.lock().drain(..).collect();
if !entries.is_empty() {
self.flush_entries(&entries);
}
}
#[must_use]
pub fn buffer_len(&self) -> usize {
self.buffer.lock().len()
}
#[must_use]
pub fn config(&self) -> &AuditLogConfig {
&self.config
}
fn flush_entries(&self, entries: &[QueryAuditEntry]) {
let max_file_size = u64::try_from(self.config.max_file_size).unwrap_or(u64::MAX);
if let Ok(metadata) = std::fs::metadata(&self.config.log_path)
&& metadata.len() >= max_file_size
&& let Err(e) = self.rotate_log()
{
eprintln!("Audit log rotation failed: {e}");
}
match OpenOptions::new()
.create(true)
.append(true)
.open(&self.config.log_path)
{
Ok(file) => {
let mut writer = BufWriter::new(file);
for entry in entries {
if let Ok(json) = serde_json::to_string(entry) {
let _ = writeln!(writer, "{json}");
}
}
let _ = writer.flush();
}
Err(e) => {
eprintln!("Audit log write failed: {e}");
}
}
}
fn rotate_log(&self) -> std::io::Result<()> {
for i in (1..self.config.rotation_count).rev() {
let old_path = self.config.log_path.with_extension(format!("jsonl.{i}"));
let new_path = self
.config
.log_path
.with_extension(format!("jsonl.{}", i + 1));
if old_path.exists() {
std::fs::rename(&old_path, &new_path)?;
}
}
let backup_path = self.config.log_path.with_extension("jsonl.1");
if self.config.log_path.exists() {
std::fs::rename(&self.config.log_path, backup_path)?;
}
Ok(())
}
}
impl Drop for QueryAuditLogger {
fn drop(&mut self) {
self.flush();
}
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::TempDir;
#[test]
fn test_query_outcome_from_status() {
let complete = QueryCompletionStatus::Complete;
assert!(QueryOutcome::from_completion_status(&complete, 10).is_success());
let cap_reached = QueryCompletionStatus::ResultCapReached {
count: 100,
limit: 100,
};
let outcome = QueryOutcome::from_completion_status(&cap_reached, 100);
assert!(outcome.has_partial_results());
}
#[test]
fn test_query_outcome_partial_results() {
assert!(!QueryOutcome::Success.has_partial_results());
assert!(
QueryOutcome::Timeout {
partial_results_count: 5
}
.has_partial_results()
);
assert!(
!QueryOutcome::Timeout {
partial_results_count: 0
}
.has_partial_results()
);
assert!(
QueryOutcome::ResultCapExceeded {
returned: 100,
limit: 100
}
.has_partial_results()
);
assert!(
QueryOutcome::MemoryLimitExceeded {
returned: 50,
memory_bytes: 1000
}
.has_partial_results()
);
assert!(
!QueryOutcome::CostLimitExceeded {
estimated: 100,
limit: 50
}
.has_partial_results()
);
}
#[test]
fn test_default_config() {
let config = AuditLogConfig::default();
assert_eq!(config.log_path, PathBuf::from(".sqry-audit.jsonl"));
assert_eq!(config.buffer_size, 100);
assert_eq!(config.flush_interval_secs, 60);
assert_eq!(config.max_file_size, 10 * 1024 * 1024);
assert_eq!(config.rotation_count, 5);
}
#[test]
fn test_detect_user_context() {
let user = AuditLogConfig::detect_user_context();
assert!(!user.is_empty() || user == "unknown");
}
#[test]
fn test_logger_creation() {
let temp = TempDir::new().unwrap();
let log_path = temp.path().join("audit.jsonl");
let config = AuditLogConfig {
log_path,
..Default::default()
};
let logger = QueryAuditLogger::new(config).unwrap();
assert_eq!(logger.buffer_len(), 0);
}
#[test]
fn test_logger_buffering() {
let temp = TempDir::new().unwrap();
let log_path = temp.path().join("audit.jsonl");
let config = AuditLogConfig {
log_path: log_path.clone(),
buffer_size: 10, flush_interval_secs: 3600,
..Default::default()
};
let logger = QueryAuditLogger::new(config).unwrap();
let entry = QueryAuditEntry {
query: "impl:Debug".to_string(),
timestamp: Utc::now(),
duration_ms: 100,
result_count: 5,
memory_usage: 1024,
predicates: vec!["impl:Debug".to_string()],
outcome: QueryOutcome::Success,
applied_limits: AppliedLimits {
timeout_ms: 30000,
result_cap: 10000,
memory_limit: 512 * 1024 * 1024,
},
user: "test".to_string(),
};
logger.log(entry);
assert_eq!(logger.buffer_len(), 1);
assert!(!log_path.exists());
logger.flush();
assert_eq!(logger.buffer_len(), 0);
assert!(log_path.exists());
}
#[test]
fn test_auto_flush_on_buffer_full() {
let temp = TempDir::new().unwrap();
let log_path = temp.path().join("audit.jsonl");
let config = AuditLogConfig {
log_path: log_path.clone(),
buffer_size: 2, flush_interval_secs: 3600,
..Default::default()
};
let logger = QueryAuditLogger::new(config).unwrap();
for i in 0..2 {
let entry = QueryAuditEntry {
query: format!("query:{i}"),
timestamp: Utc::now(),
duration_ms: 10,
result_count: i,
memory_usage: 512,
predicates: vec![],
outcome: QueryOutcome::Success,
applied_limits: AppliedLimits {
timeout_ms: 30000,
result_cap: 10000,
memory_limit: 512 * 1024 * 1024,
},
user: "test".to_string(),
};
logger.log(entry);
}
assert_eq!(logger.buffer_len(), 0);
assert!(log_path.exists());
}
#[test]
fn test_flush_on_drop() {
let temp = TempDir::new().unwrap();
let log_path = temp.path().join("audit.jsonl");
{
let config = AuditLogConfig {
log_path: log_path.clone(),
buffer_size: 100, flush_interval_secs: 3600,
..Default::default()
};
let logger = QueryAuditLogger::new(config).unwrap();
let entry = QueryAuditEntry {
query: "test".to_string(),
timestamp: Utc::now(),
duration_ms: 10,
result_count: 1,
memory_usage: 256,
predicates: vec![],
outcome: QueryOutcome::Success,
applied_limits: AppliedLimits {
timeout_ms: 30000,
result_cap: 10000,
memory_limit: 512 * 1024 * 1024,
},
user: "test".to_string(),
};
logger.log(entry);
assert!(!log_path.exists());
}
assert!(log_path.exists());
}
#[test]
fn test_log_rotation() {
let temp = TempDir::new().unwrap();
let log_path = temp.path().join("audit.jsonl");
let config = AuditLogConfig {
log_path: log_path.clone(),
buffer_size: 1, max_file_size: 50, rotation_count: 3,
flush_interval_secs: 0, ..Default::default()
};
let logger = QueryAuditLogger::new(config).unwrap();
for i in 0..5 {
let entry = QueryAuditEntry {
query: format!("query:{i}"),
timestamp: Utc::now(),
duration_ms: 10,
result_count: i,
memory_usage: 512,
predicates: vec![],
outcome: QueryOutcome::Success,
applied_limits: AppliedLimits {
timeout_ms: 30000,
result_cap: 10000,
memory_limit: 512 * 1024 * 1024,
},
user: "test".to_string(),
};
logger.log(entry);
}
assert!(log_path.exists());
}
#[test]
fn test_path_validation_writable() {
let temp = TempDir::new().unwrap();
let log_path = temp.path().join("audit.jsonl");
let config = AuditLogConfig {
log_path,
..Default::default()
};
assert!(config.validate_path().is_ok());
}
#[test]
fn test_config_builder() {
let config = AuditLogConfig::default()
.with_log_path(PathBuf::from("/tmp/test.jsonl"))
.with_buffer_size(50);
assert_eq!(config.log_path, PathBuf::from("/tmp/test.jsonl"));
assert_eq!(config.buffer_size, 50);
}
#[test]
fn test_auto_create_parent_directory() {
let temp = TempDir::new().unwrap();
let nested_path = temp
.path()
.join("subdir")
.join("nested")
.join("audit.jsonl");
let config = AuditLogConfig {
log_path: nested_path.clone(),
..Default::default()
};
let logger = QueryAuditLogger::new(config).unwrap();
assert!(nested_path.parent().unwrap().exists());
let entry = QueryAuditEntry {
query: "test".to_string(),
timestamp: Utc::now(),
duration_ms: 10,
result_count: 1,
memory_usage: 256,
predicates: vec![],
outcome: QueryOutcome::Success,
applied_limits: AppliedLimits {
timeout_ms: 30000,
result_cap: 10000,
memory_limit: 512 * 1024 * 1024,
},
user: "test".to_string(),
};
logger.log(entry);
logger.flush();
assert!(nested_path.exists());
}
}