1use crate::WparseError;
2use crate::ast::AnnFun;
3use smol_str::SmolStr;
4use std::collections::BTreeMap;
5use wp_model_core::model::{DataField, DataRecord};
6use wp_model_core::raw::RawData;
7use wp_source_types::SourceEvent;
8
9pub trait AnnotationFunc {
10 fn proc(&self, src: &SourceEvent, data: &mut DataRecord) -> Result<(), WparseError>;
11}
12
13#[derive(Clone, Debug)]
14pub struct TagAnnotation {
15 args: BTreeMap<SmolStr, SmolStr>,
16}
17
18impl AnnotationFunc for TagAnnotation {
19 fn proc(&self, _src: &SourceEvent, data: &mut DataRecord) -> Result<(), WparseError> {
20 for (key, val) in &self.args {
21 data.append(DataField::from_chars(key.clone(), val.clone()));
22 }
23 Ok(())
24 }
25}
26
27#[derive(Clone, Debug)]
28pub struct NoopAnnotation;
29
30impl AnnotationFunc for NoopAnnotation {
31 fn proc(&self, _src: &SourceEvent, _data: &mut DataRecord) -> Result<(), WparseError> {
32 Ok(())
33 }
34}
35
36#[derive(Clone, Debug)]
37pub struct RawCopy {
38 raw_key: SmolStr,
39}
40
41impl AnnotationFunc for RawCopy {
42 fn proc(&self, src: &SourceEvent, data: &mut DataRecord) -> Result<(), WparseError> {
43 match &src.payload {
44 RawData::String(raw) => {
45 data.append(DataField::from_chars(self.raw_key.clone(), raw.clone()));
46 }
47 RawData::Bytes(raw) => {
48 data.append(DataField::from_chars(
49 self.raw_key.clone(),
50 String::from_utf8_lossy(raw).into_owned(),
51 ));
52 }
53 RawData::ArcBytes(raw) => {
54 data.append(DataField::from_chars(
55 self.raw_key.clone(),
56 String::from_utf8_lossy(raw).into_owned(),
57 ));
58 }
59 }
60 Ok(())
61 }
62}
63
64#[derive(Clone, Debug)]
65pub enum AnnotationType {
66 Tag(TagAnnotation),
67 Copy(RawCopy),
68 Null(NoopAnnotation),
69}
70
71impl AnnotationFunc for AnnotationType {
72 fn proc(&self, src: &SourceEvent, data: &mut DataRecord) -> Result<(), WparseError> {
73 match self {
74 AnnotationType::Tag(func) => func.proc(src, data),
75 AnnotationType::Null(func) => func.proc(src, data),
76 AnnotationType::Copy(func) => func.proc(src, data),
77 }
78 }
79}
80
81impl AnnotationType {
82 pub fn convert(ann: &Option<AnnFun>) -> Vec<Self> {
83 let mut vec = vec![];
84 if let Some(ann) = ann {
85 if !ann.tags.is_empty() {
86 vec.push(AnnotationType::Tag(TagAnnotation {
87 args: ann.tags.clone(),
88 }));
89 }
90
91 if let Some((k, v)) = &ann.copy_raw {
92 if k == "name" {
93 vec.push(AnnotationType::Copy(RawCopy { raw_key: v.clone() }));
94 } else {
95 vec.push(AnnotationType::Null(NoopAnnotation {}))
96 }
97 }
98 } else {
99 vec.push(AnnotationType::Null(NoopAnnotation {}))
100 }
101 vec
102 }
103}
104
105#[cfg(test)]
106mod tests {
107 use super::*;
108 use crate::pkg::DEFAULT_KEY;
109 use bytes::Bytes;
110 use orion_error::dev::testing::TestAssert;
111 use std::collections::BTreeMap;
112 use std::sync::Arc;
113 use wp_model_core::model::DataRecord;
114 use wp_model_core::raw::RawData;
115 use wp_source_types::{SourceEvent, Tags};
116
117 #[test]
118 fn test_tag_fun() {
119 let ann = AnnFun {
120 tags: BTreeMap::from([("tag_1".into(), "x".into())]),
121 copy_raw: None,
122 };
123 let tag = AnnotationType::convert(&Some(ann));
124 let mut data = DataRecord::test_value();
125 let src = SourceEvent::new(
126 1,
127 DEFAULT_KEY.to_string(),
128 RawData::String("test".to_string()),
129 Tags::new().into(),
130 );
131 tag.first().unwrap().proc(&src, &mut data).assert();
132 let expected = DataField::from_chars("tag_1", "x");
133 assert_eq!(data.field("tag_1").map(|s| s.as_field()), Some(&expected));
134 }
135
136 #[test]
137 fn test_copy_fun() {
138 let ann = AnnFun {
139 tags: Default::default(),
140 copy_raw: Some(("name".into(), "raw".into())),
141 };
142 let tag = AnnotationType::convert(&Some(ann));
143 let mut data = DataRecord::test_value();
144 let src = SourceEvent::new(
145 1,
146 DEFAULT_KEY.to_string(),
147 RawData::String("test".to_string()),
148 Tags::new().into(),
149 );
150 tag.first().unwrap().proc(&src, &mut data).unwrap();
151 let expected = DataField::from_chars("raw", "test");
152 assert_eq!(data.field("raw").map(|s| s.as_field()), Some(&expected));
153 }
154
155 #[test]
156 fn test_copy_fun_handles_invalid_utf8_bytes() {
157 let tag = copy_raw_tag("raw");
158 let raw = Bytes::from_static(b"hello \xff\xfe\xc0\xaf");
159 let expected_raw = String::from_utf8_lossy(&raw).into_owned();
160 let mut data = DataRecord::test_value();
161 let src = SourceEvent::new(
162 1,
163 DEFAULT_KEY.to_string(),
164 RawData::Bytes(raw),
165 Tags::new().into(),
166 );
167
168 tag.first().unwrap().proc(&src, &mut data).unwrap();
169
170 let expected = DataField::from_chars("raw", expected_raw);
171 assert_eq!(data.field("raw").map(|s| s.as_field()), Some(&expected));
172 }
173
174 #[test]
175 fn test_copy_fun_handles_invalid_utf8_arc_bytes() {
176 let tag = copy_raw_tag("raw");
177 let raw = Arc::new(b"hello \xff\xfe\xc0\xaf".to_vec());
178 let expected_raw = String::from_utf8_lossy(raw.as_slice()).into_owned();
179 let mut data = DataRecord::test_value();
180 let src = SourceEvent::new(
181 1,
182 DEFAULT_KEY.to_string(),
183 RawData::ArcBytes(raw),
184 Tags::new().into(),
185 );
186
187 tag.first().unwrap().proc(&src, &mut data).unwrap();
188
189 let expected = DataField::from_chars("raw", expected_raw);
190 assert_eq!(data.field("raw").map(|s| s.as_field()), Some(&expected));
191 }
192
193 fn copy_raw_tag(raw_key: &str) -> Vec<AnnotationType> {
194 AnnotationType::convert(&Some(AnnFun {
195 tags: Default::default(),
196 copy_raw: Some(("name".into(), raw_key.into())),
197 }))
198 }
199}