json_surf/
utils.rs

1use std::{thread::sleep, time::Duration, time::Instant};
2use std::collections::HashMap;
3use std::path::{Path, PathBuf};
4
5use rand::{Rng};
6use rand::distributions::Alphanumeric;
7
8use serde::{Serialize};
9use serde;
10use serde_value;
11use serde_value::Value;
12
13use tantivy::schema::{Schema, TextOptions, TEXT, IntOptions, STORED, SchemaBuilder};
14
15use crate::prelude::*;
16use crate::registry::SurferFieldTypes;
17
18/// Convert a JSON serializable struct as JSON
19pub(crate) fn as_value<T>(data: &T) -> Result<Value, IndexError>
20    where
21        T: Serialize,
22
23{
24    let result = serde_value::to_value(data).map_err(|e| {
25        IndexError::new(
26            "Unable to serialize data",
27            &e.to_string(),
28        )
29    })?;
30    Ok(result)
31}
32
33/// Get all field names
34pub fn field_names(data: &Value) -> Option<Vec<String>> {
35    if let Value::Map(kv) = data {
36        let keys = kv.keys();
37        let mut fields = Vec::with_capacity(keys.len());
38        for key in keys {
39            if let Value::String(name) = key {
40                fields.push(name.to_owned());
41            }
42        }
43        return Some(fields);
44    };
45    None
46}
47
48
49/// Store and index text by default
50fn resolve_text_option(key: &str, control: Option<&HashMap<String, Control>>) -> TextOptions {
51    let default = TEXT | STORED;
52
53    match control {
54        Some(c) => {
55            let x = c.get(key);
56            if x.is_none() {
57                return default;
58            }
59            let x = x.unwrap();
60            match x {
61                Control::ControlTextOptions(opt) => {
62                    let option = opt.clone();
63                    option
64                }
65                _ => default
66            }
67        }
68        None => default
69    }
70}
71
72/// Store and index numbers by default
73fn resolve_number_option(key: &str, control: Option<&HashMap<String, Control>>) -> IntOptions {
74    let default = IntOptions::default();
75    let default = default.set_indexed();
76    let default = default.set_stored();
77    match control {
78        Some(c) => {
79            let x = c.get(key);
80            if x.is_none() {
81                return default;
82            }
83            let x = x.unwrap();
84            match x {
85                Control::ControlIntOptions(opt) => {
86                    let option = opt.clone();
87                    option
88                }
89                _ => default
90            }
91        }
92        None => default
93    }
94}
95
96/// Join to path
97pub fn join(head: &str, tail: &str) -> Option<String> {
98    let head = Path::new(head);
99    let tail = Path::new(tail);
100    let path = head.join(tail);
101    let path = path.to_str();
102    match path {
103        Some(p) => Some(p.to_string()),
104        None => None
105    }
106}
107
108/// Maps flat JSON structures
109pub(crate) fn as_schema_builder<T: Serialize>(payload: &T, control: Option<&HashMap<String, Control>>) -> Result<(SchemaBuilder, HashMap<String, SurferFieldTypes>), IndexError> {
110    let value = as_value(payload)?;
111    let data = serde_json::to_string(payload).unwrap();
112    let kv = match &value {
113        Value::Map(kv) => kv,
114        _ => {
115            return Err(IndexError::new(
116                "Unable to create schema",
117                "Expected BTree Map", )
118            );
119        }
120    };
121    let mut names = Vec::new();
122    let keys = kv.keys();
123    for key in keys {
124        match key {
125            Value::String(k) => names.push(k.clone()),
126            _ => {
127                return Err(IndexError::new(
128                    "Unable to create schema",
129                    "keys were not string", )
130                );
131            }
132        }
133    };
134
135    let mut field_names = HashMap::new();
136    let mut indexes = Vec::new();
137    let mut field_type_mappings = HashMap::<String, SurferFieldTypes>::new();
138
139    for name in names {
140        let lookup = format!("\"{}\":", name);
141        let index = data.find(lookup.as_str()).unwrap();
142        field_names.insert(index, name);
143        indexes.push(index);
144    };
145    indexes.sort();
146    let mut adjusted_field_names = Vec::new();
147
148    for index in indexes {
149        let field = field_names.get(&index).unwrap();
150        adjusted_field_names.push(field.clone())
151    }
152
153
154    if let Value::Map(kv) = &value {
155        let mut builder = Schema::builder();
156        for key in adjusted_field_names {
157            let sss = Value::String(key);
158            let value = kv.get(&sss);
159            if value.is_none() {
160                continue;
161            };
162            let value = value.unwrap();
163            if let Value::String(k) = &sss {
164                match value {
165                    Value::String(_) => {
166                        let options = resolve_text_option(k, control);
167                        builder.add_text_field(k, options);
168                        field_type_mappings.insert(k.to_string(), SurferFieldTypes::String);
169                    }
170                    Value::Bool(_) => {
171                        let options = resolve_text_option(k, control);
172                        builder.add_text_field(k, options);
173                        field_type_mappings.insert(k.to_string(), SurferFieldTypes::String);
174                    }
175                    Value::U64(_) => {
176                        let options = resolve_number_option(k, control);
177                        builder.add_u64_field(k, options);
178                        field_type_mappings.insert(k.to_string(), SurferFieldTypes::U64);
179                    }
180                    Value::U32(_) => {
181                        let options = resolve_number_option(k, control);
182                        builder.add_u64_field(k, options);
183                        field_type_mappings.insert(k.to_string(), SurferFieldTypes::U64);
184                    }
185                    Value::U16(_) => {
186                        let options = resolve_number_option(k, control);
187                        builder.add_u64_field(k, options);
188                        field_type_mappings.insert(k.to_string(), SurferFieldTypes::U64);
189                    }
190                    Value::U8(_) => {
191                        let options = resolve_number_option(k, control);
192                        builder.add_u64_field(k, options);
193                        field_type_mappings.insert(k.to_string(), SurferFieldTypes::U64);
194                    }
195                    Value::I64(_) => {
196                        let options = resolve_number_option(k, control);
197                        builder.add_i64_field(k, options);
198                        field_type_mappings.insert(k.to_string(), SurferFieldTypes::I64);
199                    }
200                    Value::I32(_) => {
201                        let options = resolve_number_option(k, control);
202                        builder.add_i64_field(k, options);
203                        field_type_mappings.insert(k.to_string(), SurferFieldTypes::I64);
204                    }
205                    Value::I16(_) => {
206                        let options = resolve_number_option(k, control);
207                        builder.add_i64_field(k, options);
208                        field_type_mappings.insert(k.to_string(), SurferFieldTypes::I64);
209                    }
210                    Value::I8(_) => {
211                        let options = resolve_number_option(k, control);
212                        builder.add_i64_field(k, options);
213                        field_type_mappings.insert(k.to_string(), SurferFieldTypes::I64);
214                    }
215                    Value::F64(_) => {
216                        let options = resolve_number_option(k, control);
217                        builder.add_f64_field(k, options);
218                        field_type_mappings.insert(k.to_string(), SurferFieldTypes::F64);
219                    }
220                    Value::F32(_) => {
221                        let options = resolve_number_option(k, control);
222                        builder.add_f64_field(k, options);
223                        field_type_mappings.insert(k.to_string(), SurferFieldTypes::F64);
224                    }
225                    Value::Seq(_) => {
226                        builder.add_bytes_field(k);
227                        field_type_mappings.insert(k.to_string(), SurferFieldTypes::Bytes);
228                    }
229                    _ => {
230                        return Err(IndexError::new(
231                            "Unable to create schema",
232                            "Unhandled value types", )
233                        );
234                    }
235                }
236            } else {
237                return Err(IndexError::new(
238                    "Unable to create schema",
239                    "keys were not string", )
240                );
241            }
242        }
243        // TODO: Throw up for empty json
244        // return Err(IndexError::new(
245        //     "Unable to create schema",
246        //     "Empty json", )
247        // );
248
249        return Ok((builder, field_type_mappings));
250    };
251    let error = IndexError::new(
252        "Unable to create schema",
253        "Invalid JSON",
254    );
255    Err(error)
256}
257
258/// List files within a dir
259pub fn ls<T: AsRef<str>>(home: T) -> Result<Vec<PathBuf>, IndexError> {
260    let paths = std::fs::read_dir(home.as_ref())?;
261    let mut entries = Vec::<PathBuf>::new();
262    for path in paths {
263        let entry = path?;
264        let value = entry.path();
265        entries.push(value);
266    };
267    Ok(entries)
268}
269
270
271/// Convenience method to get schema
272pub(crate) fn to_schema<T: Serialize>(payload: &T, control: Option<&HashMap<String, Control>>) -> Result<(Schema, HashMap<String, SurferFieldTypes>), IndexError> {
273    let (builder, mappings) = as_schema_builder(payload, control)?;
274    Ok((builder.build(), mappings))
275}
276
277/// block thread
278pub fn block_thread(sleep_in_seconds: u64) -> u64 {
279    let duration = Duration::from_secs(sleep_in_seconds);
280    let now = Instant::now();
281    sleep(duration);
282    let result = Instant::now() - now;
283    result.as_secs()
284}
285
286pub fn random_string(size: Option<usize>) -> String {
287    let size = if size.is_none() {
288        10
289    } else {
290        size.unwrap()
291    };
292    rand::thread_rng()
293        .sample_iter(&Alphanumeric)
294        .take(size)
295        .collect::<String>()
296}
297
298
299#[cfg(test)]
300mod tests {
301    use super::*;
302
303    #[derive(Serialize)]
304    struct Empty;
305
306    #[derive(Serialize)]
307    struct Emptish {
308        value: Option<String>
309    }
310
311    #[derive(Serialize)]
312    struct DataVec {
313        identity: String,
314        buffer: Vec<u8>,
315    }
316
317
318    #[test]
319    fn validate_resolve_text_option() {
320        let key = "dummy";
321        let mut control = HashMap::new();
322        let text_options = TEXT;
323        control.insert(key.to_string(), Control::ControlTextOptions(text_options));
324        let options = resolve_text_option(key, Some(&control));
325        assert_eq!(options.is_stored(), false);
326    }
327
328    #[test]
329    fn validate_resolve_default_text_option() {
330        let key = "dummy";
331        let options = resolve_text_option(key, None);
332        assert_eq!(options.is_stored(), true);
333    }
334
335    #[test]
336    fn validate_resolve_number_option() {
337        let key = "dummy";
338        let mut control = HashMap::new();
339        let int_options = IntOptions::default();
340        control.insert(key.to_string(), Control::ControlIntOptions(int_options));
341        let options = resolve_number_option(key, Some(&control));
342        assert_eq!(options.is_stored(), false);
343    }
344
345    #[test]
346    fn validate_resolve_default_number_option() {
347        let key = "dummy";
348        let options = resolve_number_option(key, None);
349        assert_eq!(options.is_stored(), true);
350    }
351
352    #[test]
353    fn invalid_resolve_text_option() {
354        let key = "dummy";
355        let mut control = HashMap::new();
356        let int_options = IntOptions::default();
357        control.insert(key.to_string(), Control::ControlIntOptions(int_options));
358        let options = resolve_text_option(key, Some(&control));
359        assert_eq!(options.is_stored(), true);
360    }
361
362    #[test]
363    fn default_resolve_text_option() {
364        let key = "dummy";
365        let mut control = HashMap::new();
366        control.insert(key.to_string(), Control::ControlTextOptions(TextOptions::default()));
367        let options = resolve_text_option(key, Some(&control));
368        assert_eq!(options.is_stored(), false);
369    }
370
371    #[test]
372    fn default_resolve_number_option() {
373        let key = "dummy";
374        let mut control = HashMap::new();
375        control.insert(key.to_string(), Control::ControlIntOptions(IntOptions::default()));
376        let options = resolve_number_option(key, Some(&control));
377        assert_eq!(options.is_stored(), false);
378    }
379
380
381    #[test]
382    fn invalid_resolve_number_option() {
383        let key = "dummy";
384        let mut control = HashMap::new();
385        let text_options = TEXT;
386        control.insert(key.to_string(), Control::ControlTextOptions(text_options));
387        let options = resolve_number_option(key, Some(&control));
388        assert_eq!(options.is_stored(), true);
389    }
390
391    #[test]
392    fn invalid_field_names() {
393        let empty = Empty;
394        let value = as_value(&empty).unwrap();
395        let result = field_names(&value);
396        assert!(result.is_none());
397    }
398
399    #[test]
400    fn validate_error_on_empty_struct() {
401        let data = Empty;
402        let value = as_value(&data);
403        assert!(value.is_ok());
404        let value = value.unwrap();
405        let result = as_schema_builder(&value, None);
406        assert!(result.is_err());
407    }
408
409    #[test]
410    fn validate_schema_builder_for_emptish() {
411        let data = Emptish {
412            value: None
413        };
414        let value = as_value(&data);
415        assert!(value.is_ok());
416        let value = value.unwrap();
417        let result = as_schema_builder(&value, None);
418        assert!(result.is_err());
419    }
420
421    #[test]
422    fn validate_schema_builder_for_vec_does_not_work() {
423        let identity = "Hello".to_string();
424        let buffer = "World".as_bytes().to_vec();
425        let data = DataVec {
426            identity,
427            buffer,
428        };
429        let result = as_value(&data);
430        let data = result.unwrap();
431        let (schema, _) = to_schema(&data, None).unwrap();
432        let data = serde_json::to_string(&data).unwrap();
433        let document = schema.parse_document(&data);
434        assert!(document.is_err())
435    }
436}