[core][feat] Allow filtering edge properties (#2186)

Author: aquamatthias

fixcore/fixcore/db/arango_query.py

@@ -877,14 +877,15 @@ def traversal_filter(cl: WithClause, in_crs: str, depth: int) -> str:
             )
             return out
 
-        def inout(in_crsr: str, start: int, until: int, edge_type: str, direction: str) -> str:
+        def inout(
+            in_crsr: str, start: int, until: int, edge_type: str, direction: str, edge_filter: Optional[Term]
+        ) -> str:
             nonlocal query_part
             in_c = ctx.next_crs("io_in")
             out = ctx.next_crs("io_out")
             out_crsr = ctx.next_crs("io_crs")
             e = ctx.next_crs("io_link")
             unique = "uniqueEdges: 'path'" if with_edges else "uniqueVertices: 'global'"
-            link_str = f", {e}" if with_edges else ""
             dir_bound = "OUTBOUND" if direction == Direction.outbound else "INBOUND"
             inout_result = (
                 # merge edge and vertex properties - will be split in the output transformer
@@ -899,12 +900,17 @@ def inout(in_crsr: str, start: int, until: int, edge_type: str, direction: str)
                 graph_cursor = in_c
                 outer_for = f"FOR {in_c} in {in_crsr} "
 
+            # optional: add the edge filter to the query
+            pre, fltr, post = term(e, edge_filter) if edge_filter else (None, None, None)
+            pre_string = " " + pre if pre else ""
+            post_string = f" AND ({post})" if post else ""
+            filter_string = "" if not fltr and not post_string else f"{pre_string} FILTER {fltr}{post_string}"
             query_part += (
                 f"LET {out} =({outer_for}"
                 # suggested by jsteemann: use crs._id instead of crs (stored in the view and more efficient)
-                f"FOR {out_crsr}{link_str} IN {start}..{until} {dir_bound} {graph_cursor}._id "
-                f"`{db.edge_collection(edge_type)}` OPTIONS {{ bfs: true, {unique} }} "
-                f"RETURN DISTINCT {inout_result}) "
+                f"FOR {out_crsr}, {e} IN {start}..{until} {dir_bound} {graph_cursor}._id "
+                f"`{db.edge_collection(edge_type)}` OPTIONS {{ bfs: true, {unique} }}{filter_string} "
+                f"RETURN DISTINCT {inout_result})"
             )
             return out
 
@@ -913,12 +919,12 @@ def navigation(in_crsr: str, nav: Navigation) -> str:
             all_walks = []
             if nav.direction == Direction.any:
                 for et in nav.edge_types:
-                    all_walks.append(inout(in_crsr, nav.start, nav.until, et, Direction.inbound))
+                    all_walks.append(inout(in_crsr, nav.start, nav.until, et, Direction.inbound, nav.edge_filter))
                 for et in nav.maybe_two_directional_outbound_edge_type or nav.edge_types:
-                    all_walks.append(inout(in_crsr, nav.start, nav.until, et, Direction.outbound))
+                    all_walks.append(inout(in_crsr, nav.start, nav.until, et, Direction.outbound, nav.edge_filter))
             else:
                 for et in nav.edge_types:
-                    all_walks.append(inout(in_crsr, nav.start, nav.until, et, nav.direction))
+                    all_walks.append(inout(in_crsr, nav.start, nav.until, et, nav.direction, nav.edge_filter))
 
             if len(all_walks) == 1:
                 return all_walks[0]

fixcore/fixcore/query/model.py

@@ -605,6 +605,7 @@ class Navigation:
     maybe_edge_types: Optional[List[EdgeType]] = None
     direction: str = Direction.outbound
     maybe_two_directional_outbound_edge_type: Optional[List[EdgeType]] = None
+    edge_filter: Optional[Term] = None
 
     @property
     def edge_types(self) -> List[EdgeType]:
@@ -617,14 +618,18 @@ def __str__(self) -> str:
         mo = self.maybe_two_directional_outbound_edge_type
         depth = ("" if start == 1 else f"[{start}]") if start == until and not mo else f"[{start}:{until_str}]"
         out_nav = ",".join(mo) if mo else ""
-        nav = f'{",".join(self.edge_types)}{depth}{out_nav}'
+        fltr = f"{{{self.edge_filter}}}" if self.edge_filter else ""
+        nav = f'{",".join(self.edge_types)}{depth}{fltr}{out_nav}'
         if self.direction == Direction.outbound:
             return f"-{nav}->"
         elif self.direction == Direction.inbound:
             return f"<-{nav}-"
         else:
             return f"<-{nav}->"
 
+    def change_variable(self, fn: Callable[[str], str]) -> Navigation:
+        return evolve(self, edge_filter=self.edge_filter.change_variable(fn)) if self.edge_filter else self
+
 
 NavigateUntilRoot = Navigation(
     start=1, until=Navigation.Max, maybe_edge_types=[EdgeTypes.default], direction=Direction.inbound
@@ -740,6 +745,7 @@ def change_variable(self, fn: Callable[[str], str]) -> Part:
             term=self.term.change_variable(fn),
             with_clause=self.with_clause.change_variable(fn) if self.with_clause else None,
             sort=[sort.change_variable(fn) for sort in self.sort],
+            navigation=self.navigation.change_variable(fn) if self.navigation else None,
         )
 
     # ancestor.some_type.reported.prop -> MergeQuery
@@ -1012,17 +1018,40 @@ def filter_with(self, clause: WithClause) -> Query:
         first_part = evolve(self.parts[0], with_clause=clause)
         return evolve(self, parts=[first_part, *self.parts[1:]])
 
-    def traverse_out(self, start: int = 1, until: int = 1, edge_type: EdgeType = EdgeTypes.default) -> Query:
-        return self.traverse(start, until, edge_type, Direction.outbound)
-
-    def traverse_in(self, start: int = 1, until: int = 1, edge_type: EdgeType = EdgeTypes.default) -> Query:
-        return self.traverse(start, until, edge_type, Direction.inbound)
-
-    def traverse_inout(self, start: int = 1, until: int = 1, edge_type: EdgeType = EdgeTypes.default) -> Query:
-        return self.traverse(start, until, edge_type, Direction.any)
+    def traverse_out(
+        self,
+        start: int = 1,
+        until: int = 1,
+        edge_type: EdgeType = EdgeTypes.default,
+        edge_filter: Optional[Term] = None,
+    ) -> Query:
+        return self.traverse(start, until, edge_type, Direction.outbound, edge_filter)
+
+    def traverse_in(
+        self,
+        start: int = 1,
+        until: int = 1,
+        edge_type: EdgeType = EdgeTypes.default,
+        edge_filter: Optional[Term] = None,
+    ) -> Query:
+        return self.traverse(start, until, edge_type, Direction.inbound, edge_filter)
+
+    def traverse_inout(
+        self,
+        start: int = 1,
+        until: int = 1,
+        edge_type: EdgeType = EdgeTypes.default,
+        edge_filter: Optional[Term] = None,
+    ) -> Query:
+        return self.traverse(start, until, edge_type, Direction.any, edge_filter)
 
     def traverse(
-        self, start: int, until: int, edge_type: EdgeType = EdgeTypes.default, direction: str = Direction.outbound
+        self,
+        start: int,
+        until: int,
+        edge_type: EdgeType = EdgeTypes.default,
+        direction: str = Direction.outbound,
+        edge_filter: Optional[Term] = None,
     ) -> Query:
         parts = self.parts.copy()
         p0 = parts[0]
@@ -1034,9 +1063,15 @@ def traverse(
                 parts[0] = evolve(p0, navigation=evolve(p0.navigation, start=start_m, until=until_m))
             # this is another traversal: so we need to start a new part
             else:
-                parts.insert(0, Part(AllTerm(), navigation=Navigation(start, until, [edge_type], direction)))
+                parts.insert(
+                    0,
+                    Part(
+                        AllTerm(),
+                        navigation=Navigation(start, until, [edge_type], direction, edge_filter=edge_filter),
+                    ),
+                )
         else:
-            parts[0] = evolve(p0, navigation=Navigation(start, until, [edge_type], direction))
+            parts[0] = evolve(p0, navigation=Navigation(start, until, [edge_type], direction, edge_filter=edge_filter))
         return evolve(self, parts=parts)
 
     def group_by(self, variables: List[AggregateVariable], funcs: List[AggregateFunction]) -> Query:

fixcore/fixcore/query/query_parser.py

@@ -224,47 +224,62 @@ def range_parser() -> Parser:
     return start, end
 
 
-edge_type_p = lexeme(regex("[A-Za-z][A-Za-z0-9_]*"))
+edge_type_p = reduce(lambda x, y: x | y, [lexeme(string(a)) for a in EdgeTypes.all])
 
 
 @make_parser
 def edge_type_parser() -> Parser:
     edge_types = yield edge_type_p.sep_by(comma_p).map(set)
-    for et in edge_types:
-        if et not in EdgeTypes.all:
-            raise AttributeError(f"Given EdgeType is not known: {et}")
     return list(edge_types)
 
 
+@make_parser
+def combined_edge_term() -> Parser:
+    left = yield simple_edge_term_p
+    result = left
+    while True:
+        op = yield bool_op_p.optional()
+        if op is None:
+            break
+        right = yield simple_edge_term_p
+        result = CombinedTerm(result, op, right)
+    return result
+
+
+leaf_edge_term_p = predicate_term | context_term | not_term | match_all_term
+simple_edge_term_p = (lparen_p >> combined_edge_term << rparen_p) | leaf_edge_term_p
+edge_term_p = l_curly_p >> combined_edge_term << r_curly_p
+
+
 @make_parser
 def edge_definition_parser() -> Parser:
     edge_types = yield edge_type_parser
     maybe_range = yield range_parser.optional()
     start, until = maybe_range if maybe_range else (1, 1)
+    edge_filter = yield edge_term_p.optional()
     after_bracket_edge_types = yield edge_type_parser
     if edge_types and after_bracket_edge_types:
         raise AttributeError("Edge types can not be defined before and after the [start,until] definition.")
-    return start, until, edge_types or after_bracket_edge_types
+    return Navigation(start, until, maybe_edge_types=edge_types or after_bracket_edge_types, edge_filter=edge_filter)
 
 
 @make_parser
 def two_directional_edge_definition_parser() -> Parser:
     edge_types = yield edge_type_parser
     maybe_range = yield range_parser.optional()
+    edge_filter = yield edge_term_p.optional()
     outbound_edge_types = yield edge_type_parser
     start, until = maybe_range if maybe_range else (1, 1)
-    return start, until, edge_types, outbound_edge_types
+    return Navigation(start, until, edge_types, Direction.any, outbound_edge_types, edge_filter)
 
 
 out_p = lexeme(string("-") >> edge_definition_parser << string("->")).map(
-    lambda nav: Navigation(nav[0], nav[1], nav[2], Direction.outbound)
+    lambda nav: evolve(nav, direction=Direction.outbound)
 )
 in_p = lexeme(string("<-") >> edge_definition_parser << string("-")).map(
-    lambda nav: Navigation(nav[0], nav[1], nav[2], Direction.inbound)
-)
-in_out_p = lexeme(string("<-") >> two_directional_edge_definition_parser << string("->")).map(
-    lambda nav: Navigation(nav[0], nav[1], nav[2], Direction.any, nav[3])
+    lambda nav: evolve(nav, direction=Direction.inbound)
 )
+in_out_p = lexeme(string("<-") >> two_directional_edge_definition_parser << string("->"))
 navigation_parser = in_out_p | out_p | in_p
 
 tag_parser = lexeme(string("#") >> literal_p).optional()

fixcore/tests/fixcore/__init__.py

@@ -13,7 +13,7 @@ def create_graph(bla_text: str, width: int = 10) -> MultiDiGraph:
 
     def add_edge(from_node: str, to_node: str, edge_type: EdgeType = EdgeTypes.default) -> None:
         key = GraphAccess.edge_key(from_node, to_node, edge_type)
-        graph.add_edge(from_node, to_node, key, edge_type=edge_type)
+        graph.add_edge(from_node, to_node, key, reported=dict(a=1, b=[{"c": 1, "d": 2}, {"c": 2, "d": 2}]))
 
     def add_node(uid: str, kind: str, node: Optional[Json] = None, replace: bool = False) -> None:
         reported = {**(node if node else to_js(Foo(uid))), "kind": kind}

fixcore/tests/fixcore/db/graphdb_test.py

@@ -2,6 +2,7 @@
 import string
 from abc import ABC, abstractmethod
 from datetime import date, datetime, timedelta
+from functools import partial
 from random import SystemRandom
 from typing import List, Optional, Any, Dict, cast, AsyncIterator, Tuple, Union, Literal
 
@@ -945,6 +946,22 @@ async def assert_security(
     await assert_security("change4", 10, 2, reopen=1, added_vulnerable=2)
 
 
+async def test_graph_edge_filter(filled_graph_db: GraphDB, foo_model: Model) -> None:
+    query_list = partial(query_list_on, filled_graph_db, foo_model)
+    assert len(await query_list("is(foo) and id==9 -{a=1}->")) == 10
+    assert len(await query_list("is(foo) and id==9 -{a=1 and b[*].{c=1 and d=2}}->")) == 10
+    assert len(await query_list("is(foo) and id==9 -{a=1 and b[*].{c=2 and d=2}}->")) == 10
+    assert len(await query_list("is(foo) and id==9 -{a=1 and b[*].{c=3 and d=2}}->")) == 0
+    assert len(await query_list("is(foo) and id==9 -{a=2 and b[*].{c=2 and d=2}}->")) == 0
+    assert len(await query_list("is(foo) and id==9 -{a=1 and b[*].{c=1 and d=2} and b[0].c=1}->")) == 10
+
+
+async def query_list_on(db: GraphDB, model: Model, q: str) -> List[Json]:
+    query = parse_query(q).on_section("reported")
+    async with await db.search_list(QueryModel(query, model)) as cursor:
+        return [entry async for entry in cursor]
+
+
 def to_json(obj: BaseResource) -> Json:
     return {"kind": obj.kind(), **to_js(obj)}
 

fixcore/tests/fixcore/hypothesis_extension.py

@@ -14,35 +14,20 @@
     lists,
     composite,
     sampled_from,
+    DrawFn,
 )
 
 from fixcore.model.resolve_in_graph import NodePath
 from fixcore.types import JsonElement, Json
 from fixcore.util import value_in_path, interleave
 
 T = TypeVar("T")
-UD = Callable[[SearchStrategy[Any]], Any]
 
 
 def optional(st: SearchStrategy[T]) -> SearchStrategy[Optional[T]]:
     return st | just(None)
 
 
-class Drawer:
-    """
-    Only here for getting a drawer for typed drawings.
-    """
-
-    def __init__(self, hypo_drawer: Callable[[SearchStrategy[Any]], Any]):
-        self._drawer = hypo_drawer
-
-    def draw(self, st: SearchStrategy[T]) -> T:
-        return cast(T, self._drawer(st))
-
-    def optional(self, st: SearchStrategy[T]) -> Optional[T]:
-        return self.draw(optional(st))
-
-
 any_ws_digits_string = text(alphabet=string.ascii_letters + " " + string.digits, min_size=0, max_size=10)
 any_string = text(alphabet=string.ascii_letters, min_size=3, max_size=10)
 kind_gen = sampled_from(["volume", "instance", "load_balancer", "volume_type"])
@@ -51,8 +36,8 @@ def optional(self, st: SearchStrategy[T]) -> Optional[T]:
 
 
 @composite
-def json_element_gen(ud: UD) -> JsonElement:
-    return cast(JsonElement, ud(json_object_gen | json_simple_element_gen | json_array_gen))
+def json_element_gen(draw: DrawFn) -> JsonElement:
+    return cast(JsonElement, draw(json_object_gen | json_simple_element_gen | json_array_gen))
 
 
 json_simple_element_gen = any_ws_digits_string | booleans() | integers(min_value=0, max_value=100000) | just(None)
@@ -61,14 +46,13 @@ def json_element_gen(ud: UD) -> JsonElement:
 
 
 @composite
-def node_gen(ud: UD) -> Json:
-    d = Drawer(ud)
-    uid = d.draw(any_string)
-    name = d.draw(any_string)
-    kind = d.draw(kind_gen)
-    reported = d.draw(json_object_gen)
-    metadata = d.draw(json_object_gen)
-    desired = d.draw(json_object_gen)
+def node_gen(draw: DrawFn) -> Json:
+    uid = draw(any_string)
+    name = draw(any_string)
+    kind = draw(kind_gen)
+    reported = draw(json_object_gen)
+    metadata = draw(json_object_gen)
+    desired = draw(json_object_gen)
     return {
         "id": uid,
         "kinds": [kind],

fixcore/tests/fixcore/model/db_updater_test.py

@@ -46,8 +46,8 @@ def to_b(a: Any) -> bytes:
 
         for node in graph.nodes():
             yield to_b(graph.nodes[node])
-        for from_node, to_node, data in graph.edges(data=True):
-            yield to_b({"from": from_node, "to": to_node, "edge_type": data["edge_type"]})
+        for from_node, to_node, key in graph.edges(keys=True):
+            yield to_b({"from": from_node, "to": to_node, "edge_type": key.edge_type})
         yield to_b(
             {"from_selector": {"node_id": "id_123"}, "to_selector": {"node_id": "id_456"}, "edge_type": "delete"}
         )