test: assert the client parses execution plans, and fix a plan parser bug

falkordb/asyncio/cluster.py

@@ -1,3 +1,4 @@
+import inspect
 import socket
 
 import redis as sync_redis  # type: ignore[import-not-found]
@@ -22,6 +23,13 @@ def Is_Cluster(conn: redis.Redis):
     if pool.connection_class is redis.UnixDomainSocketConnection:
         kwargs["unix_socket_path"] = kwargs.pop("path")
 
+    # Keep only the parameters the synchronous constructor actually accepts.
+    # redis-py stores internal state in ``connection_kwargs`` that is not part
+    # of the ``Redis.__init__`` signature — redis 8.1.0 added ``himport_registry``
+    # there — and forwarding those raises TypeError before any I/O happens.
+    accepted = inspect.signature(sync_redis.Redis.__init__).parameters
+    kwargs = {k: v for k, v in kwargs.items() if k in accepted}
+
     # Create a synchronous Redis client with the same parameters
     # as the connection pool just to keep Is_Cluster synchronous
     info = sync_redis.Redis(**kwargs).info(section="server")

falkordb/execution_plan.py

@@ -307,11 +307,15 @@ def create_operation(args):
                 self.operations[child.name].append(child)
 
                 if current:
-                    current = stack.pop()
-                    current.append_child(child)
+                    # attach the sibling to the parent and keep the parent on
+                    # the stack, so any further sibling is attached to it too
+                    parent = stack.pop()
+                    parent.append_child(child)
+                    stack.append(parent)
+                else:
+                    stack.append(child)
                 current = child
                 i += 1
-                stack.append(child)
             elif op_level == level + 1:
                 # if the operation is child of the current operation
                 # add it as child and set as current operation

tests/plan_helpers.py

@@ -0,0 +1,172 @@
+"""Assertions for execution plans.
+
+Two levels are available.
+
+``assert_plan_shape`` pins the exact operation tree — every operation name and
+how deep it sits. Use it wherever the C and Rust engines compile a query to the
+same plan, which is the common case.
+
+``assert_parsed_plan`` only checks that the client turned the reply into a tree
+faithfully: every line became one operation, nested as deep as it was indented.
+Use it for the few queries the two engines genuinely compile differently, and
+name the difference in the test.
+
+The two engines put a different driver operation at the root — C wraps a read
+plan in ``Results``, Rust wraps a write plan in ``Commit`` — while the plan
+below it is identical. ``assert_plan_shape`` skips that root, so the expected
+tree is the part of the plan the query itself describes.
+"""
+
+from typing import Iterator, List, Optional, Tuple
+
+from falkordb.execution_plan import ExecutionPlan, Operation
+
+INDENT = "    "
+
+# top level driver operations, emitted by one engine and not the other
+ENGINE_ROOT_OPS = ("Results", "Commit")
+
+
+def iter_operations(op: Operation, depth: int = 0) -> Iterator[Tuple[int, Operation]]:
+    """Yields (depth, operation) for the whole tree, depth first."""
+
+    yield depth, op
+    for child in op.children:
+        yield from iter_operations(child, depth + 1)
+
+
+def _render(op: Operation) -> str:
+    """Renders an operation the way the reply spells it."""
+
+    return f"{op.name} | {op.args}" if op.args else op.name
+
+
+def _parse_expected(expected: str) -> List[Tuple[int, str]]:
+    """Turns an indented expected plan into (depth, text) pairs."""
+
+    lines = [line for line in expected.split("\n") if line.strip()]
+    assert lines, "expected plan is empty"
+
+    base = min(len(line) - len(line.lstrip()) for line in lines)
+    parsed = []
+    for line in lines:
+        indent = len(line) - len(line.lstrip()) - base
+        assert indent % len(INDENT) == 0, f"expected plan misindented: {line!r}"
+        parsed.append((indent // len(INDENT), line.strip()))
+    return parsed
+
+
+def assert_parsed_plan(
+    plan: ExecutionPlan,
+    min_operations: int = 1,
+    expect_args: bool = False,
+) -> None:
+    """Asserts the client parsed an execution plan reply correctly."""
+
+    root = plan.structured_plan
+    assert isinstance(root, Operation)
+
+    parsed = list(iter_operations(root))
+    lines = [line for line in plan.plan if line.strip()]
+
+    # every line of the raw reply became exactly one operation
+    assert len(parsed) == len(lines)
+    assert len(parsed) >= min_operations
+
+    for (depth, op), line in zip(parsed, lines):
+        # the operation sits as deep in the tree as its line was indented,
+        # which is what makes this a test of the parser rather than the engine
+        assert depth == (len(line) - len(line.lstrip())) // len(INDENT)
+
+        # indentation and the argument separator were stripped off the name
+        assert isinstance(op.name, str)
+        assert op.name == op.name.strip()
+        assert op.name != ""
+        assert "|" not in op.name
+        assert op.name == line.split("|")[0].strip()
+
+        assert op.args is None or isinstance(op.args, str)
+        assert isinstance(op.children, list)
+
+    if expect_args:
+        assert any(op.args for _, op in parsed)
+
+
+def assert_plan_shape(plan: ExecutionPlan, expected: str) -> None:
+    """Asserts the plan is exactly ``expected``, engine root operation aside.
+
+    ``expected`` is the operation tree, indented four spaces per level::
+
+        Project
+            Cartesian Product
+                All Node Scan | (a)
+                All Node Scan | (b)
+
+    A line may name the operation on its own, or spell out its arguments after
+    a ``|`` to pin those too. Give arguments only where both engines render
+    them identically — the traverse direction, and the order of sibling scans,
+    are two that are not guaranteed to match.
+    """
+
+    # the reply was turned into a tree faithfully in the first place
+    assert_parsed_plan(plan)
+
+    expected_ops = _parse_expected(expected)
+    actual = list(iter_operations(plan.structured_plan))
+
+    # drop the engine's root operation, unless the plan is expected to have it
+    root_name = actual[0][1].name
+    expected_root = expected_ops[0][1].split("|")[0].strip()
+    if root_name in ENGINE_ROOT_OPS and expected_root != root_name:
+        actual = [(depth - 1, op) for depth, op in actual[1:]]
+
+    rendered = [
+        (depth, _render(op) if "|" in text else op.name)
+        for (depth, op), (_, text) in zip(actual, expected_ops)
+    ]
+
+    assert len(actual) == len(expected_ops) and rendered == expected_ops, (
+        "unexpected execution plan\n\nexpected:\n%s\n\ngot:\n%s"
+        % (
+            "\n".join(INDENT * d + t for d, t in expected_ops),
+            "\n".join(INDENT * d + _render(op) for d, op in actual),
+        )
+    )
+
+
+def assert_parsed_profile(
+    plan: ExecutionPlan,
+    min_operations: int = 1,
+    expect_args: bool = False,
+    records_produced: Optional[int] = None,
+) -> None:
+    """Asserts the client parsed a profile reply, statistics included."""
+
+    assert_parsed_plan(plan, min_operations, expect_args)
+    _assert_profile_stats(plan, records_produced)
+
+
+def assert_profile_shape(
+    plan: ExecutionPlan,
+    expected: str,
+    records_produced: Optional[int] = None,
+) -> None:
+    """Asserts an exact profile plan, statistics included."""
+
+    assert_plan_shape(plan, expected)
+    _assert_profile_stats(plan, records_produced)
+
+
+def _assert_profile_stats(plan: ExecutionPlan, records_produced: Optional[int]) -> None:
+    parsed = list(iter_operations(plan.structured_plan))
+    for _, op in parsed:
+        assert op.profile_stats is not None
+        assert isinstance(op.records_produced, int)
+        assert op.records_produced >= 0
+        assert isinstance(op.execution_time, float)
+        assert op.execution_time >= 0
+
+    if records_produced is not None:
+        # how many rows the query yields is a property of the query, not of the
+        # engine, so the client must report it whichever engine answered
+        assert max(op.records_produced for _, op in parsed) == records_produced

tests/test_async_db.py

@@ -1,3 +1,4 @@
+from types import SimpleNamespace
 from unittest.mock import patch
 
 import pytest
@@ -350,3 +351,38 @@ async def test_udf_flush(async_client):
     # Verify all UDFs are removed
     udfs = await db.udf_list()
     assert udfs == []
+
+
+def test_is_cluster_filters_unknown_connection_kwargs():
+    """``Is_Cluster`` must not forward pool kwargs that ``redis.Redis`` rejects.
+
+    redis-py keeps internal state in ``connection_kwargs`` that is not part of
+    the ``Redis.__init__`` signature — 8.1.0 added ``himport_registry`` and
+    several ``maint_notifications_*`` entries — and forwarding them raises
+    ``TypeError`` before any command is sent. ``Is_Cluster`` runs during
+    connection setup, so that breaks every async query.
+
+    The real ``Redis.__init__`` is exercised here (only ``info()`` is stubbed),
+    since rejecting the kwargs is precisely what used to fail.
+    """
+    from redis.asyncio import ConnectionPool
+
+    from falkordb.asyncio.cluster import Is_Cluster
+
+    pool = ConnectionPool(host="localhost", port=6379, decode_responses=True)
+    # Pin the behaviour on redis versions that don't inject anything yet.
+    pool.connection_kwargs["himport_registry"] = object()
+
+    conn = SimpleNamespace(connection_pool=pool)
+
+    with patch(
+        "falkordb.asyncio.cluster.sync_redis.Redis.info",
+        return_value={"redis_mode": "standalone"},
+    ):
+        assert Is_Cluster(conn) is False
+
+    with patch(
+        "falkordb.asyncio.cluster.sync_redis.Redis.info",
+        return_value={"redis_mode": "cluster"},
+    ):
+        assert Is_Cluster(conn) is True

tests/test_async_explain.py

@@ -3,6 +3,8 @@
 
 from falkordb.asyncio import FalkorDB
 
+from .plan_helpers import assert_parsed_plan, assert_plan_shape
+
 
 @pytest.mark.asyncio
 async def test_explain():
@@ -17,17 +19,13 @@ async def test_explain():
 
     plan = await g.explain("UNWIND range(0, 3) AS x RETURN x")
 
-    results_op = plan.structured_plan
-    assert results_op.name == "Results"
-    assert len(results_op.children) == 1
-
-    project_op = results_op.children[0]
-    assert project_op.name == "Project"
-    assert len(project_op.children) == 1
-
-    unwind_op = project_op.children[0]
-    assert unwind_op.name == "Unwind"
-    assert len(unwind_op.children) == 0
+    assert_plan_shape(
+        plan,
+        """
+        Project
+            Unwind
+        """,
+    )
 
     # close the connection pool
     await pool.aclose()
@@ -42,26 +40,43 @@ async def test_cartesian_product_explain():
     g = db.select_graph("async_explain")
     plan = await g.explain("MATCH (a), (b) RETURN *")
 
-    results_op = plan.structured_plan
-    assert results_op.name == "Results"
-    assert len(results_op.children) == 1
-
-    project_op = results_op.children[0]
-    assert project_op.name == "Project"
-    assert len(project_op.children) == 1
-
-    cp_op = project_op.children[0]
-    assert cp_op.name == "Cartesian Product"
-    assert len(cp_op.children) == 2
+    assert_plan_shape(
+        plan,
+        """
+        Project
+            Cartesian Product
+                All Node Scan | (a)
+                All Node Scan | (b)
+        """,
+    )
 
-    scan_a_op = cp_op.children[0]
-    scan_b_op = cp_op.children[1]
+    # close the connection pool
+    await pool.aclose()
 
-    assert scan_a_op.name == "All Node Scan"
-    assert len(scan_a_op.children) == 0
 
-    assert scan_b_op.name == "All Node Scan"
-    assert len(scan_b_op.children) == 0
+@pytest.mark.asyncio
+async def test_cartesian_product_explain_three_way():
+    pool = BlockingConnectionPool(
+        max_connections=16, timeout=None, decode_responses=True
+    )
+    db = FalkorDB(connection_pool=pool)
+    g = db.select_graph("async_explain")
+    plan = await g.explain("MATCH (a), (b), (c) RETURN *")
+
+    # three operations share a nesting level here, which is what the parser
+    # used to get wrong: it attached the third scan to the second instead of
+    # to the cartesian product. The engines scan the three in whichever order
+    # they like, so the arguments are left out.
+    assert_plan_shape(
+        plan,
+        """
+        Project
+            Cartesian Product
+                All Node Scan
+                All Node Scan
+                All Node Scan
+        """,
+    )
 
     # close the connection pool
     await pool.aclose()
@@ -81,37 +96,10 @@ async def test_merge():
         pass
     plan = await g.explain("MERGE (p1:person {age: 40}) MERGE (p2:person {age: 41})")
 
-    root = plan.structured_plan
-    assert root.name == "Merge"
-    assert len(root.children) == 3
-
-    merge_op = root.children[0]
-    assert merge_op.name == "Merge"
-    assert len(merge_op.children) == 2
-
-    index_scan_op = merge_op.children[0]
-    assert index_scan_op.name == "Node By Index Scan"
-    assert len(index_scan_op.children) == 0
-
-    merge_create_op = merge_op.children[1]
-    assert merge_create_op.name == "MergeCreate"
-    assert len(merge_create_op.children) == 0
-
-    index_scan_op = root.children[1]
-    assert index_scan_op.name == "Node By Index Scan"
-    assert len(index_scan_op.children) == 1
-
-    arg_op = index_scan_op.children[0]
-    assert arg_op.name == "Argument"
-    assert len(arg_op.children) == 0
-
-    merge_create_op = root.children[2]
-    assert merge_create_op.name == "MergeCreate"
-    assert len(merge_create_op.children) == 1
-
-    arg_op = merge_create_op.children[0]
-    assert arg_op.name == "Argument"
-    assert len(arg_op.children) == 0
+    # the two engines compile MERGE differently — Rust matches through an
+    # "Include Pending" operation, C emits a "MergeCreate" — so there is no
+    # single tree to assert. Check the client parsed whatever came back.
+    assert_parsed_plan(plan, min_operations=4, expect_args=True)
 
     # close the connection pool
     await pool.aclose()