enrichment/
vrl_util.rs

1//! Utilities shared between both VRL functions.
2use std::collections::BTreeMap;
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}
106
107#[allow(clippy::result_large_err)]
108pub(crate) fn is_case_sensitive(
109    arguments: &ArgumentList,
110    state: &TypeState,
111) -> Result<Case, function::Error> {
112    Ok(arguments
113        .optional_literal("case_sensitive", state)?
114        .map(|value| {
115            let case_sensitive = value
116                .as_boolean()
117                .expect("case_sensitive should be boolean"); // This will have been caught by the type checker.
118
119            if case_sensitive {
120                Case::Sensitive
121            } else {
122                Case::Insensitive
123            }
124        })
125        .unwrap_or(Case::Sensitive))
126}
127
128#[cfg(test)]
129mod tests {
130    use std::sync::{Arc, Mutex};
131
132    use chrono::{TimeZone, Utc};
133
134    use super::*;
135    use crate::test_util;
136
137    #[test]
138    fn add_indexes() {
139        let mut registry = test_util::get_table_registry();
140        let conditions =
141            BTreeMap::from([("field".into(), expression::Literal::from("value").into())]);
142        let index = add_index(&mut registry, "dummy1", Case::Insensitive, &conditions).unwrap();
143
144        assert_eq!(IndexHandle(0), index);
145    }
146
147    #[test]
148    fn add_indexes_with_dates() {
149        let indexes = Arc::new(Mutex::new(Vec::new()));
150        let dummy = test_util::DummyEnrichmentTable::new_with_index(indexes.clone());
151
152        let mut registry =
153            test_util::get_table_registry_with_tables(vec![("dummy1".to_string(), dummy)]);
154
155        let conditions = BTreeMap::from([
156            ("field1".into(), (expression::Literal::from("value")).into()),
157            (
158                "field2".into(),
159                (expression::Container::new(expression::Variant::Object(
160                    BTreeMap::from([
161                        (
162                            "from".into(),
163                            (expression::Literal::from(
164                                Utc.with_ymd_and_hms(2015, 5, 15, 0, 0, 0)
165                                    .single()
166                                    .expect("invalid timestamp"),
167                            ))
168                            .into(),
169                        ),
170                        (
171                            "to".into(),
172                            (expression::Literal::from(
173                                Utc.with_ymd_and_hms(2015, 6, 15, 0, 0, 0)
174                                    .single()
175                                    .expect("invalid timestamp"),
176                            ))
177                            .into(),
178                        ),
179                    ])
180                    .into(),
181                )))
182                .into(),
183            ),
184        ]);
185
186        let index = add_index(&mut registry, "dummy1", Case::Sensitive, &conditions).unwrap();
187
188        assert_eq!(IndexHandle(0), index);
189
190        // Ensure only the exact match has been added as an index.
191        let indexes = indexes.lock().unwrap();
192        assert_eq!(vec![vec!["field1".to_string()]], *indexes);
193    }
194}
enrichment/vrl_util.rs

enrichment/
vrl_util.rs
