diff --git a/aw-query/src/datatype.rs b/aw-query/src/datatype.rs index c2983b0c..3b043e8a 100644 --- a/aw-query/src/datatype.rs +++ b/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 for Vec<(String, Rule)> { } } -impl TryFrom<&DataType> for Vec<(Vec, Rule)> { +impl TryFrom<&DataType> for Vec { type Error = QueryError; fn try_from(value: &DataType) -> Result { value.clone().try_into() } } -impl TryFrom for Vec<(Vec, Rule)> { +fn parse_optional_priority(rule: &DataType) -> Result, 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 for Vec { type Error = QueryError; fn try_from(value: DataType) -> Result { let tagged_lists: Vec = value.try_into()?; - let mut lists: Vec<(Vec, Rule)> = Vec::with_capacity(tagged_lists.len()); + let mut lists: Vec = Vec::with_capacity(tagged_lists.len()); for list in tagged_lists { match list { DataType::List(ref l) => { @@ -244,12 +274,18 @@ impl TryFrom for Vec<(Vec, 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!( diff --git a/aw-query/src/functions.rs b/aw-query/src/functions.rs index c8f0b89f..c19d70ca 100644 --- a/aw-query/src/functions.rs +++ b/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 = args.next().unwrap().try_into()?; - let rules: Vec<(Vec, Rule)> = args.next().unwrap().try_into()?; + let rules: Vec = args.next().unwrap().try_into()?; // Run categorize let mut flooded_events = aw_transform::classify::categorize(events, &rules); // Put events back into DataType::Event container diff --git a/aw-query/tests/query.rs b/aw-query/tests/query.rs index 90cc7480..149bc2b5 100644 --- a/aw-query/tests/query.rs +++ b/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 = 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 = 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 = 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] diff --git a/aw-transform/src/classify.rs b/aw-transform/src/classify.rs index ec419cd6..8b24430d 100644 --- a/aw-transform/src/classify.rs +++ b/aw-transform/src/classify.rs @@ -113,17 +113,55 @@ impl From for Rule { } } +/// A category matching rule passed to [`categorize`]. +/// +/// `priority` is an optional integer ranking score. When set, it is used +/// instead of the depth-derived default to pick among matching rules (higher +/// wins). When `None`, ranking falls back to `depth * 10` so existing configs +/// keep their current ordering, while explicit values can slot between levels +/// (depth 1 → 10, depth 2 → 20). +pub struct CategoryRule { + pub category: Vec, + pub rule: Rule, + pub priority: Option, +} + +impl CategoryRule { + pub fn new(category: Vec, rule: Rule) -> Self { + Self { + category, + rule, + priority: None, + } + } + + pub fn with_priority(mut self, priority: i64) -> Self { + self.priority = Some(priority); + self + } +} + +impl From<(Vec, Rule)> for CategoryRule { + fn from((category, rule): (Vec, Rule)) -> Self { + Self::new(category, rule) + } +} + /// Categorizes a list of events /// /// An event can only have one category, although the category may have a hierarchy, /// for instance: "Work -> ActivityWatch -> aw-server-rust" -/// If multiple categories match, the deepest one will be chosen. +/// If multiple categories match, the highest-ranking one is chosen. +/// Ranking is the optional integer `priority` on the rule when present, +/// otherwise `depth * 10` ("the deepest one will be chosen", with room to +/// slot values between levels). Equal ranks keep the later match, matching +/// the previous depth-only `>=` comparison. /// /// Performance: builds an in-memory cache keyed on the event's data JSON so that /// events with identical data (same app/title — very common in practice) are only /// matched against the rule set once. On a month's data with 50k+ events but only /// a few hundred distinct app/title pairs this reduces regex work by >99%. -pub fn categorize(mut events: Vec, rules: &[(Vec, Rule)]) -> Vec { +pub fn categorize(mut events: Vec, rules: &[CategoryRule]) -> Vec { // Cache: serialized event data → assigned category let mut category_cache: HashMap> = HashMap::new(); let mut classified_events = Vec::with_capacity(events.len()); @@ -134,15 +172,7 @@ pub fn categorize(mut events: Vec, rules: &[(Vec, Rule)]) -> Vec< let cache_key = serde_json::to_string(&event.data).unwrap_or_default(); let category = category_cache .entry(cache_key) - .or_insert_with(|| { - let mut cat = vec!["Uncategorized".into()]; - for (c, rule) in rules { - if rule.matches(&event) { - cat = _pick_highest_ranking_category(cat, c); - } - } - cat - }) + .or_insert_with(|| _pick_category(&event, rules)) .clone(); event .data @@ -152,6 +182,28 @@ pub fn categorize(mut events: Vec, rules: &[(Vec, Rule)]) -> Vec< classified_events } +fn _pick_category(event: &Event, rules: &[CategoryRule]) -> Vec { + let mut category: Vec = vec!["Uncategorized".into()]; + // Uncategorized loses to any non-empty match, including a match with a very + // low explicit priority. i64::MIN is only used as this sentinel. + // Empty category paths are skipped so they cannot replace the fallback + // (old depth comparison: len 0 does not beat Uncategorized's len 1). + let mut rank = i64::MIN; + for class in rules { + if class.category.is_empty() { + continue; + } + if class.rule.matches(event) { + let item_rank = _effective_rank(&class.category, class.priority); + if item_rank >= rank { + category = class.category.clone(); + rank = item_rank; + } + } + } + category +} + /// Tags a list of events /// /// An event can have many tags (as opposed to only one category) which will be put into the `$tags` key of @@ -177,13 +229,12 @@ fn tag_one(mut event: Event, rules: &[(String, Rule)]) -> Event { event } -fn _pick_highest_ranking_category(acc: Vec, item: &[String]) -> Vec { - if item.len() >= acc.len() { - // If tag is category with greater or equal depth than current, then choose the new one instead. - item.to_vec() - } else { - acc - } +fn _effective_rank(category: &[String], priority: Option) -> i64 { + // Integer-only. Default is depth * 10 so explicit priorities can slot + // between nesting levels (depth 1 → 10, depth 2 → 20). Relative order of + // unprioritized rules is unchanged. + // https://github.com/ActivityWatch/aw-server-rust/pull/663#issuecomment-5481349757 + priority.unwrap_or((category.len() as i64) * 10) } #[test] @@ -260,16 +311,16 @@ fn test_categorize() { .insert("test".into(), serde_json::json!("just a test")); let mut events = vec![e]; - let rules: Vec<(Vec, Rule)> = vec![ - ( + let rules: Vec = vec![ + CategoryRule::new( vec!["Test".into()], Rule::from(Regex::new(r"test").unwrap()), ), - ( + CategoryRule::new( vec!["Test".into(), "Subtest".into()], Rule::from(Regex::new(r"test").unwrap()), ), - ( + CategoryRule::new( vec!["Other".into()], Rule::from(Regex::new(r"nonmatching").unwrap()), ), @@ -291,7 +342,7 @@ fn test_categorize_uncategorized() { .insert("test".into(), serde_json::json!("just a test")); let mut events = vec![e]; - let rules: Vec<(Vec, Rule)> = vec![( + let rules: Vec = vec![CategoryRule::new( vec!["Non-matching".into(), "test".into()], Rule::from(Regex::new(r"not going to match").unwrap()), )]; @@ -304,6 +355,153 @@ fn test_categorize_uncategorized() { ); } +#[cfg(test)] +fn event_with_data(value: &str) -> Event { + let mut e = Event::default(); + e.data.insert("test".into(), serde_json::json!(value)); + e +} + +#[cfg(test)] +fn category_of(events: &[Event]) -> &serde_json::Value { + events.first().unwrap().data.get("$category").unwrap() +} + +#[test] +fn test_categorize_depth_wins_without_priority() { + // Reported scenario from ActivityWatch/aw-server-rust#597: a deeper nested + // match still beats a shallower match when neither rule sets priority. + // Category A (depth 1) and Category B → B1 (depth 2) both match. + let events = categorize( + vec![event_with_data("just a test")], + &[ + CategoryRule::new(vec!["A".into()], Rule::from(Regex::new(r"test").unwrap())), + CategoryRule::new( + vec!["B".into(), "B1".into()], + Rule::from(Regex::new(r"test").unwrap()), + ), + ], + ); + assert_eq!(category_of(&events), &serde_json::json!(vec!["B", "B1"])); +} + +#[test] +fn test_categorize_explicit_priority_overrides_depth() { + // The same #597 tree, but A is given a higher priority than B1's default + // (depth 2 → 20). Organizational nesting no longer forces B1 to win. + let events = categorize( + vec![event_with_data("just a test")], + &[ + CategoryRule::new(vec!["A".into()], Rule::from(Regex::new(r"test").unwrap())) + .with_priority(25), + CategoryRule::new( + vec!["B".into(), "B1".into()], + Rule::from(Regex::new(r"test").unwrap()), + ), + ], + ); + assert_eq!(category_of(&events), &serde_json::json!(vec!["A"])); +} + +#[test] +fn test_categorize_inter_level_priority() { + // depth * 10 leaves integers between levels: 15 beats a depth-1 default + // (10) but loses to a depth-2 default (20). + let between = categorize( + vec![event_with_data("just a test")], + &[ + CategoryRule::new(vec!["A".into()], Rule::from(Regex::new(r"test").unwrap())), + CategoryRule::new(vec!["A2".into()], Rule::from(Regex::new(r"test").unwrap())) + .with_priority(15), + ], + ); + assert_eq!(category_of(&between), &serde_json::json!(vec!["A2"])); + + let still_loses_to_deeper = categorize( + vec![event_with_data("just a test")], + &[ + CategoryRule::new(vec!["A2".into()], Rule::from(Regex::new(r"test").unwrap())) + .with_priority(15), + CategoryRule::new( + vec!["B".into(), "B1".into()], + Rule::from(Regex::new(r"test").unwrap()), + ), + ], + ); + assert_eq!( + category_of(&still_loses_to_deeper), + &serde_json::json!(vec!["B", "B1"]) + ); +} + +#[test] +fn test_categorize_lower_priority_loses_to_default_depth() { + // A deep rule can also be demoted below a shallow rule's default + // (depth * 10) by setting an explicit lower priority on the deep rule. + let events = categorize( + vec![event_with_data("just a test")], + &[ + CategoryRule::new(vec!["A".into()], Rule::from(Regex::new(r"test").unwrap())), + CategoryRule::new( + vec!["B".into(), "B1".into()], + Rule::from(Regex::new(r"test").unwrap()), + ) + .with_priority(0), + ], + ); + assert_eq!(category_of(&events), &serde_json::json!(vec!["A"])); +} + +#[test] +fn test_categorize_equal_priority_keeps_later_match() { + // Preserve the historical `>=` later-wins rule when ranks tie. + let events = categorize( + vec![event_with_data("just a test")], + &[ + CategoryRule::new( + vec!["First".into()], + Rule::from(Regex::new(r"test").unwrap()), + ) + .with_priority(5), + CategoryRule::new( + vec!["Second".into()], + Rule::from(Regex::new(r"test").unwrap()), + ) + .with_priority(5), + ], + ); + assert_eq!(category_of(&events), &serde_json::json!(vec!["Second"])); +} + +#[test] +fn test_categorize_negative_priority_still_beats_uncategorized() { + let events = categorize( + vec![event_with_data("just a test")], + &[ + CategoryRule::new(vec!["Low".into()], Rule::from(Regex::new(r"test").unwrap())) + .with_priority(-100), + ], + ); + assert_eq!(category_of(&events), &serde_json::json!(vec!["Low"])); +} + +#[test] +fn test_categorize_empty_category_keeps_uncategorized() { + // Empty path used to lose to Uncategorized (depth 0 < 1). The MIN-rank + // sentinel must not let it replace the fallback. + let events = categorize( + vec![event_with_data("just a test")], + &[CategoryRule::new( + vec![], + Rule::from(Regex::new(r"test").unwrap()), + )], + ); + assert_eq!( + category_of(&events), + &serde_json::json!(vec!["Uncategorized"]) + ); +} + #[test] fn test_categorize_cache_correctness() { // Verifies that the deduplication cache produces the same result as @@ -326,12 +524,12 @@ fn test_categorize_cache_correctness() { .chain(std::iter::repeat(base.clone()).take(50)) .collect(); - let rules: Vec<(Vec, Rule)> = vec![ - ( + let rules: Vec = vec![ + CategoryRule::new( vec!["Browser".into()], Rule::Regex(RegexRule::new("firefox", true, Some(vec!["app".into()])).unwrap()), ), - ( + CategoryRule::new( vec!["Terminal".into()], Rule::Regex(RegexRule::new("terminal", true, Some(vec!["app".into()])).unwrap()), ),