scribe_rust/
lib.rs

1use std::sync::Arc;
2use tracing::{Event, Level as TracingLevel, Subscriber};
3use tracing_subscriber::{layer::Context, Layer};
4
5const FG_RED: &str = "\x1b[31m";
6const FG_YELLOW: &str = "\x1b[33m";
7const FG_BLUE: &str = "\x1b[34m";
8const FG_GREEN: &str = "\x1b[32m";
9const FG_GRAY: &str = "\x1b[90m";
10const FG_RESET: &str = "\x1b[0m";
11
12pub enum Color {
13    Red,
14    Yellow,
15    Blue,
16    Green,
17    Gray,
18}
19
20#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
21pub enum Level {
22    Trace,
23    Debug,
24    Info,
25    Warn,
26    Error,
27}
28
29impl std::fmt::Display for Level {
30    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
31        match self {
32            Level::Trace => write!(f, "TRACE"),
33            Level::Debug => write!(f, "DEBUG"),
34            Level::Info => write!(f, "INFO"),
35            Level::Warn => write!(f, "WARN"),
36            Level::Error => write!(f, "ERROR"),
37        }
38    }
39}
40
41impl std::str::FromStr for Level {
42    type Err = String;
43
44    fn from_str(s: &str) -> Result<Self, Self::Err> {
45        match s.to_uppercase().as_str() {
46            "TRACE" => Ok(Level::Trace),
47            "DEBUG" => Ok(Level::Debug),
48            "INFO" => Ok(Level::Info),
49            "WARN" => Ok(Level::Warn),
50            "ERROR" => Ok(Level::Error),
51            _ => Err(format!("Unknown log level: {}", s)),
52        }
53    }
54}
55
56#[derive(Debug, Clone)]
57pub struct Logger {
58    pub level: Level,
59}
60
61impl Logger {
62    pub fn new(level: Level) -> Logger {
63        Logger { level }
64    }
65
66    pub fn default() -> Arc<Logger> {
67        let level = std::env::var("LOG_LEVEL")
68            .map(|level| level.parse::<Level>().unwrap_or(Level::Info))
69            .unwrap_or(Level::Info);
70
71        Arc::new(Logger { level })
72    }
73
74    pub fn log(&self, level: Level, message: &str) {
75        if level >= self.level {
76            println!(
77                "[{}{}{}] {}",
78                match level {
79                    Level::Trace => FG_GRAY,
80                    Level::Debug => FG_BLUE,
81                    Level::Info => FG_GREEN,
82                    Level::Warn => FG_YELLOW,
83                    Level::Error => FG_RED,
84                },
85                level,
86                FG_RESET,
87                message
88            );
89        }
90    }
91
92    pub fn trace(&self, message: &str) {
93        self.log(Level::Trace, message);
94    }
95
96    pub fn debug(&self, message: &str) {
97        self.log(Level::Debug, message);
98    }
99
100    pub fn info(&self, message: &str) {
101        self.log(Level::Info, message);
102    }
103
104    pub fn warn(&self, message: &str) {
105        self.log(Level::Warn, message);
106    }
107
108    pub fn error(&self, message: &str) {
109        self.log(Level::Error, message);
110    }
111}
112
113use serde_json::{json, Value};
114use std::collections::BTreeMap;
115use tracing::field::{Field, Visit};
116
117struct FieldVisitor {
118    fields: BTreeMap<String, Value>,
119}
120
121impl FieldVisitor {
122    fn new() -> Self {
123        FieldVisitor {
124            fields: BTreeMap::new(),
125        }
126    }
127}
128
129impl Visit for FieldVisitor {
130    fn record_bool(&mut self, field: &Field, value: bool) {
131        self.fields
132            .insert(field.name().to_string(), Value::Bool(value));
133    }
134
135    fn record_i64(&mut self, field: &Field, value: i64) {
136        self.fields
137            .insert(field.name().to_string(), Value::Number(value.into()));
138    }
139
140    fn record_u64(&mut self, field: &Field, value: u64) {
141        self.fields
142            .insert(field.name().to_string(), Value::Number(value.into()));
143    }
144
145    fn record_str(&mut self, field: &Field, value: &str) {
146        self.fields
147            .insert(field.name().to_string(), Value::String(value.to_string()));
148    }
149
150    fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) {
151        self.fields.insert(
152            field.name().to_string(),
153            Value::String(format!("{:?}", value)),
154        );
155    }
156}
157
158pub struct CustomLoggerLayer {
159    pub logger: Arc<Logger>,
160}
161
162impl<S> Layer<S> for CustomLoggerLayer
163where
164    S: Subscriber,
165{
166    fn on_event(&self, event: &Event<'_>, _ctx: Context<'_, S>) {
167        let level = match *event.metadata().level() {
168            TracingLevel::TRACE => Level::Trace,
169            TracingLevel::DEBUG => Level::Debug,
170            TracingLevel::INFO => Level::Info,
171            TracingLevel::WARN => Level::Warn,
172            TracingLevel::ERROR => Level::Error,
173        };
174
175        if level >= self.logger.level {
176            // Initialize the visitor to capture event fields
177            let mut visitor = FieldVisitor::new();
178            event.record(&mut visitor);
179
180            // Attempt to retrieve the "message" field
181            let message = visitor
182                .fields
183                .get("message")
184                .and_then(|v| v.as_str())
185                .unwrap_or(event.metadata().name());
186
187            // Log the message without the JSON fields
188            self.logger.log(level, message);
189        }
190    }
191}