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}