enrichment/
vrl_util.rs

1//! Utilities shared between both VRL functions.
2use std::{collections::BTreeMap, sync::LazyLock};
3
4use vrl::{
5    diagnostic::{Label, Span},
6    prelude::*,
7};
8
9use crate::{Case, Condition, IndexHandle, TableRegistry};
10
11#[derive(Debug)]
12pub enum Error {
13    TablesNotLoaded,
14}
15
16impl fmt::Display for Error {
17    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
18        match self {
19            Error::TablesNotLoaded => write!(f, "enrichment tables not loaded"),
20        }
21    }
22}
23
24impl std::error::Error for Error {}
25
26impl DiagnosticMessage for Error {
27    fn code(&self) -> usize {
28        111
29    }
30
31    fn labels(&self) -> Vec<Label> {
32        match self {
33            Error::TablesNotLoaded => {
34                vec![Label::primary(
35                    "enrichment table error: tables not loaded".to_string(),
36                    Span::default(),
37                )]
38            }
39        }
40    }
41}
42
43/// Evaluates the condition object to search the enrichment tables with.
44pub(crate) fn evaluate_condition(key: &str, value: Value) -> ExpressionResult<Condition<'_>> {
45    Ok(match value {
46        Value::Object(map) if map.contains_key("from") && map.contains_key("to") => {
47            Condition::BetweenDates {
48                field: key,
49                from: *map
50                    .get("from")
51                    .expect("should contain from")
52                    .as_timestamp()
53                    .ok_or("from in condition must be a timestamp")?,
54                to: *map
55                    .get("to")
56                    .expect("should contain to")
57                    .as_timestamp()
58                    .ok_or("to in condition must be a timestamp")?,
59            }
60        }
61        Value::Object(map) if map.contains_key("from") => Condition::FromDate {
62            field: key,
63            from: *map
64                .get("from")
65                .expect("should contain from")
66                .as_timestamp()
67                .ok_or("from in condition must be a timestamp")?,
68        },
69        Value::Object(map) if map.contains_key("to") => Condition::ToDate {
70            field: key,
71            to: *map
72                .get("to")
73                .expect("should contain to")
74                .as_timestamp()
75                .ok_or("to in condition must be a timestamp")?,
76        },
77        _ => Condition::Equals { field: key, value },
78    })
79}
80
81/// Add an index for the given condition to the given enrichment table.
82pub(crate) fn add_index(
83    registry: &mut TableRegistry,
84    tablename: &str,
85    case: Case,
86    condition: &BTreeMap<KeyString, expression::Expr>,
87) -> std::result::Result<IndexHandle, ExpressionError> {
88    let fields = condition
89        .iter()
90        .filter_map(|(field, value)| match value {
91            expression::Expr::Container(expression::Container {
92                variant: expression::Variant::Object(map),
93            }) if (map.contains_key("from") && map.contains_key("to"))
94                || map.contains_key("from")
95                || map.contains_key("to") =>
96            {
97                None
98            }
99            _ => Some(field.as_ref()),
100        })
101        .collect::<Vec<_>>();
102    let index = registry.add_index(tablename, case, &fields)?;
103
104    Ok(index)
105}
106pub(crate) static DEFAULT_CASE_SENSITIVE: LazyLock<Value> = LazyLock::new(|| Value::Boolean(true));
107
108#[allow(clippy::result_large_err)]
109pub(crate) fn is_case_sensitive(
110    arguments: &ArgumentList,
111    state: &TypeState,
112) -> Result<Case, function::Error> {
113    let case_sensitive = arguments
114        .optional_literal("case_sensitive", state)?
115        .unwrap_or_else(|| DEFAULT_CASE_SENSITIVE.clone())
116        .as_boolean()
117        .expect("case_sensitive should be boolean"); // This will have been caught by the type checker.
118
119    Ok(if case_sensitive {
120        Case::Sensitive
121    } else {
122        Case::Insensitive
123    })
124}
125
126#[cfg(test)]
127mod tests {
128    use std::sync::{Arc, Mutex};
129
130    use chrono::{TimeZone, Utc};
131
132    use super::*;
133    use crate::test_util;
134
135    #[test]
136    fn add_indexes() {
137        let mut registry = test_util::get_table_registry();
138        let conditions =
139            BTreeMap::from([("field".into(), expression::Literal::from("value").into())]);
140        let index = add_index(&mut registry, "dummy1", Case::Insensitive, &conditions).unwrap();
141
142        assert_eq!(IndexHandle(0), index);
143    }
144
145    #[test]
146    fn add_indexes_with_dates() {
147        let indexes = Arc::new(Mutex::new(Vec::new()));
148        let dummy = test_util::DummyEnrichmentTable::new_with_index(indexes.clone());
149
150        let mut registry =
151            test_util::get_table_registry_with_tables(vec![("dummy1".to_string(), dummy)]);
152
153        let conditions = BTreeMap::from([
154            ("field1".into(), (expression::Literal::from("value")).into()),
155            (
156                "field2".into(),
157                (expression::Container::new(expression::Variant::Object(
158                    BTreeMap::from([
159                        (
160                            "from".into(),
161                            (expression::Literal::from(
162                                Utc.with_ymd_and_hms(2015, 5, 15, 0, 0, 0)
163                                    .single()
164                                    .expect("invalid timestamp"),
165                            ))
166                            .into(),
167                        ),
168                        (
169                            "to".into(),
170                            (expression::Literal::from(
171                                Utc.with_ymd_and_hms(2015, 6, 15, 0, 0, 0)
172                                    .single()
173                                    .expect("invalid timestamp"),
174                            ))
175                            .into(),
176                        ),
177                    ])
178                    .into(),
179                )))
180                .into(),
181            ),
182        ]);
183
184        let index = add_index(&mut registry, "dummy1", Case::Sensitive, &conditions).unwrap();
185
186        assert_eq!(IndexHandle(0), index);
187
188        // Ensure only the exact match has been added as an index.
189        let indexes = indexes.lock().unwrap();
190        assert_eq!(vec![vec!["field1".to_string()]], *indexes);
191    }
192}
enrichment/vrl_util.rs

enrichment/
vrl_util.rs
