pytrilogy 0.3.138__cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

trilogy/hooks/graph_hook.py

@@ -0,0 +1,139 @@
+import sys
+from os import environ
+
+import networkx as nx
+
+from trilogy.hooks.base_hook import BaseHook
+
+if not environ.get("TCL_LIBRARY"):
+    minor = sys.version_info.minor
+    if minor == 13:
+        environ["TCL_LIBRARY"] = r"C:\Program Files\Python313\tcl\tcl8.6"
+    elif minor == 12:
+        environ["TCL_LIBRARY"] = r"C:\Program Files\Python312\tcl\tcl8.6"
+    else:
+        pass
+
+
+class GraphHook(BaseHook):
+    def __init__(self):
+        super().__init__()
+        try:
+            pass
+        except ImportError:
+            raise ImportError("GraphHook requires matplotlib and scipy to be installed")
+        # https://github.com/python/cpython/issues/125235#issuecomment-2412948604
+
+    def query_graph_built(
+        self,
+        graph: nx.DiGraph,
+        target: str | None = None,
+        highlight_nodes: list[str] | None = None,
+        remove_isolates: bool = True,
+    ):
+        from matplotlib import pyplot as plt
+
+        graph = graph.copy()
+        nodes = [*graph.nodes]
+        for node in nodes:
+            if "__preql_internal" in node:
+                graph.remove_node(node)
+
+        if remove_isolates:
+            graph.remove_nodes_from(list(nx.isolates(graph)))
+
+        color_map = []
+        highlight_nodes = highlight_nodes or []
+        for node in graph:
+            if node in highlight_nodes:
+                color_map.append("orange")
+            elif str(node).startswith("ds"):
+                color_map.append("blue")
+            else:
+                color_map.append("green")
+
+        pos = nx.spring_layout(graph)
+        kwargs = {}
+
+        if target:
+            edge_colors = []
+            descendents = nx.descendants(graph, target)
+            for edge in graph.edges():
+                if edge[0] == target:
+                    edge_colors.append("blue")
+                elif edge[1] == target:
+                    edge_colors.append("blue")
+                elif edge[1] in descendents:
+                    edge_colors.append("green")
+                else:
+                    edge_colors.append("black")
+            kwargs["edge_color"] = edge_colors
+
+        # Draw the graph without labels first
+        nx.draw(
+            graph,
+            pos=pos,
+            node_color=color_map,
+            connectionstyle="arc3, rad = 0.1",
+            with_labels=False,  # Important: don't draw labels with nx.draw
+            **kwargs
+        )
+
+        # Draw labels with manual spacing
+        self._draw_labels_with_manual_spacing(graph, pos)
+
+        plt.show()
+
+    def _draw_labels_with_manual_spacing(self, graph, pos):
+        import numpy as np
+
+        pos_labels = {}
+        node_positions = list(pos.values())
+
+        # Calculate average distance between nodes to determine spacing
+        if len(node_positions) > 1:
+            distances = []
+            for i, (x1, y1) in enumerate(node_positions):
+                for j, (x2, y2) in enumerate(node_positions[i + 1 :], i + 1):
+                    dist = np.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)
+                    distances.append(dist)
+
+            avg_distance = np.mean(distances)
+            min_spacing = max(
+                0.1, avg_distance * 0.3
+            )  # Minimum spacing as fraction of average distance
+        else:
+            min_spacing = 0.1
+
+        # Simple spacing algorithm - offset labels that are too close
+        for i, node in enumerate(graph.nodes()):
+            x, y = pos[node]
+
+            # Check for nearby labels and adjust position
+            adjusted_x, adjusted_y = x, y
+            for j, other_node in enumerate(
+                list(graph.nodes())[:i]
+            ):  # Only check previous nodes
+                other_x, other_y = pos_labels.get(other_node, pos[other_node])
+                distance = np.sqrt(
+                    (adjusted_x - other_x) ** 2 + (adjusted_y - other_y) ** 2
+                )
+
+                if distance < min_spacing:
+                    # Calculate offset direction
+                    if distance > 0:
+                        offset_x = (adjusted_x - other_x) / distance * min_spacing
+                        offset_y = (adjusted_y - other_y) / distance * min_spacing
+                    else:
+                        # If nodes are at exact same position, use random offset
+                        angle = np.random.random() * 2 * np.pi
+                        offset_x = np.cos(angle) * min_spacing
+                        offset_y = np.sin(angle) * min_spacing
+
+                    adjusted_x = other_x + offset_x
+                    adjusted_y = other_y + offset_y
+
+            pos_labels[node] = (adjusted_x, adjusted_y)
+
+        # Draw the labels at adjusted positions
+        nx.draw_networkx_labels(graph, pos=pos_labels, font_size=10)

trilogy/hooks/query_debugger.py

@@ -0,0 +1,166 @@
+from enum import Enum
+from logging import DEBUG, StreamHandler
+from typing import Union
+from uuid import uuid4
+
+from trilogy.constants import logger
+from trilogy.core.models.build import BuildDatasource
+from trilogy.core.models.execute import (
+    CTE,
+    QueryDatasource,
+    UnionCTE,
+)
+from trilogy.core.processing.nodes import StrategyNode
+from trilogy.core.statements.author import SelectStatement
+from trilogy.hooks.base_hook import BaseHook
+
+
+class PrintMode(Enum):
+    OFF = False
+    BASIC = True
+    FULL = 3
+
+
+class DebuggingHook(BaseHook):
+    def __init__(
+        self,
+        level=DEBUG,
+        max_depth: int | None = None,
+        process_ctes: PrintMode | bool = True,
+        process_nodes: PrintMode | bool = True,
+        process_datasources: PrintMode | bool = True,
+        process_other: bool = True,
+    ):
+        if not any([isinstance(x, StreamHandler) for x in logger.handlers]):
+            logger.addHandler(StreamHandler())
+        logger.setLevel(level)
+
+        self.max_depth = max_depth
+        self.process_ctes = PrintMode(process_ctes)
+        self.process_nodes = PrintMode(process_nodes)
+        self.process_datasources = PrintMode(process_datasources)
+        self.process_other = PrintMode(process_other)
+        self.messages: list[str] = []
+        self.uuid = uuid4()
+        from trilogy.dialect.bigquery import BigqueryDialect
+
+        self.renderer = BigqueryDialect()
+
+    def print(self, *args):
+        merged = " ".join([str(x) for x in args])
+        self.messages.append(merged)
+
+    def write(self):
+        with open(f"debug_{self.uuid}.log", "w") as f:
+            f.write("\n".join(self.messages))
+
+    def process_select_info(self, select: SelectStatement):
+        if self.process_datasources != PrintMode.OFF:
+            self.print(f"grain: {str(select.grain)}")
+
+    def process_root_datasource(self, datasource: QueryDatasource):
+        if self.process_datasources != PrintMode.OFF:
+            printed = self.print_recursive_resolved(
+                datasource, self.process_datasources
+            )
+            for row in printed:
+                self.print("".join([str(v) for v in row]))
+
+    def process_root_cte(self, cte: CTE | UnionCTE):
+        if self.process_ctes != PrintMode.OFF:
+            self.print_recursive_ctes(cte, max_depth=self.max_depth)
+
+    def process_root_strategy_node(self, node: StrategyNode):
+        if self.process_nodes != PrintMode.OFF:
+            printed = self.print_recursive_nodes(node, mode=self.process_nodes)
+            for row in printed:
+                # logger.info("".join([str(v) for v in row]))
+                self.print("".join([str(v) for v in row]))
+
+    def print_recursive_resolved(
+        self,
+        input: Union[QueryDatasource, BuildDatasource],
+        mode: PrintMode,
+        depth: int = 0,
+    ):
+        extra = []
+        if isinstance(input, QueryDatasource):
+            if input.joins:
+                extra.append("join")
+            if input.condition:
+                extra.append("filter")
+        if input.group_required:
+            extra.append("group")
+        output = [c.address for c in input.output_concepts[:3]]
+        if len(input.output_concepts) > 3:
+            output.append("...")
+        display = [
+            (
+                "  " * depth,
+                input.__class__.__name__,
+                "<",
+                ",".join(extra),
+                ">",
+                # [c.address for c in input.input_concepts],
+                "->",
+                output,
+            )
+        ]
+        if isinstance(input, QueryDatasource):
+            for child in input.datasources:
+                display += self.print_recursive_resolved(
+                    child, mode=mode, depth=depth + 1
+                )
+        return display
+
+    def print_recursive_ctes(
+        self, input: CTE | UnionCTE, depth: int = 0, max_depth: int | None = None
+    ):
+        if max_depth and depth > max_depth:
+            return
+        select_statement = [c.address for c in input.output_columns]
+        self.print(
+            "  " * depth, input.name, "->", input.group_to_grain, "->", select_statement
+        )
+        sql = self.renderer.render_cte(input).statement
+        for line in sql.split("\n"):
+            logger.debug("  " * (depth) + line)
+        if isinstance(input, CTE):
+            for child in input.parent_ctes:
+                self.print_recursive_ctes(child, depth + 1)
+        elif isinstance(input, UnionCTE):
+            for child in input.parent_ctes:
+                for parent in child.parent_ctes:
+                    self.print_recursive_ctes(parent, depth + 1)
+
+    def print_recursive_nodes(
+        self, input: StrategyNode, mode: PrintMode = PrintMode.BASIC, depth: int = 0
+    ):
+        resolved = input.resolve()
+        if mode == PrintMode.FULL:
+            display = [
+                [
+                    "  " * depth,
+                    input,
+                    "->",
+                    resolved.grain,
+                    "->",
+                    [c.address for c in resolved.output_concepts],
+                ]
+            ]
+        elif mode == PrintMode.BASIC:
+            display = [
+                [
+                    "  " * depth,
+                    input,
+                    "->",
+                    resolved.grain,
+                ]
+            ]
+        for child in input.parents:
+            display += self.print_recursive_nodes(
+                child,
+                mode=mode,
+                depth=depth + 1,
+            )
+        return display

trilogy/parser.py

@@ -0,0 +1,10 @@
+from typing import Optional
+
+from trilogy.core.models.environment import Environment
+from trilogy.parsing.parse_engine import parse_text
+
+
+def parse(
+    input: str, environment: Optional[Environment] = None
+) -> tuple[Environment, list]:
+    return parse_text(input, environment=environment)

trilogy/parsing/README.md

@@ -0,0 +1,21 @@
+# Grammar Overview
+
+```
+ !start: ( block | show_statement )*
+    block: statement _TERMINATOR PARSE_COMMENT?
+    ?statement: concept
+    | datasource
+    | function
+    | multi_select_statement
+    | select_statement
+    | persist_statement
+    | rowset_derivation_statement
+    | import_statement
+    | copy_statement
+    | merge_statement
+    | rawsql_statement
+    
+    _TERMINATOR:  ";"i /\s*/
+```
+
+Full grammar in trilogy.lark.