Add support for Key-Value pairs (#362)

ellttBen · estk · web-flow · commit 945a2bbc3e2d · 2025-05-25T22:58:24.000-07:00
* Add log::kv support

* Update log version

* Switch to a direct deserialize implementation for json kv attributes encoding. Factor pattern parsing code between MDC and key_value targets

---------

Co-authored-by: estk &lt;esims89@gmail.com&gt;
diff --git a/Cargo.toml b/Cargo.toml
@@ -48,6 +48,7 @@ console_writer = ["ansi_writer", "libc", "winapi"]
 simple_writer = []
 threshold_filter = []
 background_rotation = []
+log_kv = ["log/kv"]
 
 all_components = [
     "console_appender",
@@ -78,7 +79,7 @@ chrono = { version = "0.4.23", optional = true, features = [
 flate2 = { version = "1.0", optional = true }
 fnv = "1.0"
 humantime = { version = "2.1", optional = true }
-log = { version = "0.4.20", features = ["std"] }
+log = { version = "0.4.21", features = ["std"] }
 log-mdc = { version = "0.1", optional = true }
 serde = { version = "1.0.196", optional = true, features = ["derive"] }
 serde-value = { version = "0.7", optional = true }
diff --git a/src/encode/json.rs b/src/encode/json.rs
@@ -24,6 +24,11 @@
 //!     }
 //! }
 //! ```
+//! If the `log_kv` feature is enabled, an additional `attributes` field will
+//! contain a map of the record's [log::kv][log_kv] structured logging
+//! attributes.
+//!
+//! [log_kv]: https://docs.rs/log/latest/log/kv/index.html
 
 use chrono::{
     format::{DelayedFormat, Fixed, Item},
@@ -76,6 +81,8 @@ impl JsonEncoder {
             thread: thread.name(),
             thread_id: thread_id::get(),
             mdc: Mdc,
+            #[cfg(feature = "log_kv")]
+            attributes: kv::Attributes(record.key_values()),
         };
         message.serialize(&mut serde_json::Serializer::new(&mut *w))?;
         w.write_all(NEWLINE.as_bytes())?;
@@ -106,6 +113,8 @@ struct Message<'a> {
     thread: Option<&'a str>,
     thread_id: usize,
     mdc: Mdc,
+    #[cfg(feature = "log_kv")]
+    attributes: kv::Attributes<'a>,
 }
 
 fn ser_display<T, S>(v: &T, s: S) -> Result<S::Ok, S::Error>
@@ -162,6 +171,40 @@ impl Deserialize for JsonEncoderDeserializer {
         Ok(Box::<JsonEncoder>::default())
     }
 }
+#[cfg(feature = "log_kv")]
+mod kv {
+    use log::kv::VisitSource;
+    use serde::ser::{self, Error, SerializeMap};
+
+    pub(crate) struct Attributes<'a>(pub &'a dyn log::kv::Source);
+
+    pub(crate) struct SerializerVisitor<T: ser::SerializeMap>(pub T);
+
+    impl<'kvs, T: ser::SerializeMap> VisitSource<'kvs> for SerializerVisitor<T> {
+        fn visit_pair(
+            &mut self,
+            key: log::kv::Key<'kvs>,
+            value: log::kv::Value<'kvs>,
+        ) -> Result<(), log::kv::Error> {
+            self.0
+                .serialize_entry(key.as_str(), &value.to_string())
+                .map_err(|e| log::kv::Error::boxed(e.to_string()))?;
+            Ok(())
+        }
+    }
+
+    impl<'a> ser::Serialize for Attributes<'a> {
+        fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+        where
+            S: serde::Serializer,
+        {
+            let map = serializer.serialize_map(Some(self.0.count()))?;
+            let mut visitor = SerializerVisitor(map);
+            self.0.visit(&mut visitor).map_err(S::Error::custom)?;
+            visitor.0.end()
+        }
+    }
+}
 
 #[cfg(test)]
 #[cfg(feature = "simple_writer")]
@@ -189,26 +232,35 @@ mod test {
 
         let encoder = JsonEncoder::new();
 
+        let mut record_builder = Record::builder();
+        record_builder
+            .level(level)
+            .target(target)
+            .module_path(Some(module_path))
+            .file(Some(file))
+            .line(Some(line));
+
+        #[cfg(feature = "log_kv")]
+        record_builder.key_values(&[("log_foo", "log_bar")]);
+
         let mut buf = vec![];
         encoder
             .encode_inner(
                 &mut SimpleWriter(&mut buf),
                 time,
-                &Record::builder()
-                    .level(level)
-                    .target(target)
-                    .module_path(Some(module_path))
-                    .file(Some(file))
-                    .line(Some(line))
-                    .args(format_args!("{}", message))
-                    .build(),
+                &record_builder.args(format_args!("{}", message)).build(),
             )
             .unwrap();
 
+        #[cfg(feature = "log_kv")]
+        let expected_attributes = ",\"attributes\":{\"log_foo\":\"log_bar\"}";
+        #[cfg(not(feature = "log_kv"))]
+        let expected_attributes = "";
+
         let expected = format!(
             "{{\"time\":\"{}\",\"level\":\"{}\",\"message\":\"{}\",\"module_path\":\"{}\",\
             \"file\":\"{}\",\"line\":{},\"target\":\"{}\",\
-            \"thread\":\"{}\",\"thread_id\":{},\"mdc\":{{\"foo\":\"bar\"}}}}",
+            \"thread\":\"{}\",\"thread_id\":{},\"mdc\":{{\"foo\":\"bar\"}}{}}}",
             time.to_rfc3339(),
             level,
             message,
@@ -218,6 +270,7 @@ mod test {
             target,
             thread,
             thread_id::get(),
+            expected_attributes
         );
         assert_eq!(expected, String::from_utf8(buf).unwrap().trim());
     }
diff --git a/src/encode/pattern/mod.rs b/src/encode/pattern/mod.rs
@@ -73,6 +73,14 @@
 //!     defaults to the empty string.
 //!     * `{X(user_id)}` - `123e4567-e89b-12d3-a456-426655440000`
 //!     * `{X(nonexistent_key)(no mapping)}` - `no mapping`
+//! * `K`, `key_value` - A value from a [log::kv][log_kv] structured logging
+//!     record attributes. The first argument specifies the key, and the second
+//!     argument specifies the default value if the key is not present in the
+//!     log record's attributes. The second argument is optional, and defaults
+//!     to the empty string. This formatter requires the `log_kv` feature to be
+//!     enabled.
+//!     * `{K(user_id)}` - `123e4567-e89b-12d3-a456-426655440000`
+//!     * `{K(nonexistent_key)(no mapping)}` - `no mapping`
 //! * An "unnamed" formatter simply formats its argument, applying the format
 //!     specification.
 //!     * `{({l} {m})}` - `INFO hello`
@@ -120,6 +128,7 @@
 //! level `DEBUG` will be truncated to `DEBUG hello, wo`.
 //!
 //! [MDC]: https://crates.io/crates/log-mdc
+//! [log_kv]: https://docs.rs/log/latest/log/kv/index.html
 
 use chrono::{Local, Utc};
 use derivative::Derivative;
@@ -135,6 +144,8 @@ use crate::encode::{
 #[cfg(feature = "config_parsing")]
 use crate::config::{Deserialize, Deserializers};
 
+use self::parser::Formatter;
+
 mod parser;
 
 thread_local!(
@@ -496,46 +507,25 @@ impl<'a> From<Piece<'a>> for Chunk {
                 "P" | "pid" => no_args(&formatter.args, parameters, FormattedChunk::ProcessId),
                 "i" | "tid" => no_args(&formatter.args, parameters, FormattedChunk::SystemThreadId),
                 "t" | "target" => no_args(&formatter.args, parameters, FormattedChunk::Target),
-                "X" | "mdc" => {
-                    if formatter.args.len() > 2 {
-                        return Chunk::Error("expected at most two arguments".to_owned());
-                    }
-
-                    let key = match formatter.args.first() {
-                        Some(arg) => {
-                            if let Some(arg) = arg.first() {
-                                match arg {
-                                    Piece::Text(key) => key.to_owned(),
-                                    Piece::Error(ref e) => return Chunk::Error(e.clone()),
-                                    _ => return Chunk::Error("invalid MDC key".to_owned()),
-                                }
-                            } else {
-                                return Chunk::Error("invalid MDC key".to_owned());
-                            }
-                        }
-                        None => return Chunk::Error("missing MDC key".to_owned()),
-                    };
-
-                    let default = match formatter.args.get(1) {
-                        Some(arg) => {
-                            if let Some(arg) = arg.first() {
-                                match arg {
-                                    Piece::Text(key) => key.to_owned(),
-                                    Piece::Error(ref e) => return Chunk::Error(e.clone()),
-                                    _ => return Chunk::Error("invalid MDC default".to_owned()),
-                                }
-                            } else {
-                                return Chunk::Error("invalid MDC default".to_owned());
-                            }
-                        }
-                        None => "",
-                    };
-
-                    Chunk::Formatted {
-                        chunk: FormattedChunk::Mdc(key.into(), default.into()),
+                "X" | "mdc" => match kv_parsing(&formatter) {
+                    Err(e) => Chunk::Error(format!("MDC: {e}")),
+                    Ok((key, default)) => Chunk::Formatted {
+                        chunk: FormattedChunk::Mdc(key, default),
                         params: parameters,
-                    }
-                }
+                    },
+                },
+                #[cfg(feature = "log_kv")]
+                "K" | "key_value" => match kv_parsing(&formatter) {
+                    Err(e) => Chunk::Error(format!("key_value: {e}")),
+                    Ok((key, default)) => Chunk::Formatted {
+                        chunk: FormattedChunk::Kv(key, default),
+                        params: parameters,
+                    },
+                },
+                #[cfg(not(feature = "log_kv"))]
+                "K" | "key_value" => Chunk::Error(
+                    "The log_kv feature is required to parse the key_value argument".to_owned(),
+                ),
                 "" => {
                     if formatter.args.len() != 1 {
                         return Chunk::Error("expected exactly one argument".to_owned());
@@ -568,6 +558,43 @@ fn no_args(arg: &[Vec<Piece>], params: Parameters, chunk: FormattedChunk) -> Chu
     }
 }
 
+fn kv_parsing<'a>(formatter: &'a Formatter) -> Result<(String, String), &'a str> {
+    if formatter.args.len() > 2 {
+        return Err("expected at most two arguments");
+    }
+
+    let key = match formatter.args.first() {
+        Some(arg) => {
+            if let Some(arg) = arg.first() {
+                match arg {
+                    Piece::Text(key) => key.to_owned(),
+                    Piece::Error(ref e) => return Err(e),
+                    _ => return Err("invalid key"),
+                }
+            } else {
+                return Err("invalid key");
+            }
+        }
+        None => return Err("missing key"),
+    };
+
+    let default = match formatter.args.get(1) {
+        Some(arg) => {
+            if let Some(arg) = arg.first() {
+                match arg {
+                    Piece::Text(key) => key.to_owned(),
+                    Piece::Error(ref e) => return Err(e),
+                    _ => return Err("invalid default"),
+                }
+            } else {
+                return Err("invalid default");
+            }
+        }
+        None => "",
+    };
+    Ok((key.into(), default.into()))
+}
+
 #[derive(Clone, Eq, PartialEq, Hash, Debug)]
 enum Timezone {
     Utc,
@@ -593,6 +620,8 @@ enum FormattedChunk {
     Debug(Vec<Chunk>),
     Release(Vec<Chunk>),
     Mdc(String, String),
+    #[cfg(feature = "log_kv")]
+    Kv(String, String),
 }
 
 impl FormattedChunk {
@@ -666,6 +695,15 @@ impl FormattedChunk {
             FormattedChunk::Mdc(ref key, ref default) => {
                 log_mdc::get(key, |v| write!(w, "{}", v.unwrap_or(default)))
             }
+            #[cfg(feature = "log_kv")]
+            FormattedChunk::Kv(ref key, ref default) => {
+                use log::kv::ToKey;
+                if let Some(v) = record.key_values().get(key.to_key()) {
+                    write!(w, "{v}")
+                } else {
+                    write!(w, "{default}")
+                }
+            }
         }
     }
 }
@@ -751,6 +789,8 @@ mod tests {
     #[cfg(feature = "simple_writer")]
     use std::thread;
 
+    #[cfg(feature = "log_kv")]
+    use super::Parser;
     use super::{Chunk, PatternEncoder};
     #[cfg(feature = "simple_writer")]
     use crate::encode::writer::simple::SimpleWriter;
@@ -1031,6 +1071,71 @@ mod tests {
         assert_eq!(buf, b"missing value");
     }
 
+    #[test]
+    #[cfg(all(feature = "simple_writer", feature = "log_kv"))]
+    fn test_kv() {
+        let pw = PatternEncoder::new("{K(user_id)}");
+        let kv = [("user_id", "kv value")];
+
+        let mut buf = vec![];
+        pw.encode(
+            &mut SimpleWriter(&mut buf),
+            &Record::builder().key_values(&kv).build(),
+        )
+        .unwrap();
+
+        assert_eq!(buf, b"kv value");
+    }
+
+    #[test]
+    #[cfg(all(feature = "simple_writer", feature = "log_kv"))]
+    fn test_kv_missing_default() {
+        let pw = PatternEncoder::new("{K(user_id)}");
+
+        let mut buf = vec![];
+        pw.encode(&mut SimpleWriter(&mut buf), &Record::builder().build())
+            .unwrap();
+
+        assert_eq!(buf, b"");
+    }
+
+    #[test]
+    #[cfg(all(feature = "simple_writer", feature = "log_kv"))]
+    fn test_kv_missing_custom() {
+        let pw = PatternEncoder::new("{K(user_id)(missing value)}");
+
+        let mut buf = vec![];
+        pw.encode(&mut SimpleWriter(&mut buf), &Record::builder().build())
+            .unwrap();
+
+        assert_eq!(buf, b"missing value");
+    }
+
+    #[test]
+    #[cfg(feature = "log_kv")]
+    fn test_kv_from_piece_to_chunk() {
+        let tests = vec![
+            (
+                "[{K(user_id)(foobar)(test):<5.5}]",
+                "expected at most two arguments",
+            ),
+            ("[{K({l user_id):<5.5}]", "expected '}'"),
+            ("[{K({l} user_id):<5.5}]", "key_value: invalid key"),
+            ("[{K:<5.5}]", "key_value: missing key"),
+            ("[{K(user_id)({l):<5.5}]", "expected '}'"),
+            ("[{K(user_id)({l}):<5.5}]", "key_value: invalid default"),
+            ("[{K(user_id)():<5.5} {M}]", "key_value: invalid default"),
+        ];
+
+        for (pattern, error_msg) in tests {
+            let chunks: Vec<Chunk> = Parser::new(pattern).map(From::from).collect();
+            match chunks.get(1).unwrap() {
+                Chunk::Error(err) => assert!(err.contains(error_msg)),
+                _ => panic!(),
+            }
+        }
+    }
+
     #[test]
     #[cfg(feature = "simple_writer")]
     fn debug_release() {
