enrichment/
vrl_util.rs

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

enrichment/
vrl_util.rs
