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