structure execution plan

redisgraph/execution_plan.py

@@ -0,0 +1,61 @@
+class Operation:
+    def __init__(self, name, args=None):
+        self.name = name
+        self.args = args
+        self.children = []
+
+    def append_child(self, child):
+        self.children.append(child)
+        return self
+
+    def __eq__(self, o: object) -> bool:
+        if not isinstance(o, Operation):
+            return False
+
+        if self.name != o.name or self.args != o.args or len(self.children) != len(o.children):
+            return False
+
+        for i in range(len(self.children)):
+            if not self.children[i] == o.children[i]:
+                return False
+
+        return True
+
+
+class ExecutionPlan:
+    def __init__(self, plan):
+        self.plan = plan
+        self.structured_plan = self._operation_tree()
+
+    def __str__(self) -> str:
+        return "\n".join(self.plan)
+
+    def _operation_tree(self):
+        i = 0
+        level = 0
+        stack = []
+        current = None
+
+        while i < len(self.plan):
+            op = self.plan[i]
+            op_level = op.count("    ")
+            if op_level == level:
+                args = op.split("|")
+                current = Operation(args[0].strip(), None if len(args) == 1 else args[1].strip())
+                i += 1
+            elif op_level == level + 1:
+                args = op.split("|")
+                child = Operation(args[0].strip(), None if len(args) == 1 else args[1].strip())
+                current.children.append(child)
+                stack.append(current)
+                current = child
+                level += 1
+                i += 1
+            elif op_level < level:
+                levels_back = level - op_level + 1
+                for _ in range(levels_back):
+                    current = stack.pop()
+                level -= levels_back
+            else:
+                raise Exception("corrupted plan")
+        return stack[0]

redisgraph/graph.py

@@ -3,6 +3,7 @@
 from redisgraph.util import random_string, quote_string, stringify_param_value
 from redisgraph.query_result import QueryResult
 from redisgraph.exceptions import VersionMismatchException
+from redisgraph.execution_plan import ExecutionPlan
 
 
 class Graph:
@@ -215,9 +216,6 @@ def query(self, q, params=None, timeout=None, read_only=False):
             # re-issue query
             return self.query(q, params, timeout, read_only)
 
-    def _execution_plan_to_string(self, plan):
-        return "\n".join(plan)
-
     def execution_plan(self, query, params=None):
         """
         Get the execution plan for given query,
@@ -231,7 +229,7 @@ def execution_plan(self, query, params=None):
             query = self._build_params_header(params) + query
 
         plan = self.redis_con.execute_command("GRAPH.EXPLAIN", self.name, query)
-        return self._execution_plan_to_string(plan)
+        return ExecutionPlan(plan)
 
     def delete(self):
         """

tests/functional/test_all.py

@@ -1,4 +1,5 @@
 import unittest
+from redisgraph.execution_plan import Operation
 from tests.utils import base
 
 import redis
@@ -252,7 +253,15 @@ def test_execution_plan(self):
 
         result = redis_graph.execution_plan("MATCH (r:Rider)-[:rides]->(t:Team) WHERE t.name = $name RETURN r.name, t.name, $params", {'name': 'Yehuda'})
         expected = "Results\n    Project\n        Conditional Traverse | (t:Team)->(r:Rider)\n            Filter\n                Node By Label Scan | (t:Team)"
-        self.assertEqual(result, expected)
+        self.assertEqual(str(result), expected)
+
+        expected = Operation('Results') \
+            .append_child(Operation('Project')
+                          .append_child(Operation('Conditional Traverse', "(t:Team)->(r:Rider)")
+                                        .append_child(Operation("Filter")
+                                                      .append_child(Operation('Node By Label Scan', "(t:Team)")))))
+
+        self.assertEqual(result.structured_plan, expected)
 
         redis_graph.delete()