✨ draw post execution plans as trees

[joocer]

To help with understanding query execution,, provide the ability to draw an executed query plan, ChatGPT has suggested this code will help draw the structure

from collections import defaultdict
from typing import List, Tuple, Dict

class QueryPlanTree:
    def __init__(self, edges: List[Tuple[str, str, str]]):
        """
        Constructs a query plan tree from a list of edges.

        Parameters:
            edges: List of tuples (source, relationship, target)
        """
        self.edges = edges
        self.children = defaultdict(list)
        self.parents = {}
        
        # Build adjacency list from edges
        for src, rel, tgt in edges:
            self.children[src].append((rel, tgt))
            self.parents[tgt] = src
        
        # Find root (a node that has no incoming edges)
        self.root = None
        for src, _, tgt in edges:
            if src not in self.parents:
                self.root = src
                break

    def _build_tree(self, node: str, prefix="") -> List[str]:
        """
        Recursively builds the ASCII tree representation.
        """
        lines = [
            f"{prefix}┌────────────┐",
            f"{prefix}│ {node:^10} │",
            f"{prefix}└────────────┘"
        ]
        
        children = self.children.get(node, [])
        if not children:
            return lines

        # Prepare child branches
        subtrees = [self._build_tree(child, prefix + "        ") for _, child in children]

        # Combine the current node with its children
        result = []
        result.extend(lines)
        result.append(prefix + "      ▲")
        result.append(prefix + " ┌───────┴────────┐")

        # Format children horizontally
        for i, subtree in enumerate(subtrees):
            if i > 0:
                result.append(prefix + " │                │")
            result.extend(subtree)

        return result

    def draw(self) -> str:
        """
        Returns the ASCII tree as a string.
        """
        if not self.root:
            return "No root node found."
        return "\n".join(self._build_tree(self.root))


# Example edges representing a query plan
edges = [
    ("Join", "parent", "Scan A"),
    ("Join", "parent", "Join_2"),
    ("Join_2", "parent", "Scan B"),
    ("Join_2", "parent", "Scan C")
]

# Create and render the tree
query_tree = QueryPlanTree(edges)
print(query_tree.draw())