rivet_logger/processors/
psr_log_message.rs

1use std::collections::BTreeMap;
2
3use time::macros::format_description;
4
5use crate::logger::{BoxError, LogRecord, LogValue, Processor};
6
7pub struct PsrLogMessage {
8    date_format: Option<String>,
9    remove_used_context_fields: bool,
10}
11
12impl PsrLogMessage {
13    pub const SIMPLE_DATE: &str =
14        "[year]-[month]-[day]T[hour]:[minute]:[second].[subsecond digits:6][offset_hour sign:mandatory]:[offset_minute]";
15
16    pub fn new(date_format: Option<String>, remove_used_context_fields: bool) -> Self {
17        Self {
18            date_format,
19            remove_used_context_fields,
20        }
21    }
22}
23
24impl Default for PsrLogMessage {
25    fn default() -> Self {
26        Self::new(None, false)
27    }
28}
29
30impl Processor for PsrLogMessage {
31    fn process(&self, mut record: LogRecord) -> Result<LogRecord, BoxError> {
32        if !record.message.contains('{') {
33            return Ok(record);
34        }
35
36        let mut replacements = BTreeMap::new();
37        let mut used_keys = Vec::new();
38
39        for (key, value) in &record.context {
40            let placeholder = format!("{{{key}}}");
41            if !record.message.contains(&placeholder) {
42                continue;
43            }
44
45            replacements.insert(
46                placeholder,
47                render_value(value, self.date_format.as_deref()),
48            );
49            used_keys.push(key.clone());
50        }
51
52        for (placeholder, value) in replacements {
53            record.message = record.message.replace(&placeholder, &value);
54        }
55
56        if self.remove_used_context_fields {
57            for key in used_keys {
58                record.context.remove(&key);
59            }
60        }
61
62        Ok(record)
63    }
64}
65
66fn render_value(value: &LogValue, date_format: Option<&str>) -> String {
67    match value {
68        LogValue::Null => String::new(),
69        LogValue::Bool(value) => value.to_string(),
70        LogValue::I64(value) => value.to_string(),
71        LogValue::U64(value) => value.to_string(),
72        LogValue::F64(value) => value.to_string(),
73        LogValue::String(value) => value.clone(),
74        LogValue::DateTime(value) => {
75            if let Some(pattern) = date_format {
76                if let Ok(parsed) = time::format_description::parse(pattern) {
77                    if let Ok(formatted) = value.format(&parsed) {
78                        return formatted;
79                    }
80                }
81            }
82
83            value
84                .format(format_description!(
85                    "[year]-[month]-[day]T[hour]:[minute]:[second].[subsecond digits:6][offset_hour sign:mandatory]:[offset_minute]"
86                ))
87                .unwrap_or_else(|_| value.to_string())
88        }
89        LogValue::Array(items) => render_array(items, date_format),
90        LogValue::Map(map) => render_map(map, date_format),
91        LogValue::Deferred(_) => "[deferred]".to_string(),
92    }
93}
94
95fn render_array(items: &[LogValue], date_format: Option<&str>) -> String {
96    let values = items
97        .iter()
98        .map(|item| render_json_value(item, date_format))
99        .collect::<Vec<_>>()
100        .join(",");
101
102    format!("[{values}]")
103}
104
105fn render_map(map: &BTreeMap<String, LogValue>, date_format: Option<&str>) -> String {
106    let values = map
107        .iter()
108        .map(|(key, value)| {
109            let key = escape_json_string(key);
110            let value = render_json_value(value, date_format);
111            format!("\"{key}\":{value}")
112        })
113        .collect::<Vec<_>>()
114        .join(",");
115
116    format!("{{{values}}}")
117}
118
119fn render_json_value(value: &LogValue, date_format: Option<&str>) -> String {
120    match value {
121        LogValue::Null => "null".to_string(),
122        LogValue::Bool(value) => value.to_string(),
123        LogValue::I64(value) => value.to_string(),
124        LogValue::U64(value) => value.to_string(),
125        LogValue::F64(value) => value.to_string(),
126        LogValue::String(value) => format!("\"{}\"", escape_json_string(value)),
127        LogValue::DateTime(_) => format!("\"{}\"", render_value(value, date_format)),
128        LogValue::Array(values) => render_array(values, date_format),
129        LogValue::Map(values) => render_map(values, date_format),
130        LogValue::Deferred(_) => "\"[deferred]\"".to_string(),
131    }
132}
133
134fn escape_json_string(value: &str) -> String {
135    value
136        .replace('\\', "\\\\")
137        .replace('"', "\\\"")
138        .replace('\n', "\\n")
139}
140
141#[cfg(test)]
142mod tests {
143    use std::collections::BTreeMap;
144
145    use crate::logger::{Context, Level};
146
147    use super::*;
148
149    fn record(message: &str, context: Context) -> LogRecord {
150        LogRecord {
151            datetime: time::OffsetDateTime::now_utc(),
152            channel: "test".to_string(),
153            level: Level::Info,
154            message: message.to_string(),
155            context,
156            extra: BTreeMap::new(),
157        }
158    }
159
160    #[test]
161    fn interpolates_placeholders() {
162        let mut context = BTreeMap::new();
163        context.insert("user".to_string(), LogValue::from("joseph"));
164        let processor = PsrLogMessage::default();
165
166        let processed = processor
167            .process(record("hello {user}", context))
168            .expect("processor should interpolate");
169
170        assert_eq!(processed.message, "hello joseph");
171    }
172
173    #[test]
174    fn removes_used_context_when_enabled() {
175        let mut context = BTreeMap::new();
176        context.insert("user".to_string(), LogValue::from("joseph"));
177        let processor = PsrLogMessage::new(None, true);
178
179        let processed = processor
180            .process(record("hello {user}", context))
181            .expect("processor should interpolate");
182
183        assert!(!processed.context.contains_key("user"));
184    }
185}
rivet_logger/processors/psr_log_message.rs

rivet_logger/processors/
psr_log_message.rs
