jupiter/idb/
idb_yaml_set_loader.rs

1//! Imports YAML files as sets into InfoGraphDB.
2//!
3//! The file is expected to contain a map (hash) which contains lists defining the sets.
4//!
5//! ## Example data:
6//! ```yaml
7//! set1: [A, B, C]
8//! set2: [X, Y]
9//! ```
10//! ## Example loader:
11//! ```yaml
12//! file: 'path/to/file.json'
13//! loader: 'idb-yaml-sets'
14//! namespace: 'target namespace to import into'
15//! ```
16use crate::idb::set::Set;
17use crate::idb::{Database, DatabaseCommand};
18use crate::platform::Platform;
19use crate::repository::loader::{Loader, LoaderInfo};
20use anyhow::Context;
21use linked_hash_map::OccupiedEntry;
22use std::fmt::{Display, Formatter};
23use std::sync::Arc;
24use yaml_rust::Yaml;
25
26/// Represents the global loader instance.
27pub struct IdbYamlSetLoader {
28    platform: Arc<Platform>,
29}
30
31impl IdbYamlSetLoader {
32    /// Creates a new loader to be passed into [Repository::register_loader].
33    pub fn new(platform: Arc<Platform>) -> Self {
34        IdbYamlSetLoader { platform }
35    }
36}
37
38impl Display for IdbYamlSetLoader {
39    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
40        write!(f, "IDB-YAML-SET")
41    }
42}
43
44#[async_trait::async_trait]
45impl Loader for IdbYamlSetLoader {
46    async fn file_changed(&self, loader_info: &LoaderInfo) -> anyhow::Result<()> {
47        let data = tokio::fs::read_to_string(loader_info.get_data())
48            .await
49            .context("Unable to read data file")?;
50        let mut rows = yaml_rust::YamlLoader::load_from_str(data.as_str())
51            .context("Cannot parse the given YAML data.")?;
52
53        // If only one yaml object is present, and it's an array -> unwrap it. This was most probably
54        // a JSON file like [{obj1}, {obj2}...]...
55        if rows.len() == 1 && rows.first().unwrap().is_array() {
56            rows = rows.remove(0).into_vec().unwrap();
57        }
58
59        let source = loader_info.file_name().to_string();
60        let sets = self.load_sets(rows);
61        for (name, set) in sets {
62            self.register_set(source.clone(), name, set).await?;
63        }
64
65        Ok(())
66    }
67
68    fn platform(&self) -> &Arc<Platform> {
69        &self.platform
70    }
71
72    async fn file_deleted(&self, loader_info: &LoaderInfo) -> anyhow::Result<()> {
73        let source = loader_info.file_name().to_string();
74
75        self.platform()
76            .require::<Database>()
77            .perform(DatabaseCommand::DropSets(source))
78            .await
79            .context("Failed to drop set.")?;
80
81        Ok(())
82    }
83}
84
85impl IdbYamlSetLoader {
86    fn load_sets(&self, rows: Vec<Yaml>) -> Vec<(String, Set)> {
87        let mut result = Vec::new();
88        for row in rows {
89            if let Yaml::Hash(mut map) = row {
90                for entry in map.entries() {
91                    if let Some((name, set)) = self.transform(entry) {
92                        result.push((name, set));
93                    }
94                }
95            }
96        }
97
98        result
99    }
100
101    fn transform(&self, entry: OccupiedEntry<Yaml, Yaml>) -> Option<(String, Set)> {
102        if let Yaml::String(key) = entry.key() {
103            if let Yaml::Array(list) = entry.get() {
104                let mut set = Set::default();
105                for item in list {
106                    if let Some(str) = item.as_str() {
107                        set.add(str.to_string());
108                    }
109                }
110
111                Some((key.to_owned(), set))
112            } else {
113                None
114            }
115        } else {
116            None
117        }
118    }
119
120    async fn register_set(&self, source: String, key: String, set: Set) -> anyhow::Result<()> {
121        self.platform()
122            .require::<Database>()
123            .perform(DatabaseCommand::CreateSet(source, key, set))
124            .await
125            .context("Failed to create set.")
126    }
127}
jupiter/idb/idb_yaml_set_loader.rs

jupiter/idb/
idb_yaml_set_loader.rs
