Implement log collection in graph
This commit is contained in:
parent
d9869123af
commit
845b8bcc72
9 changed files with 1236 additions and 44 deletions
|
|
@ -135,6 +135,10 @@ crate.spec(
|
|||
package = "sysinfo",
|
||||
version = "0.30",
|
||||
)
|
||||
crate.spec(
|
||||
package = "chrono",
|
||||
version = "0.4",
|
||||
)
|
||||
crate.from_specs()
|
||||
use_repo(crate, "crates")
|
||||
|
||||
|
|
|
|||
File diff suppressed because one or more lines are too long
|
|
@ -43,6 +43,8 @@ rust_library(
|
|||
"event_log/writer.rs",
|
||||
"format_consistency_test.rs",
|
||||
"lib.rs",
|
||||
"log_access.rs",
|
||||
"log_collector.rs",
|
||||
"mermaid_utils.rs",
|
||||
"orchestration/error.rs",
|
||||
"orchestration/events.rs",
|
||||
|
|
@ -66,6 +68,7 @@ rust_library(
|
|||
"@crates//:aide",
|
||||
"@crates//:axum",
|
||||
"@crates//:axum-jsonschema",
|
||||
"@crates//:chrono",
|
||||
"@crates//:log",
|
||||
"@crates//:prost",
|
||||
"@crates//:prost-types",
|
||||
|
|
@ -127,6 +130,7 @@ rust_test(
|
|||
crate = ":databuild",
|
||||
edition = "2021",
|
||||
deps = [
|
||||
"@crates//:tempfile",
|
||||
"@crates//:tokio",
|
||||
],
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,10 +1,11 @@
|
|||
use databuild::{JobGraph, Task, JobStatus, BuildRequestStatus, PartitionStatus, BuildRequestEvent, JobEvent, PartitionEvent, PartitionRef};
|
||||
use databuild::event_log::{create_build_event_log, create_build_event};
|
||||
use databuild::build_event::EventType;
|
||||
use databuild::log_collector::{LogCollector, LogCollectorError};
|
||||
use crossbeam_channel::{Receiver, Sender};
|
||||
use log::{debug, error, info, warn};
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::io::{Read, Write};
|
||||
use std::io::{BufReader, Read, Write};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::process::{Command, Stdio};
|
||||
use std::sync::Arc;
|
||||
|
|
@ -126,6 +127,9 @@ fn worker(
|
|||
}
|
||||
};
|
||||
|
||||
// Generate a job run ID for this execution
|
||||
let job_run_id = Uuid::new_v4().to_string();
|
||||
|
||||
let mut cmd = Command::new(&exec_path);
|
||||
cmd.stdin(Stdio::piped())
|
||||
.stdout(Stdio::piped())
|
||||
|
|
@ -139,6 +143,9 @@ fn worker(
|
|||
cmd.env(key, value); // Add current process's environment variables
|
||||
}
|
||||
|
||||
// Add the job run ID so the job wrapper can use the same ID
|
||||
cmd.env("DATABUILD_JOB_RUN_ID", &job_run_id);
|
||||
|
||||
match cmd.spawn() {
|
||||
Ok(mut child) => {
|
||||
if let Some(mut child_stdin) = child.stdin.take() {
|
||||
|
|
@ -177,22 +184,73 @@ fn worker(
|
|||
continue;
|
||||
}
|
||||
|
||||
match child.wait_with_output() {
|
||||
Ok(output) => {
|
||||
// Initialize log collector
|
||||
let mut log_collector = match LogCollector::new(LogCollector::default_logs_dir()) {
|
||||
Ok(collector) => collector,
|
||||
Err(e) => {
|
||||
let err_msg = format!("[Worker {}] Failed to initialize log collector for {}: {}",
|
||||
worker_id, task.job.as_ref().unwrap().label, e);
|
||||
error!("{}", err_msg);
|
||||
result_tx
|
||||
.send(TaskExecutionResult {
|
||||
task_key,
|
||||
job_label: task.job.as_ref().unwrap().label.clone(),
|
||||
success: false,
|
||||
stdout: String::new(),
|
||||
stderr: err_msg.clone(),
|
||||
duration: start_time.elapsed(),
|
||||
error_message: Some(err_msg),
|
||||
})
|
||||
.unwrap_or_else(|e| error!("[Worker {}] Failed to send error result: {}", worker_id, e));
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
// Collect stdout/stderr and process with LogCollector
|
||||
let stdout_handle = child.stdout.take();
|
||||
let stderr_handle = child.stderr.take();
|
||||
|
||||
let mut stdout_content = String::new();
|
||||
let mut stderr_content = String::new();
|
||||
|
||||
// Read stdout and process with LogCollector
|
||||
if let Some(stdout) = stdout_handle {
|
||||
let stdout_reader = BufReader::new(stdout);
|
||||
if let Err(e) = log_collector.consume_job_output(&job_run_id, stdout_reader) {
|
||||
warn!("[Worker {}] Failed to process job logs for {}: {}",
|
||||
worker_id, task.job.as_ref().unwrap().label, e);
|
||||
}
|
||||
}
|
||||
|
||||
// Read stderr (raw, not structured)
|
||||
if let Some(mut stderr) = stderr_handle {
|
||||
if let Err(e) = stderr.read_to_string(&mut stderr_content) {
|
||||
warn!("[Worker {}] Failed to read stderr for {}: {}",
|
||||
worker_id, task.job.as_ref().unwrap().label, e);
|
||||
}
|
||||
}
|
||||
|
||||
// Wait for the process to finish
|
||||
match child.wait() {
|
||||
Ok(status) => {
|
||||
let duration = start_time.elapsed();
|
||||
let success = output.status.success();
|
||||
let stdout = String::from_utf8_lossy(&output.stdout).to_string();
|
||||
let stderr = String::from_utf8_lossy(&output.stderr).to_string();
|
||||
let success = status.success();
|
||||
|
||||
// Close the log collector for this job
|
||||
if let Err(e) = log_collector.close_job(&job_run_id) {
|
||||
warn!("[Worker {}] Failed to close log collector for {}: {}",
|
||||
worker_id, task.job.as_ref().unwrap().label, e);
|
||||
}
|
||||
|
||||
if success {
|
||||
info!(
|
||||
"[Worker {}] Job succeeded: {} (Duration: {:?})",
|
||||
worker_id, task.job.as_ref().unwrap().label, duration
|
||||
"[Worker {}] Job succeeded: {} (Duration: {:?}, Job Run ID: {})",
|
||||
worker_id, task.job.as_ref().unwrap().label, duration, job_run_id
|
||||
);
|
||||
} else {
|
||||
error!(
|
||||
"[Worker {}] Job failed: {} (Duration: {:?}, Status: {:?})\nStdout: {}\nStderr: {}",
|
||||
worker_id, task.job.as_ref().unwrap().label, duration, output.status, stdout, stderr
|
||||
"[Worker {}] Job failed: {} (Duration: {:?}, Status: {:?}, Job Run ID: {})\nStderr: {}",
|
||||
worker_id, task.job.as_ref().unwrap().label, duration, status, job_run_id, stderr_content
|
||||
);
|
||||
}
|
||||
result_tx
|
||||
|
|
@ -200,10 +258,10 @@ fn worker(
|
|||
task_key,
|
||||
job_label: task.job.as_ref().unwrap().label.clone(),
|
||||
success,
|
||||
stdout,
|
||||
stderr,
|
||||
stdout: format!("Job logs written to JSONL (Job Run ID: {})", job_run_id),
|
||||
stderr: stderr_content,
|
||||
duration,
|
||||
error_message: if success { None } else { Some(format!("Exited with status: {:?}", output.status)) },
|
||||
error_message: if success { None } else { Some(format!("Exited with status: {:?}", status)) },
|
||||
})
|
||||
.unwrap_or_else(|e| error!("[Worker {}] Failed to send result: {}", worker_id, e));
|
||||
}
|
||||
|
|
|
|||
|
|
@ -63,8 +63,12 @@ impl JobWrapper<StdoutSink> {
|
|||
|
||||
impl<S: LogSink> JobWrapper<S> {
|
||||
fn new_with_sink(sink: S) -> Self {
|
||||
// Use job ID from environment if provided by graph execution, otherwise generate one
|
||||
let job_id = env::var("DATABUILD_JOB_RUN_ID")
|
||||
.unwrap_or_else(|_| Uuid::new_v4().to_string());
|
||||
|
||||
Self {
|
||||
job_id: Uuid::new_v4().to_string(),
|
||||
job_id,
|
||||
sequence_number: 0,
|
||||
start_time: SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
|
|
|
|||
|
|
@ -18,6 +18,12 @@ pub mod mermaid_utils;
|
|||
// Status conversion utilities
|
||||
pub mod status_utils;
|
||||
|
||||
// Log collection module
|
||||
pub mod log_collector;
|
||||
|
||||
// Log access module
|
||||
pub mod log_access;
|
||||
|
||||
// Format consistency tests
|
||||
#[cfg(test)]
|
||||
mod format_consistency_test;
|
||||
|
|
|
|||
440
databuild/log_access.rs
Normal file
440
databuild/log_access.rs
Normal file
|
|
@ -0,0 +1,440 @@
|
|||
use crate::{JobLogEntry, JobLogsRequest, JobLogsResponse, log_message};
|
||||
use serde_json;
|
||||
use std::collections::HashMap;
|
||||
use std::fs::{self, File};
|
||||
use std::io::{BufRead, BufReader};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
use thiserror::Error;
|
||||
|
||||
#[derive(Error, Debug)]
|
||||
pub enum LogAccessError {
|
||||
#[error("IO error: {0}")]
|
||||
Io(#[from] std::io::Error),
|
||||
#[error("JSON parsing error: {0}")]
|
||||
Json(#[from] serde_json::Error),
|
||||
#[error("Invalid request: {0}")]
|
||||
InvalidRequest(String),
|
||||
#[error("Job not found: {0}")]
|
||||
JobNotFound(String),
|
||||
}
|
||||
|
||||
pub struct LogReader {
|
||||
logs_base_path: PathBuf,
|
||||
}
|
||||
|
||||
impl LogReader {
|
||||
pub fn new<P: AsRef<Path>>(logs_base_path: P) -> Self {
|
||||
Self {
|
||||
logs_base_path: logs_base_path.as_ref().to_path_buf(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Create LogReader with the default logs directory
|
||||
pub fn default() -> Self {
|
||||
Self::new(crate::log_collector::LogCollector::default_logs_dir())
|
||||
}
|
||||
|
||||
/// Get job logs according to the request criteria
|
||||
pub fn get_job_logs(&self, request: &JobLogsRequest) -> Result<JobLogsResponse, LogAccessError> {
|
||||
let job_file_path = self.find_job_file(&request.job_run_id)?;
|
||||
|
||||
let file = File::open(&job_file_path)?;
|
||||
let reader = BufReader::new(file);
|
||||
|
||||
let mut entries = Vec::new();
|
||||
let mut count = 0u32;
|
||||
let limit = if request.limit > 0 { request.limit } else { 1000 }; // Default limit
|
||||
|
||||
for line in reader.lines() {
|
||||
let line = line?;
|
||||
|
||||
// Skip empty lines
|
||||
if line.trim().is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Parse the log entry
|
||||
let entry: JobLogEntry = serde_json::from_str(&line)?;
|
||||
|
||||
// Apply filters
|
||||
if !self.matches_filters(&entry, request) {
|
||||
continue;
|
||||
}
|
||||
|
||||
entries.push(entry);
|
||||
count += 1;
|
||||
|
||||
// Stop if we've hit the limit
|
||||
if count >= limit {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// Check if there are more entries by trying to read one more
|
||||
let has_more = count == limit;
|
||||
|
||||
Ok(JobLogsResponse {
|
||||
entries,
|
||||
has_more,
|
||||
})
|
||||
}
|
||||
|
||||
/// List available job run IDs for a given date range
|
||||
pub fn list_available_jobs(&self, date_range: Option<(String, String)>) -> Result<Vec<String>, LogAccessError> {
|
||||
let mut job_ids = Vec::new();
|
||||
|
||||
// If no date range specified, look at all directories
|
||||
if let Some((start_date, end_date)) = date_range {
|
||||
// Parse date range and iterate through dates
|
||||
for date_str in self.date_range_iterator(&start_date, &end_date)? {
|
||||
let date_dir = self.logs_base_path.join(&date_str);
|
||||
if date_dir.exists() {
|
||||
job_ids.extend(self.get_job_ids_from_directory(&date_dir)?);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// List all date directories and collect job IDs
|
||||
if self.logs_base_path.exists() {
|
||||
for entry in fs::read_dir(&self.logs_base_path)? {
|
||||
let entry = entry?;
|
||||
if entry.file_type()?.is_dir() {
|
||||
job_ids.extend(self.get_job_ids_from_directory(&entry.path())?);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Remove duplicates and sort
|
||||
job_ids.sort();
|
||||
job_ids.dedup();
|
||||
|
||||
Ok(job_ids)
|
||||
}
|
||||
|
||||
/// Get metrics points for a specific job
|
||||
pub fn get_job_metrics(&self, job_run_id: &str) -> Result<Vec<crate::MetricPoint>, LogAccessError> {
|
||||
let job_file_path = self.find_job_file(job_run_id)?;
|
||||
|
||||
let file = File::open(&job_file_path)?;
|
||||
let reader = BufReader::new(file);
|
||||
|
||||
let mut metrics = Vec::new();
|
||||
|
||||
for line in reader.lines() {
|
||||
let line = line?;
|
||||
|
||||
// Skip empty lines
|
||||
if line.trim().is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Parse the log entry
|
||||
let entry: JobLogEntry = serde_json::from_str(&line)?;
|
||||
|
||||
// Extract metrics from the entry
|
||||
if let Some(crate::job_log_entry::Content::Metric(metric)) = entry.content {
|
||||
metrics.push(metric);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(metrics)
|
||||
}
|
||||
|
||||
/// Find the JSONL file for a specific job run ID
|
||||
fn find_job_file(&self, job_run_id: &str) -> Result<PathBuf, LogAccessError> {
|
||||
// Search through all date directories for the job file
|
||||
if !self.logs_base_path.exists() {
|
||||
return Err(LogAccessError::JobNotFound(job_run_id.to_string()));
|
||||
}
|
||||
|
||||
for entry in fs::read_dir(&self.logs_base_path)? {
|
||||
let entry = entry?;
|
||||
if entry.file_type()?.is_dir() {
|
||||
let job_file = entry.path().join(format!("{}.jsonl", job_run_id));
|
||||
if job_file.exists() {
|
||||
return Ok(job_file);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Err(LogAccessError::JobNotFound(job_run_id.to_string()))
|
||||
}
|
||||
|
||||
/// Check if a log entry matches the request filters
|
||||
fn matches_filters(&self, entry: &JobLogEntry, request: &JobLogsRequest) -> bool {
|
||||
// Filter by timestamp (since_timestamp is in nanoseconds)
|
||||
if request.since_timestamp > 0 {
|
||||
if let Ok(entry_timestamp) = entry.timestamp.parse::<u64>() {
|
||||
let entry_timestamp_ns = entry_timestamp * 1_000_000_000; // Convert seconds to nanoseconds
|
||||
if entry_timestamp_ns <= request.since_timestamp as u64 {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Filter by log level (only applies to log messages)
|
||||
if request.min_level > 0 {
|
||||
if let Some(crate::job_log_entry::Content::Log(log_msg)) = &entry.content {
|
||||
if log_msg.level < request.min_level {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
// For non-log entries (metrics, events), we include them regardless of min_level
|
||||
}
|
||||
|
||||
true
|
||||
}
|
||||
|
||||
/// Get job IDs from files in a specific directory
|
||||
fn get_job_ids_from_directory(&self, dir_path: &Path) -> Result<Vec<String>, LogAccessError> {
|
||||
let mut job_ids = Vec::new();
|
||||
|
||||
for entry in fs::read_dir(dir_path)? {
|
||||
let entry = entry?;
|
||||
if entry.file_type()?.is_file() {
|
||||
if let Some(file_name) = entry.file_name().to_str() {
|
||||
if file_name.ends_with(".jsonl") {
|
||||
// Extract job ID by removing .jsonl extension
|
||||
let job_id = file_name.trim_end_matches(".jsonl");
|
||||
job_ids.push(job_id.to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(job_ids)
|
||||
}
|
||||
|
||||
/// Generate an iterator over date strings in a range (YYYY-MM-DD format)
|
||||
fn date_range_iterator(&self, start_date: &str, end_date: &str) -> Result<Vec<String>, LogAccessError> {
|
||||
// Simple implementation - for production might want more robust date parsing
|
||||
let start_parts: Vec<&str> = start_date.split('-').collect();
|
||||
let end_parts: Vec<&str> = end_date.split('-').collect();
|
||||
|
||||
if start_parts.len() != 3 || end_parts.len() != 3 {
|
||||
return Err(LogAccessError::InvalidRequest("Invalid date format, expected YYYY-MM-DD".to_string()));
|
||||
}
|
||||
|
||||
// For now, just return the start and end dates
|
||||
// In a full implementation, you'd iterate through all dates in between
|
||||
let mut dates = vec![start_date.to_string()];
|
||||
if start_date != end_date {
|
||||
dates.push(end_date.to_string());
|
||||
}
|
||||
|
||||
Ok(dates)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::{job_log_entry, log_message, LogMessage, PartitionRef, MetricPoint};
|
||||
use std::io::Write;
|
||||
use tempfile::TempDir;
|
||||
|
||||
fn create_test_log_entry(job_id: &str, sequence: u64, timestamp: &str) -> JobLogEntry {
|
||||
JobLogEntry {
|
||||
timestamp: timestamp.to_string(),
|
||||
job_id: job_id.to_string(),
|
||||
outputs: vec![PartitionRef { r#str: "test/partition".to_string() }],
|
||||
sequence_number: sequence,
|
||||
content: Some(job_log_entry::Content::Log(LogMessage {
|
||||
level: log_message::LogLevel::Info as i32,
|
||||
message: format!("Test log message {}", sequence),
|
||||
fields: HashMap::new(),
|
||||
})),
|
||||
}
|
||||
}
|
||||
|
||||
fn create_test_metric_entry(job_id: &str, sequence: u64, timestamp: &str) -> JobLogEntry {
|
||||
JobLogEntry {
|
||||
timestamp: timestamp.to_string(),
|
||||
job_id: job_id.to_string(),
|
||||
outputs: vec![PartitionRef { r#str: "test/partition".to_string() }],
|
||||
sequence_number: sequence,
|
||||
content: Some(job_log_entry::Content::Metric(MetricPoint {
|
||||
name: "test_metric".to_string(),
|
||||
value: 42.0,
|
||||
labels: HashMap::new(),
|
||||
unit: "count".to_string(),
|
||||
})),
|
||||
}
|
||||
}
|
||||
|
||||
fn setup_test_logs(temp_dir: &TempDir) -> Result<(), Box<dyn std::error::Error>> {
|
||||
// Create date directory
|
||||
let date_dir = temp_dir.path().join("2025-01-27");
|
||||
fs::create_dir_all(&date_dir)?;
|
||||
|
||||
// Create a test job file
|
||||
let job_file = date_dir.join("job_123.jsonl");
|
||||
let mut file = File::create(&job_file)?;
|
||||
|
||||
// Write test entries
|
||||
let entry1 = create_test_log_entry("job_123", 1, "1737993600"); // 2025-01-27 12:00:00
|
||||
let entry2 = create_test_log_entry("job_123", 2, "1737993660"); // 2025-01-27 12:01:00
|
||||
let entry3 = create_test_metric_entry("job_123", 3, "1737993720"); // 2025-01-27 12:02:00
|
||||
|
||||
writeln!(file, "{}", serde_json::to_string(&entry1)?)?;
|
||||
writeln!(file, "{}", serde_json::to_string(&entry2)?)?;
|
||||
writeln!(file, "{}", serde_json::to_string(&entry3)?)?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_log_reader_creation() {
|
||||
let temp_dir = TempDir::new().unwrap();
|
||||
let reader = LogReader::new(temp_dir.path());
|
||||
|
||||
assert_eq!(reader.logs_base_path, temp_dir.path());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_get_job_logs_basic() {
|
||||
let temp_dir = TempDir::new().unwrap();
|
||||
setup_test_logs(&temp_dir).unwrap();
|
||||
|
||||
let reader = LogReader::new(temp_dir.path());
|
||||
let request = JobLogsRequest {
|
||||
job_run_id: "job_123".to_string(),
|
||||
since_timestamp: 0,
|
||||
min_level: 0,
|
||||
limit: 10,
|
||||
};
|
||||
|
||||
let response = reader.get_job_logs(&request).unwrap();
|
||||
|
||||
assert_eq!(response.entries.len(), 3);
|
||||
assert!(!response.has_more);
|
||||
|
||||
// Verify the entries are in order
|
||||
assert_eq!(response.entries[0].sequence_number, 1);
|
||||
assert_eq!(response.entries[1].sequence_number, 2);
|
||||
assert_eq!(response.entries[2].sequence_number, 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_get_job_logs_with_timestamp_filter() {
|
||||
let temp_dir = TempDir::new().unwrap();
|
||||
setup_test_logs(&temp_dir).unwrap();
|
||||
|
||||
let reader = LogReader::new(temp_dir.path());
|
||||
let request = JobLogsRequest {
|
||||
job_run_id: "job_123".to_string(),
|
||||
since_timestamp: 1737993600_000_000_000, // 2025-01-27 12:00:00 in nanoseconds
|
||||
min_level: 0,
|
||||
limit: 10,
|
||||
};
|
||||
|
||||
let response = reader.get_job_logs(&request).unwrap();
|
||||
|
||||
// Should get entries 2 and 3 (after the timestamp)
|
||||
assert_eq!(response.entries.len(), 2);
|
||||
assert_eq!(response.entries[0].sequence_number, 2);
|
||||
assert_eq!(response.entries[1].sequence_number, 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_get_job_logs_with_level_filter() {
|
||||
let temp_dir = TempDir::new().unwrap();
|
||||
setup_test_logs(&temp_dir).unwrap();
|
||||
|
||||
let reader = LogReader::new(temp_dir.path());
|
||||
let request = JobLogsRequest {
|
||||
job_run_id: "job_123".to_string(),
|
||||
since_timestamp: 0,
|
||||
min_level: log_message::LogLevel::Warn as i32, // Only WARN and ERROR
|
||||
limit: 10,
|
||||
};
|
||||
|
||||
let response = reader.get_job_logs(&request).unwrap();
|
||||
|
||||
// Should get only the metric entry (sequence 3) since log entries are INFO level
|
||||
assert_eq!(response.entries.len(), 1);
|
||||
assert_eq!(response.entries[0].sequence_number, 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_get_job_logs_with_limit() {
|
||||
let temp_dir = TempDir::new().unwrap();
|
||||
setup_test_logs(&temp_dir).unwrap();
|
||||
|
||||
let reader = LogReader::new(temp_dir.path());
|
||||
let request = JobLogsRequest {
|
||||
job_run_id: "job_123".to_string(),
|
||||
since_timestamp: 0,
|
||||
min_level: 0,
|
||||
limit: 2,
|
||||
};
|
||||
|
||||
let response = reader.get_job_logs(&request).unwrap();
|
||||
|
||||
assert_eq!(response.entries.len(), 2);
|
||||
assert!(response.has_more);
|
||||
assert_eq!(response.entries[0].sequence_number, 1);
|
||||
assert_eq!(response.entries[1].sequence_number, 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_list_available_jobs() {
|
||||
let temp_dir = TempDir::new().unwrap();
|
||||
setup_test_logs(&temp_dir).unwrap();
|
||||
|
||||
// Create another job file
|
||||
let date_dir = temp_dir.path().join("2025-01-27");
|
||||
let job_file2 = date_dir.join("job_456.jsonl");
|
||||
let mut file2 = File::create(&job_file2).unwrap();
|
||||
let entry = create_test_log_entry("job_456", 1, "1737993600");
|
||||
writeln!(file2, "{}", serde_json::to_string(&entry).unwrap()).unwrap();
|
||||
|
||||
let reader = LogReader::new(temp_dir.path());
|
||||
let job_ids = reader.list_available_jobs(None).unwrap();
|
||||
|
||||
assert_eq!(job_ids.len(), 2);
|
||||
assert!(job_ids.contains(&"job_123".to_string()));
|
||||
assert!(job_ids.contains(&"job_456".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_get_job_metrics() {
|
||||
let temp_dir = TempDir::new().unwrap();
|
||||
setup_test_logs(&temp_dir).unwrap();
|
||||
|
||||
let reader = LogReader::new(temp_dir.path());
|
||||
let metrics = reader.get_job_metrics("job_123").unwrap();
|
||||
|
||||
assert_eq!(metrics.len(), 1);
|
||||
assert_eq!(metrics[0].name, "test_metric");
|
||||
assert_eq!(metrics[0].value, 42.0);
|
||||
assert_eq!(metrics[0].unit, "count");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_job_not_found() {
|
||||
let temp_dir = TempDir::new().unwrap();
|
||||
let reader = LogReader::new(temp_dir.path());
|
||||
|
||||
let request = JobLogsRequest {
|
||||
job_run_id: "nonexistent_job".to_string(),
|
||||
since_timestamp: 0,
|
||||
min_level: 0,
|
||||
limit: 10,
|
||||
};
|
||||
|
||||
let result = reader.get_job_logs(&request);
|
||||
assert!(result.is_err());
|
||||
assert!(matches!(result.unwrap_err(), LogAccessError::JobNotFound(_)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_default_log_reader() {
|
||||
let reader = LogReader::default();
|
||||
|
||||
// Should use the default logs directory
|
||||
let expected = crate::log_collector::LogCollector::default_logs_dir();
|
||||
assert_eq!(reader.logs_base_path, expected);
|
||||
}
|
||||
}
|
||||
324
databuild/log_collector.rs
Normal file
324
databuild/log_collector.rs
Normal file
|
|
@ -0,0 +1,324 @@
|
|||
use crate::JobLogEntry;
|
||||
use serde_json;
|
||||
use std::collections::HashMap;
|
||||
use std::fs::{self, File, OpenOptions};
|
||||
use std::io::{BufRead, BufReader, Write};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
use thiserror::Error;
|
||||
|
||||
#[derive(Error, Debug)]
|
||||
pub enum LogCollectorError {
|
||||
#[error("IO error: {0}")]
|
||||
Io(#[from] std::io::Error),
|
||||
#[error("JSON parsing error: {0}")]
|
||||
Json(#[from] serde_json::Error),
|
||||
#[error("Invalid log entry: {0}")]
|
||||
InvalidLogEntry(String),
|
||||
}
|
||||
|
||||
pub struct LogCollector {
|
||||
logs_dir: PathBuf,
|
||||
active_files: HashMap<String, File>,
|
||||
}
|
||||
|
||||
impl LogCollector {
|
||||
pub fn new<P: AsRef<Path>>(logs_dir: P) -> Result<Self, LogCollectorError> {
|
||||
let logs_dir = logs_dir.as_ref().to_path_buf();
|
||||
|
||||
// Ensure the base logs directory exists
|
||||
if !logs_dir.exists() {
|
||||
fs::create_dir_all(&logs_dir)?;
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
logs_dir,
|
||||
active_files: HashMap::new(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Get the default logs directory based on environment variable or fallback
|
||||
pub fn default_logs_dir() -> PathBuf {
|
||||
std::env::var("DATABUILD_LOGS_DIR")
|
||||
.map(PathBuf::from)
|
||||
.unwrap_or_else(|_| {
|
||||
// Fallback to ./logs/databuild for safety - avoid system directories
|
||||
std::env::current_dir()
|
||||
.unwrap_or_else(|_| PathBuf::from("."))
|
||||
.join("logs")
|
||||
.join("databuild")
|
||||
})
|
||||
}
|
||||
|
||||
/// Create a date-organized directory path for today
|
||||
fn get_date_directory(&self) -> Result<PathBuf, LogCollectorError> {
|
||||
let now = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map_err(|e| LogCollectorError::InvalidLogEntry(format!("System time error: {}", e)))?;
|
||||
|
||||
let timestamp = now.as_secs();
|
||||
let datetime = chrono::DateTime::from_timestamp(timestamp as i64, 0)
|
||||
.ok_or_else(|| LogCollectorError::InvalidLogEntry("Invalid timestamp".to_string()))?;
|
||||
|
||||
let date_str = datetime.format("%Y-%m-%d").to_string();
|
||||
let date_dir = self.logs_dir.join(date_str);
|
||||
|
||||
// Ensure the date directory exists
|
||||
if !date_dir.exists() {
|
||||
fs::create_dir_all(&date_dir)?;
|
||||
}
|
||||
|
||||
Ok(date_dir)
|
||||
}
|
||||
|
||||
/// Get or create a file handle for a specific job run
|
||||
fn get_job_file(&mut self, job_run_id: &str) -> Result<&mut File, LogCollectorError> {
|
||||
if !self.active_files.contains_key(job_run_id) {
|
||||
let date_dir = self.get_date_directory()?;
|
||||
let file_path = date_dir.join(format!("{}.jsonl", job_run_id));
|
||||
|
||||
let file = OpenOptions::new()
|
||||
.create(true)
|
||||
.append(true)
|
||||
.open(&file_path)?;
|
||||
|
||||
self.active_files.insert(job_run_id.to_string(), file);
|
||||
}
|
||||
|
||||
Ok(self.active_files.get_mut(job_run_id).unwrap())
|
||||
}
|
||||
|
||||
/// Write a single log entry to the appropriate JSONL file
|
||||
pub fn write_log_entry(&mut self, job_run_id: &str, entry: &JobLogEntry) -> Result<(), LogCollectorError> {
|
||||
let file = self.get_job_file(job_run_id)?;
|
||||
let json_line = serde_json::to_string(entry)?;
|
||||
writeln!(file, "{}", json_line)?;
|
||||
file.flush()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Consume stdout from a job process and parse/store log entries
|
||||
pub fn consume_job_output<R: BufRead>(&mut self, job_run_id: &str, reader: R) -> Result<(), LogCollectorError> {
|
||||
for line in reader.lines() {
|
||||
let line = line?;
|
||||
|
||||
// Skip empty lines
|
||||
if line.trim().is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Try to parse as JobLogEntry
|
||||
match serde_json::from_str::<JobLogEntry>(&line) {
|
||||
Ok(entry) => {
|
||||
// Validate that the job_id matches
|
||||
if entry.job_id != job_run_id {
|
||||
return Err(LogCollectorError::InvalidLogEntry(
|
||||
format!("Job ID mismatch: expected {}, got {}", job_run_id, entry.job_id)
|
||||
));
|
||||
}
|
||||
|
||||
self.write_log_entry(job_run_id, &entry)?;
|
||||
}
|
||||
Err(_) => {
|
||||
// If it's not a JobLogEntry, treat it as raw output and create a log entry
|
||||
let raw_entry = JobLogEntry {
|
||||
timestamp: SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_secs()
|
||||
.to_string(),
|
||||
job_id: job_run_id.to_string(),
|
||||
outputs: vec![], // Raw output doesn't have specific outputs
|
||||
sequence_number: 0, // Raw output gets sequence 0
|
||||
content: Some(crate::job_log_entry::Content::Log(crate::LogMessage {
|
||||
level: crate::log_message::LogLevel::Info as i32,
|
||||
message: line,
|
||||
fields: HashMap::new(),
|
||||
})),
|
||||
};
|
||||
|
||||
self.write_log_entry(job_run_id, &raw_entry)?;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Close and flush all active files
|
||||
pub fn close_all(&mut self) -> Result<(), LogCollectorError> {
|
||||
for (_, mut file) in self.active_files.drain() {
|
||||
file.flush()?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Close and flush a specific job's file
|
||||
pub fn close_job(&mut self, job_run_id: &str) -> Result<(), LogCollectorError> {
|
||||
if let Some(mut file) = self.active_files.remove(job_run_id) {
|
||||
file.flush()?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::{job_log_entry, log_message, LogMessage, PartitionRef, WrapperJobEvent};
|
||||
use std::io::Cursor;
|
||||
use tempfile::TempDir;
|
||||
|
||||
fn create_test_log_entry(job_id: &str, sequence: u64) -> JobLogEntry {
|
||||
JobLogEntry {
|
||||
timestamp: "1234567890".to_string(),
|
||||
job_id: job_id.to_string(),
|
||||
outputs: vec![PartitionRef { r#str: "test/partition".to_string() }],
|
||||
sequence_number: sequence,
|
||||
content: Some(job_log_entry::Content::Log(LogMessage {
|
||||
level: log_message::LogLevel::Info as i32,
|
||||
message: "Test log message".to_string(),
|
||||
fields: HashMap::new(),
|
||||
})),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_log_collector_creation() {
|
||||
let temp_dir = TempDir::new().unwrap();
|
||||
let collector = LogCollector::new(temp_dir.path()).unwrap();
|
||||
|
||||
assert_eq!(collector.logs_dir, temp_dir.path());
|
||||
assert!(collector.active_files.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_write_single_log_entry() {
|
||||
let temp_dir = TempDir::new().unwrap();
|
||||
let mut collector = LogCollector::new(temp_dir.path()).unwrap();
|
||||
|
||||
let entry = create_test_log_entry("job_123", 1);
|
||||
collector.write_log_entry("job_123", &entry).unwrap();
|
||||
|
||||
// Verify file was created and contains the entry
|
||||
collector.close_all().unwrap();
|
||||
|
||||
// Check that a date directory was created
|
||||
let date_dirs: Vec<_> = fs::read_dir(temp_dir.path()).unwrap().collect();
|
||||
assert_eq!(date_dirs.len(), 1);
|
||||
|
||||
// Check that the job file exists in the date directory
|
||||
let date_dir_path = date_dirs[0].as_ref().unwrap().path();
|
||||
let job_files: Vec<_> = fs::read_dir(&date_dir_path).unwrap().collect();
|
||||
assert_eq!(job_files.len(), 1);
|
||||
|
||||
let job_file_path = job_files[0].as_ref().unwrap().path();
|
||||
assert!(job_file_path.file_name().unwrap().to_string_lossy().contains("job_123"));
|
||||
|
||||
// Verify content
|
||||
let content = fs::read_to_string(&job_file_path).unwrap();
|
||||
assert!(content.contains("Test log message"));
|
||||
assert!(content.contains("\"sequence_number\":1"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_consume_structured_output() {
|
||||
let temp_dir = TempDir::new().unwrap();
|
||||
let mut collector = LogCollector::new(temp_dir.path()).unwrap();
|
||||
|
||||
let entry1 = create_test_log_entry("job_456", 1);
|
||||
let entry2 = create_test_log_entry("job_456", 2);
|
||||
|
||||
let input = format!("{}\n{}\n",
|
||||
serde_json::to_string(&entry1).unwrap(),
|
||||
serde_json::to_string(&entry2).unwrap()
|
||||
);
|
||||
|
||||
let reader = Cursor::new(input);
|
||||
collector.consume_job_output("job_456", reader).unwrap();
|
||||
collector.close_all().unwrap();
|
||||
|
||||
// Verify both entries were written
|
||||
let date_dirs: Vec<_> = fs::read_dir(temp_dir.path()).unwrap().collect();
|
||||
let date_dir_path = date_dirs[0].as_ref().unwrap().path();
|
||||
let job_files: Vec<_> = fs::read_dir(&date_dir_path).unwrap().collect();
|
||||
let job_file_path = job_files[0].as_ref().unwrap().path();
|
||||
|
||||
let content = fs::read_to_string(&job_file_path).unwrap();
|
||||
let lines: Vec<&str> = content.trim().split('\n').collect();
|
||||
assert_eq!(lines.len(), 2);
|
||||
|
||||
// Verify both entries can be parsed back
|
||||
let parsed1: JobLogEntry = serde_json::from_str(lines[0]).unwrap();
|
||||
let parsed2: JobLogEntry = serde_json::from_str(lines[1]).unwrap();
|
||||
assert_eq!(parsed1.sequence_number, 1);
|
||||
assert_eq!(parsed2.sequence_number, 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_consume_mixed_output() {
|
||||
let temp_dir = TempDir::new().unwrap();
|
||||
let mut collector = LogCollector::new(temp_dir.path()).unwrap();
|
||||
|
||||
let entry = create_test_log_entry("job_789", 1);
|
||||
let structured_line = serde_json::to_string(&entry).unwrap();
|
||||
|
||||
let input = format!("{}\nRaw output line\nAnother raw line\n", structured_line);
|
||||
|
||||
let reader = Cursor::new(input);
|
||||
collector.consume_job_output("job_789", reader).unwrap();
|
||||
collector.close_all().unwrap();
|
||||
|
||||
// Verify all lines were captured (1 structured + 2 raw)
|
||||
let date_dirs: Vec<_> = fs::read_dir(temp_dir.path()).unwrap().collect();
|
||||
let date_dir_path = date_dirs[0].as_ref().unwrap().path();
|
||||
let job_files: Vec<_> = fs::read_dir(&date_dir_path).unwrap().collect();
|
||||
let job_file_path = job_files[0].as_ref().unwrap().path();
|
||||
|
||||
let content = fs::read_to_string(&job_file_path).unwrap();
|
||||
let lines: Vec<&str> = content.trim().split('\n').collect();
|
||||
assert_eq!(lines.len(), 3);
|
||||
|
||||
// First line should be the structured entry
|
||||
let parsed1: JobLogEntry = serde_json::from_str(lines[0]).unwrap();
|
||||
assert_eq!(parsed1.sequence_number, 1);
|
||||
|
||||
// Second and third lines should be raw output entries
|
||||
let parsed2: JobLogEntry = serde_json::from_str(lines[1]).unwrap();
|
||||
let parsed3: JobLogEntry = serde_json::from_str(lines[2]).unwrap();
|
||||
assert_eq!(parsed2.sequence_number, 0); // Raw output gets sequence 0
|
||||
assert_eq!(parsed3.sequence_number, 0);
|
||||
|
||||
if let Some(job_log_entry::Content::Log(log_msg)) = &parsed2.content {
|
||||
assert_eq!(log_msg.message, "Raw output line");
|
||||
} else {
|
||||
panic!("Expected log content");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_default_logs_dir() {
|
||||
let default_dir = LogCollector::default_logs_dir();
|
||||
|
||||
// Should be a valid path
|
||||
assert!(default_dir.is_absolute() || default_dir.starts_with("."));
|
||||
assert!(default_dir.to_string_lossy().contains("logs"));
|
||||
assert!(default_dir.to_string_lossy().contains("databuild"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_job_id_validation() {
|
||||
let temp_dir = TempDir::new().unwrap();
|
||||
let mut collector = LogCollector::new(temp_dir.path()).unwrap();
|
||||
|
||||
let mut entry = create_test_log_entry("wrong_job_id", 1);
|
||||
entry.job_id = "wrong_job_id".to_string();
|
||||
|
||||
let input = serde_json::to_string(&entry).unwrap();
|
||||
let reader = Cursor::new(input);
|
||||
|
||||
let result = collector.consume_job_output("expected_job_id", reader);
|
||||
assert!(result.is_err());
|
||||
assert!(result.unwrap_err().to_string().contains("Job ID mismatch"));
|
||||
}
|
||||
}
|
||||
File diff suppressed because one or more lines are too long
Loading…
Reference in a new issue