Skip to main content

wpl/ast/
ann_func.rs

1use crate::ast::AnnFun;
2use crate::{WparseError, WparseReason};
3use orion_error::UvsFrom;
4use orion_error::conversion::ToStructError;
5use smol_str::SmolStr;
6use std::collections::BTreeMap;
7use wp_connector_api::SourceEvent;
8use wp_model_core::model::{DataField, DataRecord};
9use wp_model_core::raw::RawData;
10
11pub trait AnnotationFunc {
12    fn proc(&self, src: &SourceEvent, data: &mut DataRecord) -> Result<(), WparseError>;
13}
14
15#[derive(Clone, Debug)]
16pub struct TagAnnotation {
17    args: BTreeMap<SmolStr, SmolStr>,
18}
19
20impl AnnotationFunc for TagAnnotation {
21    fn proc(&self, _src: &SourceEvent, data: &mut DataRecord) -> Result<(), WparseError> {
22        for (key, val) in &self.args {
23            data.append(DataField::from_chars(key.clone(), val.clone()));
24        }
25        Ok(())
26    }
27}
28
29#[derive(Clone, Debug)]
30pub struct NoopAnnotation;
31
32impl AnnotationFunc for NoopAnnotation {
33    fn proc(&self, _src: &SourceEvent, _data: &mut DataRecord) -> Result<(), WparseError> {
34        Ok(())
35    }
36}
37
38#[derive(Clone, Debug)]
39pub struct RawCopy {
40    raw_key: SmolStr,
41}
42
43impl AnnotationFunc for RawCopy {
44    fn proc(&self, src: &SourceEvent, data: &mut DataRecord) -> Result<(), WparseError> {
45        match &src.payload {
46            RawData::String(raw) => {
47                data.append(DataField::from_chars(self.raw_key.clone(), raw.clone()));
48            }
49            RawData::Bytes(raw) => match std::str::from_utf8(raw) {
50                Ok(str) => {
51                    data.append(DataField::from_chars(self.raw_key.clone(), str.to_string()));
52                }
53                Err(e) => {
54                    return Err(WparseReason::from_data()
55                        .to_err()
56                        .with_detail(format!("[u8] to string error :{}", e)));
57                }
58            },
59            RawData::ArcBytes(raw) => match std::str::from_utf8(raw) {
60                Ok(str) => {
61                    data.append(DataField::from_chars(self.raw_key.clone(), str.to_string()));
62                }
63                Err(e) => {
64                    return Err(WparseReason::from_data()
65                        .to_err()
66                        .with_detail(format!("ArcBytes to string error :{}", e)));
67                }
68            },
69        }
70        Ok(())
71    }
72}
73
74#[derive(Clone, Debug)]
75pub enum AnnotationType {
76    Tag(TagAnnotation),
77    Copy(RawCopy),
78    Null(NoopAnnotation),
79}
80
81impl AnnotationFunc for AnnotationType {
82    fn proc(&self, src: &SourceEvent, data: &mut DataRecord) -> Result<(), WparseError> {
83        match self {
84            AnnotationType::Tag(func) => func.proc(src, data),
85            AnnotationType::Null(func) => func.proc(src, data),
86            AnnotationType::Copy(func) => func.proc(src, data),
87        }
88    }
89}
90
91impl AnnotationType {
92    pub fn convert(ann: &Option<AnnFun>) -> Vec<Self> {
93        let mut vec = vec![];
94        if let Some(ann) = ann {
95            if !ann.tags.is_empty() {
96                vec.push(AnnotationType::Tag(TagAnnotation {
97                    args: ann.tags.clone(),
98                }));
99            }
100
101            if let Some((k, v)) = &ann.copy_raw {
102                if k == "name" {
103                    vec.push(AnnotationType::Copy(RawCopy { raw_key: v.clone() }));
104                } else {
105                    vec.push(AnnotationType::Null(NoopAnnotation {}))
106                }
107            }
108        } else {
109            vec.push(AnnotationType::Null(NoopAnnotation {}))
110        }
111        vec
112    }
113}
114
115#[cfg(test)]
116mod tests {
117    use super::*;
118    use crate::pkg::DEFAULT_KEY;
119    use orion_error::testcase::TestAssert;
120    use std::collections::BTreeMap;
121    use wp_connector_api::{SourceEvent, Tags};
122    use wp_model_core::model::DataRecord;
123    use wp_model_core::raw::RawData;
124
125    #[test]
126    fn test_tag_fun() {
127        let ann = AnnFun {
128            tags: BTreeMap::from([("tag_1".into(), "x".into())]),
129            copy_raw: None,
130        };
131        let tag = AnnotationType::convert(&Some(ann));
132        let mut data = DataRecord::test_value();
133        let src = SourceEvent::new(
134            1,
135            DEFAULT_KEY.to_string(),
136            RawData::String("test".to_string()),
137            Tags::new().into(),
138        );
139        tag.first().unwrap().proc(&src, &mut data).assert();
140        let expected = DataField::from_chars("tag_1", "x");
141        assert_eq!(data.field("tag_1").map(|s| s.as_field()), Some(&expected));
142    }
143
144    #[test]
145    fn test_copy_fun() {
146        let ann = AnnFun {
147            tags: Default::default(),
148            copy_raw: Some(("name".into(), "raw".into())),
149        };
150        let tag = AnnotationType::convert(&Some(ann));
151        let mut data = DataRecord::test_value();
152        let src = SourceEvent::new(
153            1,
154            DEFAULT_KEY.to_string(),
155            RawData::String("test".to_string()),
156            Tags::new().into(),
157        );
158        tag.first().unwrap().proc(&src, &mut data).unwrap();
159        let expected = DataField::from_chars("raw", "test");
160        assert_eq!(data.field("raw").map(|s| s.as_field()), Some(&expected));
161    }
162}