feat(classify): optional priority field for category rules

aw_query/functions.py

@@ -333,12 +333,18 @@ def q2_nop():
 @q2_function(categorize)
 @q2_typecheck
 def q2_categorize(events: list, classes: list):
-    classes = [(_cls, Rule(rule_dict)) for _cls, rule_dict in classes]
+    try:
+        classes = [(_cls, Rule(rule_dict)) for _cls, rule_dict in classes]
+    except ValueError as exc:
+        raise QueryFunctionException(str(exc)) from None
     return categorize(events, classes)
 
 
 @q2_function(tag)
 @q2_typecheck
 def q2_tag(events: list, classes: list):
-    classes = [(_cls, Rule(rule_dict)) for _cls, rule_dict in classes]
+    try:
+        classes = [(_cls, Rule(rule_dict)) for _cls, rule_dict in classes]
+    except ValueError as exc:
+        raise QueryFunctionException(str(exc)) from None
     return tag(events, classes)

aw_transform/classify.py

@@ -1,6 +1,5 @@
 import json
 from typing import Pattern, List, Iterable, Tuple, Dict, Optional, Any
-from functools import reduce
 import re
 
 from aw_core import Event
@@ -10,14 +9,29 @@
 Category = List[str]
 
 
+def _parse_optional_priority(rules: Dict[str, Any]) -> Optional[int]:
+    if "priority" in rules:
+        val = rules["priority"]
+    elif "weight" in rules:
+        val = rules["weight"]
+    else:
+        return None
+    # bool is a subclass of int
+    if isinstance(val, bool) or not isinstance(val, int):
+        raise ValueError("priority/weight must be an integer")
+    return val
+
+
 class Rule:
     regex: Optional[Pattern]
     select_keys: Optional[List[str]]
     ignore_case: bool
+    priority: Optional[int]
 
     def __init__(self, rules: Dict[str, Any]) -> None:
         self.select_keys = rules.get("select_keys", None)
         self.ignore_case = rules.get("ignore_case", False)
+        self.priority = _parse_optional_priority(rules)
 
         # NOTE: Also checks that the regex isn't an empty string (which would erroneously match everything)
         regex_str = rules.get("regex", None)
@@ -52,15 +66,15 @@ def categorize(
             key = str(id(e.data))
         if key not in cache:
             cache[key] = _pick_category(
-                [_cls for _cls, rule in classes if rule.match(e)]
+                [(_cls, rule) for _cls, rule in classes if rule.match(e)]
             )
         e.data["$category"] = list(cache[key])
     return events
 
 
 def _categorize_one(e: Event, classes: List[Tuple[Category, Rule]]) -> Event:
     e.data["$category"] = _pick_category(
-        [_cls for _cls, rule in classes if rule.match(e)]
+        [(_cls, rule) for _cls, rule in classes if rule.match(e)]
     )
     return e
 
@@ -83,11 +97,26 @@ def _tag_one(e: Event, classes: List[Tuple[Tag, Rule]]) -> Event:
     return e
 
 
-def _pick_category(tags: Iterable[Category]) -> Category:
-    return reduce(_pick_deepest_cat, tags, ["Uncategorized"])
-
-
-def _pick_deepest_cat(t1: Category, t2: Category) -> Category:
-    # t1 will be the accumulator when used in reduce
-    # Always bias against t1, since it could be "Uncategorized"
-    return t2 if len(t2) >= len(t1) else t1
+def _effective_rank(category: Category, rule: Rule) -> int:
+    # 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.
+    if rule.priority is not None:
+        return rule.priority
+    return len(category) * 10
+
+
+def _pick_category(matches: Iterable[Tuple[Category, Rule]]) -> Category:
+    category: Category = ["Uncategorized"]
+    rank: Optional[int] = None
+    for cat, rule in matches:
+        if not cat:
+            continue
+        item_rank = _effective_rank(cat, rule)
+        # None means no match yet, so any non-empty category wins — including
+        # an explicit priority below a signed 64-bit floor. Equal ranks keep
+        # the later match (same contract as the old depth-only `>=`).
+        if rank is None or item_rank >= rank:
+            category = cat
+            rank = item_rank
+    return category

tests/test_query2.py

@@ -338,6 +338,19 @@ def test_query2_merge_subwatcher_fields_invalid_conflict():
         query(qname, example_query, starttime, endtime, ds)
 
 
+def test_query2_categorize_invalid_priority():
+    ds = mock_ds
+    qname = "asd"
+    starttime = iso8601.parse_date("1970-01-01")
+    endtime = iso8601.parse_date("1970-01-02")
+    example_query = """
+        events = [];
+        RETURN = categorize(events, [[["test"], {"regex": "test", "priority": "high"}]]);
+    """
+    with pytest.raises(QueryFunctionException, match="integer"):
+        query(qname, example_query, starttime, endtime, ds)
+
+
 @pytest.mark.parametrize("datastore", param_datastore_objects())
 def test_query2_function_in_function(datastore):
     qname = "asd"
@@ -614,6 +627,39 @@ def test_query2_query_categorize(datastore):
         datastore.delete_bucket(bid)
 
 
+@pytest.mark.parametrize("datastore", param_datastore_objects())
+def test_query2_query_categorize_priority(datastore):
+    bid = "test_bucket_priority"
+    qname = "test"
+    starttime = iso8601.parse_date("1970")
+    endtime = starttime + timedelta(hours=1)
+
+    example_query = rf"""
+    events = query_bucket("{bid}");
+    events = categorize(events, [
+                [["A"], {{"regex": "test", "priority": 25}}],
+                [["B", "B1"], {{"regex": "test"}}]
+            ]);
+    RETURN = events;
+    """
+    try:
+        bucket = datastore.create_bucket(
+            bucket_id=bid, type="test", client="test", hostname="test", name="asd"
+        )
+        bucket.insert(
+            Event(
+                data={"label": "test"},
+                timestamp=starttime,
+                duration=timedelta(seconds=1),
+            )
+        )
+        result = query(qname, example_query, starttime, endtime, datastore)
+        assert len(result) == 1
+        assert result[0].data["$category"] == ["A"]
+    finally:
+        datastore.delete_bucket(bid)
+
+
 @pytest.mark.parametrize("datastore", param_datastore_objects())
 def test_query2_flood_with_pulsetime(datastore):
     """Test that flood() accepts an optional pulsetime argument (fixes WorkReport query bug)."""

tests/test_transforms.py

@@ -418,6 +418,121 @@ def test_categorize():
     assert events[3].data["$category"] == ["Uncategorized"]
 
 
+def _event(value: str = "just a test") -> Event:
+    return Event(timestamp=datetime.now(timezone.utc), duration=0, data={"key": value})
+
+
+def test_categorize_depth_wins_without_priority():
+    events = categorize(
+        [_event()],
+        [
+            (["A"], Rule({"regex": "test"})),
+            (["B", "B1"], Rule({"regex": "test"})),
+        ],
+    )
+    assert events[0].data["$category"] == ["B", "B1"]
+
+
+def test_categorize_explicit_priority_overrides_depth():
+    # Default for B1 is depth 2 → 20; 25 beats it.
+    events = categorize(
+        [_event()],
+        [
+            (["A"], Rule({"regex": "test", "priority": 25})),
+            (["B", "B1"], Rule({"regex": "test"})),
+        ],
+    )
+    assert events[0].data["$category"] == ["A"]
+
+
+def test_categorize_weight_alias():
+    events = categorize(
+        [_event()],
+        [
+            (["A"], Rule({"regex": "test", "weight": 25})),
+            (["B", "B1"], Rule({"regex": "test"})),
+        ],
+    )
+    assert events[0].data["$category"] == ["A"]
+
+
+def test_categorize_inter_level_priority():
+    between = categorize(
+        [_event()],
+        [
+            (["A"], Rule({"regex": "test"})),
+            (["A2"], Rule({"regex": "test", "priority": 15})),
+        ],
+    )
+    assert between[0].data["$category"] == ["A2"]
+
+    still_loses_to_deeper = categorize(
+        [_event()],
+        [
+            (["A2"], Rule({"regex": "test", "priority": 15})),
+            (["B", "B1"], Rule({"regex": "test"})),
+        ],
+    )
+    assert still_loses_to_deeper[0].data["$category"] == ["B", "B1"]
+
+
+def test_categorize_lower_priority_loses_to_default_depth():
+    events = categorize(
+        [_event()],
+        [
+            (["A"], Rule({"regex": "test"})),
+            (["B", "B1"], Rule({"regex": "test", "priority": 0})),
+        ],
+    )
+    assert events[0].data["$category"] == ["A"]
+
+
+def test_categorize_equal_priority_keeps_later_match():
+    events = categorize(
+        [_event()],
+        [
+            (["First"], Rule({"regex": "test", "priority": 5})),
+            (["Second"], Rule({"regex": "test", "priority": 5})),
+        ],
+    )
+    assert events[0].data["$category"] == ["Second"]
+
+
+def test_categorize_negative_priority_still_beats_uncategorized():
+    events = categorize(
+        [_event()],
+        [(["Low"], Rule({"regex": "test", "priority": -100}))],
+    )
+    assert events[0].data["$category"] == ["Low"]
+
+
+def test_categorize_priority_below_i64_min_still_beats_uncategorized():
+    # Python ints are unbounded; a signed-64-bit fallback sentinel would
+    # incorrectly keep Uncategorized for values below -(2**63).
+    events = categorize(
+        [_event()],
+        [(["Low"], Rule({"regex": "test", "priority": -(2**63) - 1}))],
+    )
+    assert events[0].data["$category"] == ["Low"]
+
+
+def test_categorize_empty_category_keeps_uncategorized():
+    events = categorize(
+        [_event()],
+        [([], Rule({"regex": "test"}))],
+    )
+    assert events[0].data["$category"] == ["Uncategorized"]
+
+
+def test_rule_invalid_priority():
+    with pytest.raises(ValueError, match="integer"):
+        Rule({"regex": "test", "priority": 1.5})
+    with pytest.raises(ValueError, match="integer"):
+        Rule({"regex": "test", "priority": "high"})
+    with pytest.raises(ValueError, match="integer"):
+        Rule({"regex": "test", "priority": True})
+
+
 def test_categorize_cache_correctness():
     """Cache reuses category for identical data; distinct data gets its own category."""
     now = datetime.now(timezone.utc)