feat(classify): optional priority field for category rules

aw-query/src/datatype.rs

@@ -4,7 +4,7 @@ use std::fmt;
 use super::functions;
 use super::QueryError;
 use aw_models::Event;
-use aw_transform::classify::{RegexRule, Rule};
+use aw_transform::classify::{CategoryRule, RegexRule, Rule};
 
 use serde::{Serialize, Serializer};
 use serde_json::value::Value;
@@ -224,18 +224,48 @@ impl TryFrom<DataType> for Vec<(String, Rule)> {
     }
 }
 
-impl TryFrom<&DataType> for Vec<(Vec<String>, Rule)> {
+impl TryFrom<&DataType> for Vec<CategoryRule> {
     type Error = QueryError;
     fn try_from(value: &DataType) -> Result<Self, Self::Error> {
         value.clone().try_into()
     }
 }
 
-impl TryFrom<DataType> for Vec<(Vec<String>, Rule)> {
+fn parse_optional_priority(rule: &DataType) -> Result<Option<i64>, QueryError> {
+    let obj = match rule {
+        DataType::Dict(dict) => dict,
+        _ => return Ok(None),
+    };
+    let val = match obj.get("priority").or_else(|| obj.get("weight")) {
+        Some(v) => v,
+        None => return Ok(None),
+    };
+    match val {
+        DataType::Number(n) => {
+            if !n.is_finite() {
+                return Err(QueryError::InvalidFunctionParameters(
+                    "priority/weight must be a finite integer".to_string(),
+                ));
+            }
+            let as_int = *n as i64;
+            if (as_int as f64) != *n {
+                return Err(QueryError::InvalidFunctionParameters(
+                    "priority/weight must be an integer".to_string(),
+                ));
+            }
+            Ok(Some(as_int))
+        }
+        _ => Err(QueryError::InvalidFunctionParameters(
+            "priority/weight must be an integer".to_string(),
+        )),
+    }
+}
+
+impl TryFrom<DataType> for Vec<CategoryRule> {
     type Error = QueryError;
     fn try_from(value: DataType) -> Result<Self, Self::Error> {
         let tagged_lists: Vec<DataType> = value.try_into()?;
-        let mut lists: Vec<(Vec<String>, Rule)> = Vec::with_capacity(tagged_lists.len());
+        let mut lists: Vec<CategoryRule> = Vec::with_capacity(tagged_lists.len());
         for list in tagged_lists {
             match list {
                 DataType::List(ref l) => {
@@ -244,12 +274,18 @@ impl TryFrom<DataType> for Vec<(Vec<String>, Rule)> {
                         None => return Err(QueryError::InvalidFunctionParameters(
                             format!("Expected function parameter of type list of (category, rule) tuples, list contains {l:?}")))
                     };
-                    let rule: Rule = match l.get(1) {
-                        Some(rule) => rule.try_into()?,
+                    let rule_data = match l.get(1) {
+                        Some(rule) => rule,
                         None => return Err(QueryError::InvalidFunctionParameters(
                             format!("Expected function parameter of type list of (category, rule) tuples, list contains {l:?}")))
                     };
-                    lists.push((category, rule));
+                    let priority = parse_optional_priority(rule_data)?;
+                    let rule: Rule = rule_data.try_into()?;
+                    lists.push(CategoryRule {
+                        category,
+                        rule,
+                        priority,
+                    });
                 }
                 ref invalid_type => {
                     return Err(QueryError::InvalidFunctionParameters(format!(

aw-query/src/functions.rs

@@ -118,7 +118,7 @@ pub fn fill_env(env: &mut VarEnv) {
 mod qfunctions {
     use aw_datastore::Datastore;
     use aw_models::Event;
-    use aw_transform::classify::Rule;
+    use aw_transform::classify::{CategoryRule, Rule};
 
     use super::validate;
     use crate::DataType;
@@ -290,7 +290,7 @@ mod qfunctions {
         validate::args_length(&args, 2)?;
         let mut args = args.into_iter();
         let events: Vec<Event> = args.next().unwrap().try_into()?;
-        let rules: Vec<(Vec<String>, Rule)> = args.next().unwrap().try_into()?;
+        let rules: Vec<CategoryRule> = args.next().unwrap().try_into()?;
         // Run categorize
         let mut flooded_events = aw_transform::classify::categorize(events, &rules);
         // Put events back into DataType::Event container

aw-query/tests/query.rs

@@ -390,6 +390,57 @@ mod query_tests {
         assert_eq!(cats, &serde_json::json!(vec!["Test", "Subtest"]));
     }
 
+    #[test]
+    fn test_categorize_priority_overrides_depth() {
+        let ds = setup_datastore_populated();
+        let interval = TimeInterval::new_from_string(TIME_INTERVAL).unwrap();
+
+        // Without priority, the deeper rule wins (existing contract).
+        let code = format!(
+            r#"
+            events = query_bucket("{}");
+            events = categorize(events, [[["A"], {{ "type": "regex", "regex": "^value$" }}], [["B", "B1"], {{ "type": "regex", "regex": "^value$" }}]]);
+            return  events;"#,
+            "testid"
+        );
+        let result: DataType = aw_query::query(&code, &interval, &ds).unwrap();
+        let events: Vec<Event> = Vec::try_from(&result).unwrap();
+        assert_eq!(
+            events.first().unwrap().data.get("$category").unwrap(),
+            &serde_json::json!(vec!["B", "B1"])
+        );
+
+        // With an explicit priority on A, organizational depth no longer wins.
+        let code = format!(
+            r#"
+            events = query_bucket("{}");
+            events = categorize(events, [[["A"], {{ "type": "regex", "regex": "^value$", "priority": 25 }}], [["B", "B1"], {{ "type": "regex", "regex": "^value$" }}]]);
+            return  events;"#,
+            "testid"
+        );
+        let result: DataType = aw_query::query(&code, &interval, &ds).unwrap();
+        let events: Vec<Event> = Vec::try_from(&result).unwrap();
+        assert_eq!(
+            events.first().unwrap().data.get("$category").unwrap(),
+            &serde_json::json!(vec!["A"])
+        );
+
+        // `weight` is accepted as an alias for `priority`.
+        let code = format!(
+            r#"
+            events = query_bucket("{}");
+            events = categorize(events, [[["A"], {{ "type": "regex", "regex": "^value$", "weight": 25 }}], [["B", "B1"], {{ "type": "regex", "regex": "^value$" }}]]);
+            return  events;"#,
+            "testid"
+        );
+        let result: DataType = aw_query::query(&code, &interval, &ds).unwrap();
+        let events: Vec<Event> = Vec::try_from(&result).unwrap();
+        assert_eq!(
+            events.first().unwrap().data.get("$category").unwrap(),
+            &serde_json::json!(vec!["A"])
+        );
+    }
+
     #[test]
     fn test_tag() {
         let ds = setup_datastore_populated();
@@ -547,6 +598,22 @@ mod query_tests {
             return  events;"#;
         let res = aw_query::query(code, &interval, &ds);
         assert_err_type!(res, QueryError::RegexCompileError(_));
+
+        // Test categorize rule with non-integer priority
+        let code = r#"
+            events = [];
+            events = categorize(events, [[["test"], { "type": "regex", "regex": "test", "priority": 1.5 }]]);
+            return  events;"#;
+        let res = aw_query::query(code, &interval, &ds);
+        assert_err_type!(res, QueryError::InvalidFunctionParameters(_));
+
+        // Test categorize rule with non-numeric priority
+        let code = r#"
+            events = [];
+            events = categorize(events, [[["test"], { "type": "regex", "regex": "test", "priority": "high" }]]);
+            return  events;"#;
+        let res = aw_query::query(code, &interval, &ds);
+        assert_err_type!(res, QueryError::InvalidFunctionParameters(_));
     }
 
     #[test]