1use std::{collections::btree_map, fmt::Write as _, iter, slice, sync::LazyLock};
2
3use regex::Regex;
4use serde::{Serialize, Serializer};
5use vrl::path::PathPrefix;
6
7use crate::event::{KeyString, ObjectMap, Value};
8
9static IS_VALID_PATH_SEGMENT: LazyLock<Regex> =
10 LazyLock::new(|| Regex::new(r"^[a-zA-Z0-9_]+$").unwrap());
11
12pub fn all_fields(fields: &ObjectMap) -> FieldsIter {
15 FieldsIter::new(None, fields, true)
16}
17
18pub fn all_fields_unquoted(fields: &ObjectMap) -> FieldsIter {
21 FieldsIter::new(None, fields, false)
22}
23
24pub fn all_metadata_fields(fields: &ObjectMap) -> FieldsIter {
27 FieldsIter::new(Some(PathPrefix::Metadata), fields, true)
28}
29
30pub fn all_fields_non_object_root(value: &Value) -> FieldsIter {
32 FieldsIter::non_object(value)
33}
34
35pub fn all_fields_skip_array_elements(fields: &ObjectMap) -> FieldsIter {
38 FieldsIter::new_with_skip_array_elements(fields)
39}
40
41#[derive(Clone, Debug)]
42enum LeafIter<'a> {
43 Root((&'a Value, bool)),
44 Map(btree_map::Iter<'a, KeyString, Value>),
45 Array(iter::Enumerate<slice::Iter<'a, Value>>),
46}
47
48#[derive(Clone, Copy, Debug)]
49enum PathComponent<'a> {
50 Key(&'a KeyString),
51 Index(usize),
52}
53
54#[derive(Clone)]
58pub struct FieldsIter<'a> {
59 path_prefix: Option<PathPrefix>,
61 stack: Vec<LeafIter<'a>>,
63 path: Vec<PathComponent<'a>>,
65 skip_array_elements: bool,
67 quote_invalid_fields: bool,
69}
70
71impl<'a> FieldsIter<'a> {
72 fn new(
73 path_prefix: Option<PathPrefix>,
74 fields: &'a ObjectMap,
75 quote_invalid_fields: bool,
76 ) -> FieldsIter<'a> {
77 FieldsIter {
78 path_prefix,
79 stack: vec![LeafIter::Map(fields.iter())],
80 path: vec![],
81 skip_array_elements: false,
82 quote_invalid_fields,
83 }
84 }
85
86 fn non_object(value: &'a Value) -> FieldsIter<'a> {
89 FieldsIter {
90 path_prefix: None,
91 stack: vec![LeafIter::Root((value, false))],
92 path: vec![],
93 skip_array_elements: false,
94 quote_invalid_fields: true,
95 }
96 }
97
98 fn new_with_skip_array_elements(fields: &'a ObjectMap) -> FieldsIter<'a> {
99 FieldsIter {
100 path_prefix: None,
101 stack: vec![LeafIter::Map(fields.iter())],
102 path: vec![],
103 skip_array_elements: true,
104 quote_invalid_fields: true,
105 }
106 }
107
108 fn push(&mut self, value: &'a Value, component: PathComponent<'a>) -> Option<&'a Value> {
109 match value {
110 Value::Object(map) if !map.is_empty() => {
111 self.stack.push(LeafIter::Map(map.iter()));
112 self.path.push(component);
113 None
114 }
115 Value::Array(array) if !array.is_empty() => {
116 if self.skip_array_elements {
117 Some(value)
118 } else {
119 self.stack.push(LeafIter::Array(array.iter().enumerate()));
120 self.path.push(component);
121 None
122 }
123 }
124 _ => Some(value),
125 }
126 }
127
128 fn pop(&mut self) {
129 self.stack.pop();
130 self.path.pop();
131 }
132
133 fn make_path(&mut self, component: PathComponent<'a>) -> KeyString {
134 let mut res = match self.path_prefix {
135 None => String::new(),
136 Some(prefix) => match prefix {
137 PathPrefix::Event => String::from("."),
138 PathPrefix::Metadata => String::from("%"),
139 },
140 };
141 let mut path_iter = self.path.iter().chain(iter::once(&component)).peekable();
142 loop {
143 match path_iter.next() {
144 None => break res.into(),
145 Some(PathComponent::Key(key)) => {
146 if self.quote_invalid_fields && !IS_VALID_PATH_SEGMENT.is_match(key) {
147 write!(res, "\"{key}\"").expect("write to String never fails");
148 } else {
149 res.push_str(key);
150 }
151 }
152 Some(PathComponent::Index(index)) => {
153 write!(res, "[{index}]").expect("write to String never fails");
154 }
155 }
156 if let Some(PathComponent::Key(_)) = path_iter.peek() {
157 res.push('.');
158 }
159 }
160 }
161}
162
163impl<'a> Iterator for FieldsIter<'a> {
164 type Item = (KeyString, &'a Value);
165
166 fn next(&mut self) -> Option<Self::Item> {
167 loop {
168 match self.stack.last_mut() {
169 None => return None,
170 Some(LeafIter::Map(map_iter)) => match map_iter.next() {
171 None => self.pop(),
172 Some((key, value)) => {
173 if let Some(scalar_value) = self.push(value, PathComponent::Key(key)) {
174 return Some((self.make_path(PathComponent::Key(key)), scalar_value));
175 }
176 }
177 },
178 Some(LeafIter::Array(array_iter)) => match array_iter.next() {
179 None => self.pop(),
180 Some((index, value)) => {
181 if let Some(scalar_value) = self.push(value, PathComponent::Index(index)) {
182 return Some((
183 self.make_path(PathComponent::Index(index)),
184 scalar_value,
185 ));
186 }
187 }
188 },
189 Some(LeafIter::Root((value, visited))) => {
190 let result = (!*visited).then(|| ("message".into(), *value));
191 *visited = true;
192 break result;
193 }
194 }
195 }
196 }
197}
198
199impl Serialize for FieldsIter<'_> {
200 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
201 where
202 S: Serializer,
203 {
204 serializer.collect_map(self.clone())
205 }
206}
207
208#[cfg(test)]
209mod test {
210 use serde_json::json;
211 use similar_asserts::assert_eq;
212
213 use super::{super::test::fields_from_json, *};
214
215 #[test]
216 fn keys_simple() {
217 let fields = fields_from_json(json!({
218 "field2": 3,
219 "field1": 4,
220 "field3": 5
221 }));
222 let expected: Vec<_> = vec![
223 ("field1", &Value::Integer(4)),
224 ("field2", &Value::Integer(3)),
225 ("field3", &Value::Integer(5)),
226 ]
227 .into_iter()
228 .map(|(k, v)| (k.into(), v))
229 .collect();
230
231 let collected: Vec<_> = all_fields(&fields).collect();
232 assert_eq!(collected, expected);
233 }
234
235 fn special_fields() -> ObjectMap {
236 fields_from_json(json!({
237 "a-b": 1,
238 "a*b": 2,
239 "a b": 3,
240 ".a .b*": 4,
241 "\"a\"": 5,
242 }))
243 }
244
245 #[test]
246 fn keys_special_quoted() {
247 let fields = special_fields();
248 let mut collected: Vec<_> = all_fields(&fields).collect();
249 collected.sort_by(|(a, _), (b, _)| a.cmp(b));
250
251 let mut expected: Vec<(KeyString, &Value)> = vec![
252 ("\"a-b\"", &Value::Integer(1)),
253 ("\"a*b\"", &Value::Integer(2)),
254 ("\"a b\"", &Value::Integer(3)),
255 ("\".a .b*\"", &Value::Integer(4)),
256 ("\"\"a\"\"", &Value::Integer(5)),
257 ]
258 .into_iter()
259 .map(|(k, v)| (k.into(), v))
260 .collect();
261 expected.sort_by(|(a, _), (b, _)| a[1..].cmp(&b[1..]));
263
264 assert_eq!(collected, expected);
265 }
266
267 #[test]
268 fn keys_special_unquoted() {
269 let fields = special_fields();
270 let mut collected: Vec<_> = all_fields_unquoted(&fields).collect();
271 collected.sort_by(|(a, _), (b, _)| a.cmp(b));
272
273 let mut expected: Vec<(KeyString, &Value)> = vec![
274 ("a-b", &Value::Integer(1)),
275 ("a*b", &Value::Integer(2)),
276 ("a b", &Value::Integer(3)),
277 (".a .b*", &Value::Integer(4)),
278 ("\"a\"", &Value::Integer(5)),
279 ]
280 .into_iter()
281 .map(|(k, v)| (k.into(), v))
282 .collect();
283 expected.sort_by(|(a, _), (b, _)| a.cmp(b));
284
285 assert_eq!(collected, expected);
286 }
287
288 #[test]
289 fn metadata_keys_simple() {
290 let fields = fields_from_json(json!({
291 "field_1": 1,
292 "field_0": 0,
293 "field_2": 2
294 }));
295 let expected: Vec<_> = vec![
296 ("%field_0", &Value::Integer(0)),
297 ("%field_1", &Value::Integer(1)),
298 ("%field_2", &Value::Integer(2)),
299 ]
300 .into_iter()
301 .map(|(k, v)| (k.into(), v))
302 .collect();
303
304 let collected: Vec<_> = all_metadata_fields(&fields).collect();
305 assert_eq!(collected, expected);
306 }
307
308 fn nested_fields() -> ObjectMap {
309 fields_from_json(json!({
310 "a": {
311 "b": {
312 "c": 5
313 },
314 "a": 4,
315 "array": [null, 3, {
316 "x": 1
317 }, [2]]
318 },
319 "a.b.c": 6,
320 "d": {},
321 "e": [],
322 }))
323 }
324
325 #[test]
326 fn keys_nested_quoted() {
327 let fields = nested_fields();
328 let expected: Vec<_> = vec![
329 ("a.a", Value::Integer(4)),
330 ("a.array[0]", Value::Null),
331 ("a.array[1]", Value::Integer(3)),
332 ("a.array[2].x", Value::Integer(1)),
333 ("a.array[3][0]", Value::Integer(2)),
334 ("a.b.c", Value::Integer(5)),
335 ("\"a.b.c\"", Value::Integer(6)),
336 ("d", Value::Object(ObjectMap::new())),
337 ("e", Value::Array(Vec::new())),
338 ]
339 .into_iter()
340 .map(|(k, v)| (k.into(), v))
341 .collect();
342
343 let collected: Vec<_> = all_fields(&fields).map(|(k, v)| (k, v.clone())).collect();
344 assert_eq!(collected, expected);
345 }
346
347 #[test]
348 fn keys_nested_unquoted() {
349 let fields = nested_fields();
350 let expected: Vec<_> = vec![
351 ("a.a", Value::Integer(4)),
352 ("a.array[0]", Value::Null),
353 ("a.array[1]", Value::Integer(3)),
354 ("a.array[2].x", Value::Integer(1)),
355 ("a.array[3][0]", Value::Integer(2)),
356 ("a.b.c", Value::Integer(5)),
357 ("a.b.c", Value::Integer(6)),
358 ("d", Value::Object(ObjectMap::new())),
359 ("e", Value::Array(Vec::new())),
360 ]
361 .into_iter()
362 .map(|(k, v)| (k.into(), v))
363 .collect();
364
365 let collected: Vec<_> = all_fields_unquoted(&fields)
366 .map(|(k, v)| (k, v.clone()))
367 .collect();
368 assert_eq!(collected, expected);
369 }
370
371 #[test]
372 fn test_non_object_root() {
373 let value = Value::Integer(3);
374 let collected: Vec<_> = all_fields_non_object_root(&value)
375 .map(|(k, v)| (k.into(), v.clone()))
376 .collect();
377 assert_eq!(collected, vec![("message".to_owned(), value)]);
378 }
379}