wpl/ast/
ann_func.rs

1use crate::WparseError;
2use crate::WplEvaluator;
3use crate::ast::AnnFun;
4use smol_str::SmolStr;
5use std::collections::BTreeMap;
6use wp_model_core::model::{DataField, DataRecord};
7use wp_model_core::raw::RawData;
8use wp_source_types::SourceEvent;
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) => {
49                data.append(DataField::from_chars(
50                    self.raw_key.clone(),
51                    String::from_utf8_lossy(raw).into_owned(),
52                ));
53            }
54            RawData::ArcBytes(raw) => {
55                data.append(DataField::from_chars(
56                    self.raw_key.clone(),
57                    String::from_utf8_lossy(raw).into_owned(),
58                ));
59            }
60        }
61        Ok(())
62    }
63}
64
65/// 将原始 payload 复制给指定 rule 的 parser 解析，把产出的字段并入当前 record。
66///
67/// `target` 在构建期由 motor 层注入（按 `rule_name` 解析同包 rule 的 parser）；
68/// 未注入时（如 editor/station 解析路径）`proc` 直接 no-op。
69/// 目标 rule 只解析，不执行其自身注解。
70#[derive(Clone)]
71pub struct CopyEventParseAnnotation {
72    pub rule_name: SmolStr,
73    pub target: Option<WplEvaluator>,
74}
75
76impl std::fmt::Debug for CopyEventParseAnnotation {
77    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
78        f.debug_struct("CopyEventParseAnnotation")
79            .field("rule_name", &self.rule_name)
80            .field("target_set", &self.target.is_some())
81            .finish()
82    }
83}
84
85impl AnnotationFunc for CopyEventParseAnnotation {
86    fn proc(&self, src: &SourceEvent, data: &mut DataRecord) -> Result<(), WparseError> {
87        let Some(target) = &self.target else {
88            return Ok(());
89        };
90        // 复制 src.payload 喂给目标 rule 的 parser（proc_ref 避免每条事件 clone payload）
91        let (target_rec, _left) = target.proc_ref(src.event_id, &src.payload, 0)?;
92        // 把目标 rule 解析产出的字段并入当前 record
93        data.merge(target_rec);
94        Ok(())
95    }
96}
97
98#[derive(Clone, Debug)]
99pub enum AnnotationType {
100    Tag(TagAnnotation),
101    Copy(RawCopy),
102    Null(NoopAnnotation),
103    CopyEventParse(CopyEventParseAnnotation),
104}
105
106impl AnnotationFunc for AnnotationType {
107    fn proc(&self, src: &SourceEvent, data: &mut DataRecord) -> Result<(), WparseError> {
108        match self {
109            AnnotationType::Tag(func) => func.proc(src, data),
110            AnnotationType::Null(func) => func.proc(src, data),
111            AnnotationType::Copy(func) => func.proc(src, data),
112            AnnotationType::CopyEventParse(func) => func.proc(src, data),
113        }
114    }
115}
116
117impl AnnotationType {
118    pub fn convert(ann: &Option<AnnFun>) -> Vec<Self> {
119        let mut vec = vec![];
120        if let Some(ann) = ann {
121            if !ann.tags.is_empty() {
122                vec.push(AnnotationType::Tag(TagAnnotation {
123                    args: ann.tags.clone(),
124                }));
125            }
126
127            if let Some((k, v)) = &ann.copy_raw {
128                if k == "name" {
129                    vec.push(AnnotationType::Copy(RawCopy { raw_key: v.clone() }));
130                } else {
131                    vec.push(AnnotationType::Null(NoopAnnotation {}))
132                }
133            }
134
135            if let Some(rule) = &ann.copy_event_parse {
136                // target 留空，构建期由 motor 层按 rule_name 注入目标 rule 的 parser
137                vec.push(AnnotationType::CopyEventParse(CopyEventParseAnnotation {
138                    rule_name: rule.clone(),
139                    target: None,
140                }));
141            }
142        } else {
143            vec.push(AnnotationType::Null(NoopAnnotation {}))
144        }
145        vec
146    }
147}
148
149#[cfg(test)]
150mod tests {
151    use super::*;
152    use crate::pkg::DEFAULT_KEY;
153    use bytes::Bytes;
154    use orion_error::dev::testing::TestAssert;
155    use std::collections::BTreeMap;
156    use std::sync::Arc;
157    use wp_model_core::model::DataRecord;
158    use wp_model_core::raw::RawData;
159    use wp_source_types::{SourceEvent, Tags};
160
161    #[test]
162    fn test_tag_fun() {
163        let ann = AnnFun {
164            tags: BTreeMap::from([("tag_1".into(), "x".into())]),
165            copy_raw: None,
166            copy_event_parse: None,
167            no_match: false,
168        };
169        let tag = AnnotationType::convert(&Some(ann));
170        let mut data = DataRecord::test_value();
171        let src = SourceEvent::new(
172            1,
173            DEFAULT_KEY.to_string(),
174            RawData::String("test".to_string()),
175            Tags::new().into(),
176        );
177        tag.first().unwrap().proc(&src, &mut data).assert();
178        let expected = DataField::from_chars("tag_1", "x");
179        assert_eq!(data.field("tag_1").map(|s| s.as_field()), Some(&expected));
180    }
181
182    #[test]
183    fn test_copy_fun() {
184        let ann = AnnFun {
185            tags: Default::default(),
186            copy_raw: Some(("name".into(), "raw".into())),
187            copy_event_parse: None,
188            no_match: false,
189        };
190        let tag = AnnotationType::convert(&Some(ann));
191        let mut data = DataRecord::test_value();
192        let src = SourceEvent::new(
193            1,
194            DEFAULT_KEY.to_string(),
195            RawData::String("test".to_string()),
196            Tags::new().into(),
197        );
198        tag.first().unwrap().proc(&src, &mut data).unwrap();
199        let expected = DataField::from_chars("raw", "test");
200        assert_eq!(data.field("raw").map(|s| s.as_field()), Some(&expected));
201    }
202
203    #[test]
204    fn test_copy_fun_handles_invalid_utf8_bytes() {
205        let tag = copy_raw_tag("raw");
206        let raw = Bytes::from_static(b"hello \xff\xfe\xc0\xaf");
207        let expected_raw = String::from_utf8_lossy(&raw).into_owned();
208        let mut data = DataRecord::test_value();
209        let src = SourceEvent::new(
210            1,
211            DEFAULT_KEY.to_string(),
212            RawData::Bytes(raw),
213            Tags::new().into(),
214        );
215
216        tag.first().unwrap().proc(&src, &mut data).unwrap();
217
218        let expected = DataField::from_chars("raw", expected_raw);
219        assert_eq!(data.field("raw").map(|s| s.as_field()), Some(&expected));
220    }
221
222    #[test]
223    fn test_copy_fun_handles_invalid_utf8_arc_bytes() {
224        let tag = copy_raw_tag("raw");
225        let raw = Arc::new(b"hello \xff\xfe\xc0\xaf".to_vec());
226        let expected_raw = String::from_utf8_lossy(raw.as_slice()).into_owned();
227        let mut data = DataRecord::test_value();
228        let src = SourceEvent::new(
229            1,
230            DEFAULT_KEY.to_string(),
231            RawData::ArcBytes(raw),
232            Tags::new().into(),
233        );
234
235        tag.first().unwrap().proc(&src, &mut data).unwrap();
236
237        let expected = DataField::from_chars("raw", expected_raw);
238        assert_eq!(data.field("raw").map(|s| s.as_field()), Some(&expected));
239    }
240
241    fn copy_raw_tag(raw_key: &str) -> Vec<AnnotationType> {
242        AnnotationType::convert(&Some(AnnFun {
243            tags: Default::default(),
244            copy_raw: Some(("name".into(), raw_key.into())),
245            copy_event_parse: None,
246            no_match: false,
247        }))
248    }
249
250    /// 构造一个注入了 target parser 的 copy_event_parse 注解，
251    /// target rule 把 payload 整体捕获为 raw 字段。
252    fn copy_event_parse_with_target(rule_code: &str, rule_name: &str) -> Vec<AnnotationType> {
253        let target = WplEvaluator::from_code(rule_code).expect("build target evaluator");
254        let mut funcs = AnnotationType::convert(&Some(AnnFun {
255            tags: Default::default(),
256            copy_raw: None,
257            copy_event_parse: Some(rule_name.into()),
258            no_match: false,
259        }));
260        for ann in &mut funcs {
261            if let AnnotationType::CopyEventParse(c) = ann {
262                c.target = Some(target.clone());
263            }
264        }
265        funcs
266    }
267
268    #[test]
269    fn test_copy_event_parse_merges_target_fields() {
270        // 目标 rule：解析 JSON {"raw":"..."} 产出名为 raw 的 chars 字段
271        let funcs =
272            copy_event_parse_with_target(r#"rule raw_event { (json(chars@raw)) }"#, "raw_event");
273        let mut data = DataRecord::test_value();
274        let src = SourceEvent::new(
275            1,
276            DEFAULT_KEY.to_string(),
277            RawData::String(r#"{ "raw": "payload-content" }"#.to_string()),
278            Tags::new().into(),
279        );
280        funcs.first().unwrap().proc(&src, &mut data).unwrap();
281        let expected = DataField::from_chars("raw", "payload-content");
282        assert_eq!(data.field("raw").map(|s| s.as_field()), Some(&expected));
283    }
284
285    #[test]
286    fn test_copy_event_parse_noop_without_target() {
287        // 未注入 target（editor/station 路径）时应 no-op，不报错也不写字段
288        let funcs = AnnotationType::convert(&Some(AnnFun {
289            tags: Default::default(),
290            copy_raw: None,
291            copy_event_parse: Some("raw_event".into()),
292            no_match: false,
293        }));
294        let mut data = DataRecord::test_value();
295        let src = SourceEvent::new(
296            1,
297            DEFAULT_KEY.to_string(),
298            RawData::String("payload-content".to_string()),
299            Tags::new().into(),
300        );
301        funcs.first().unwrap().proc(&src, &mut data).unwrap();
302        // 未注入 target 时 raw 字段不应存在
303        assert!(data.field("raw").is_none());
304    }
305}
wpl/ast/ann_func.rs

wpl/ast/
ann_func.rs
