diff --git a/aw_query/functions.py b/aw_query/functions.py index 75b620b..00a85fc 100644 --- a/aw_query/functions.py +++ b/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) diff --git a/aw_transform/classify.py b/aw_transform/classify.py index 7f3ba4f..35a33b8 100644 --- a/aw_transform/classify.py +++ b/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,7 +66,7 @@ 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 @@ -60,7 +74,7 @@ def categorize( 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 diff --git a/tests/test_query2.py b/tests/test_query2.py index 98e25d6..f7e864d 100644 --- a/tests/test_query2.py +++ b/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).""" diff --git a/tests/test_transforms.py b/tests/test_transforms.py index b7bf883..c30e97d 100644 --- a/tests/test_transforms.py +++ b/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)