pytrilogy 0.0.1.102__py3-none-any.whl

trilogy/dialect/sql_server.py

@@ -0,0 +1,89 @@
+from typing import Mapping, Callable, Any
+
+from jinja2 import Template
+from trilogy.utility import string_to_hash
+
+
+from trilogy.core.enums import FunctionType, WindowType
+from trilogy.core.models import (
+    ProcessedQuery,
+    ProcessedQueryPersist,
+    ProcessedShowStatement,
+)
+from trilogy.dialect.base import BaseDialect
+
+WINDOW_FUNCTION_MAP: Mapping[WindowType, Callable[[Any, Any, Any], str]] = {}
+
+FUNCTION_MAP = {
+    FunctionType.COUNT: lambda args: f"count({args[0]})",
+    FunctionType.SUM: lambda args: f"sum({args[0]})",
+    FunctionType.AVG: lambda args: f"avg({args[0]})",
+    FunctionType.LENGTH: lambda args: f"length({args[0]})",
+    FunctionType.LIKE: lambda args: (
+        f" CASE WHEN {args[0]} like {args[1]} THEN True ELSE False END"
+    ),
+    FunctionType.CONCAT: lambda args: (
+        f"CONCAT({','.join([f''' '{a}' ''' for a in args])})"
+    ),
+}
+
+# if an aggregate function is called on a source that is at the same grain as the aggregate
+# we may return a static value
+FUNCTION_GRAIN_MATCH_MAP = {
+    **FUNCTION_MAP,
+    FunctionType.COUNT: lambda args: "1",
+    FunctionType.SUM: lambda args: f"{args[0]}",
+    FunctionType.AVG: lambda args: f"{args[0]}",
+}
+
+TSQL_TEMPLATE = Template(
+    """{%- if ctes %}
+WITH {% for cte in ctes %}
+{{cte.name}} as ({{cte.statement}}){% if not loop.last %},{% endif %}{% endfor %}{% endif %}
+SELECT
+{%- if limit is not none %}
+TOP {{ limit }}{% endif %}
+{%- for select in select_columns %}
+    {{ select }}{% if not loop.last %},{% endif %}{% endfor %}
+{% if base %}FROM
+    {{ base }}{% endif %}{% if joins %}
+{% for join in joins %}
+{{ join }}
+{% endfor %}{% endif %}
+{% if where %}WHERE
+    {{ where }}
+{% endif %}
+{%- if group_by %}
+GROUP BY {% for group in group_by %}
+    {{group}}{% if not loop.last %},{% endif %}
+{% endfor %}{% endif %}
+{%- if order_by %}
+ORDER BY {% for order in order_by %}
+    {{ order }}{% if not loop.last %},{% endif %}
+{% endfor %}{% endif %}
+"""
+)
+
+MAX_IDENTIFIER_LENGTH = 128
+
+
+class SqlServerDialect(BaseDialect):
+    WINDOW_FUNCTION_MAP = {**BaseDialect.WINDOW_FUNCTION_MAP, **WINDOW_FUNCTION_MAP}
+    FUNCTION_MAP = {**BaseDialect.FUNCTION_MAP, **FUNCTION_MAP}
+    FUNCTION_GRAIN_MATCH_MAP = {
+        **BaseDialect.FUNCTION_GRAIN_MATCH_MAP,
+        **FUNCTION_GRAIN_MATCH_MAP,
+    }
+    QUOTE_CHARACTER = '"'
+    SQL_TEMPLATE = TSQL_TEMPLATE
+
+    def compile_statement(
+        self, query: ProcessedQuery | ProcessedQueryPersist | ProcessedShowStatement
+    ) -> str:
+        base = super().compile_statement(query)
+        if isinstance(base, (ProcessedQuery, ProcessedQueryPersist)):
+            for cte in query.ctes:
+                if len(cte.name) > MAX_IDENTIFIER_LENGTH:
+                    new_name = f"rhash_{string_to_hash(cte.name)}"
+                    base = base.replace(cte.name, new_name)
+        return base

trilogy/engine.py

@@ -0,0 +1,48 @@
+from sqlalchemy.engine import Engine, Connection, CursorResult
+from typing import Protocol
+
+
+class EngineResult(Protocol):
+    pass
+
+    def fetchall(self) -> list[tuple]:
+        pass
+
+
+class EngineConnection(Protocol):
+    pass
+
+    def execute(self, statement: str) -> EngineResult:
+        pass
+
+
+class ExecutionEngine(Protocol):
+    pass
+
+    def connect(self) -> EngineConnection:
+        pass
+
+
+### Begin default SQLAlchemy implementation
+class SqlAlchemyResult(EngineResult):
+    def __init__(self, result: CursorResult):
+        self.result = result
+
+    def fetchall(self):
+        return self.result.fetchall()
+
+
+class SqlAlchemyConnection(EngineConnection):
+    def __init__(self, connection: Connection):
+        self.connection = connection
+
+    def execute(self, statement: str) -> SqlAlchemyResult:
+        return SqlAlchemyResult(self.connection.execute(statement))
+
+
+class SqlAlchemyEngine(ExecutionEngine):
+    def __init__(self, engine: Engine):
+        self.engine = engine
+
+    def connect(self) -> SqlAlchemyConnection:
+        return SqlAlchemyConnection(self.engine.connect())

trilogy/executor.py

@@ -0,0 +1,242 @@
+from typing import List, Optional, Any
+from functools import singledispatchmethod
+from sqlalchemy import text
+from sqlalchemy.engine import Engine, CursorResult
+
+from trilogy.constants import logger
+from trilogy.core.models import (
+    Environment,
+    ProcessedQuery,
+    ProcessedShowStatement,
+    ProcessedQueryPersist,
+    MultiSelectStatement,
+    SelectStatement,
+    PersistStatement,
+    ShowStatement,
+    Concept,
+)
+from trilogy.dialect.base import BaseDialect
+from trilogy.dialect.enums import Dialects
+from trilogy.parser import parse_text
+from trilogy.hooks.base_hook import BaseHook
+
+from dataclasses import dataclass
+
+
+@dataclass
+class MockResult:
+    values: list[Any]
+    columns: list[str]
+
+    def fetchall(self):
+        return self.values
+
+    def keys(self):
+        return self.columns
+
+
+def generate_result_set(columns: List[Concept], output_data: list[Any]) -> MockResult:
+    names = [x.address.replace(".", "_") for x in columns]
+    return MockResult(
+        values=[dict(zip(names, [row])) for row in output_data], columns=names
+    )
+
+
+class Executor(object):
+    def __init__(
+        self,
+        dialect: Dialects,
+        engine: Engine,
+        environment: Optional[Environment] = None,
+        hooks: List[BaseHook] | None = None,
+    ):
+        self.dialect: Dialects = dialect
+        self.engine = engine
+        self.environment = environment or Environment()
+        self.generator: BaseDialect
+        self.logger = logger
+        self.hooks = hooks
+        if self.dialect == Dialects.BIGQUERY:
+            from trilogy.dialect.bigquery import BigqueryDialect
+
+            self.generator = BigqueryDialect()
+        elif self.dialect == Dialects.SQL_SERVER:
+            from trilogy.dialect.sql_server import SqlServerDialect
+
+            self.generator = SqlServerDialect()
+        elif self.dialect == Dialects.DUCK_DB:
+            from trilogy.dialect.duckdb import DuckDBDialect
+
+            self.generator = DuckDBDialect()
+        elif self.dialect == Dialects.PRESTO:
+            from trilogy.dialect.presto import PrestoDialect
+
+            self.generator = PrestoDialect()
+        elif self.dialect == Dialects.TRINO:
+            from trilogy.dialect.presto import TrinoDialect
+
+            self.generator = TrinoDialect()
+        elif self.dialect == Dialects.POSTGRES:
+            from trilogy.dialect.postgres import PostgresDialect
+
+            self.generator = PostgresDialect()
+        elif self.dialect == Dialects.SNOWFLAKE:
+
+            from trilogy.dialect.snowflake import SnowflakeDialect
+
+            self.generator = SnowflakeDialect()
+        else:
+            raise ValueError(f"Unsupported dialect {self.dialect}")
+        self.connection = self.engine.connect()
+
+    def execute_statement(self, statement) -> Optional[CursorResult]:
+        if not isinstance(statement, (ProcessedQuery, ProcessedQueryPersist)):
+            return None
+        return self.execute_query(statement)
+
+    @singledispatchmethod
+    def execute_query(self, query) -> CursorResult:
+        raise NotImplementedError("Cannot execute type {}".format(type(query)))
+
+    @execute_query.register
+    def _(self, query: SelectStatement | PersistStatement) -> CursorResult:
+        sql = self.generator.generate_queries(
+            self.environment, [query], hooks=self.hooks
+        )
+        return self.execute_query(sql[0])
+
+    @execute_query.register
+    def _(self, query: ProcessedShowStatement) -> CursorResult:
+        return generate_result_set(
+            query.output_columns,
+            [
+                self.generator.compile_statement(x)
+                for x in query.output_values
+                if isinstance(x, ProcessedQuery)
+            ],
+        )
+
+    @execute_query.register
+    def _(self, query: ProcessedQuery | ProcessedQueryPersist) -> CursorResult:
+        sql = self.generator.compile_statement(query)
+        # connection = self.engine.connect()
+        output = self.connection.execute(text(sql))
+        if isinstance(query, ProcessedQueryPersist):
+            self.environment.add_datasource(query.datasource)
+        return output
+
+    @singledispatchmethod
+    def generate_sql(self, command: ProcessedQuery | str) -> list[str]:
+        raise NotImplementedError(
+            "Cannot generate sql for type {}".format(type(command))
+        )
+
+    @generate_sql.register  # type: ignore
+    def _(self, command: ProcessedQuery) -> List[str]:
+        output = []
+        compiled_sql = self.generator.compile_statement(command)
+        output.append(compiled_sql)
+        return output
+
+    @generate_sql.register  # type: ignore
+    def _(self, command: MultiSelectStatement) -> List[str]:
+        output = []
+        sql = self.generator.generate_queries(
+            self.environment, [command], hooks=self.hooks
+        )
+        for statement in sql:
+            compiled_sql = self.generator.compile_statement(statement)
+            output.append(compiled_sql)
+
+        output.append(compiled_sql)
+        return output
+
+    @generate_sql.register  # type: ignore
+    def _(self, command: SelectStatement) -> List[str]:
+        output = []
+        sql = self.generator.generate_queries(
+            self.environment, [command], hooks=self.hooks
+        )
+        for statement in sql:
+            compiled_sql = self.generator.compile_statement(statement)
+            output.append(compiled_sql)
+        return output
+
+    @generate_sql.register  # type: ignore
+    def _(self, command: str) -> List[str]:
+        """generate SQL for execution"""
+        _, parsed = parse_text(command, self.environment)
+        generatable = [
+            x for x in parsed if isinstance(x, (SelectStatement, PersistStatement))
+        ]
+        sql = self.generator.generate_queries(
+            self.environment, generatable, hooks=self.hooks
+        )
+        output = []
+        for statement in sql:
+            if isinstance(statement, ProcessedShowStatement):
+                continue
+            compiled_sql = self.generator.compile_statement(statement)
+            output.append(compiled_sql)
+        return output
+
+    def parse_text(
+        self, command: str, persist: bool = False
+    ) -> List[ProcessedQuery | ProcessedQueryPersist | ProcessedShowStatement]:
+        """Process a preql text command"""
+        _, parsed = parse_text(command, self.environment)
+        generatable = [
+            x
+            for x in parsed
+            if isinstance(
+                x,
+                (
+                    SelectStatement,
+                    PersistStatement,
+                    MultiSelectStatement,
+                    ShowStatement,
+                ),
+            )
+        ]
+        sql = []
+        while generatable:
+            t = generatable.pop(0)
+            x = self.generator.generate_queries(
+                self.environment, [t], hooks=self.hooks
+            )[0]
+            if persist and isinstance(x, ProcessedQueryPersist):
+                self.environment.add_datasource(x.datasource)
+            sql.append(x)
+        return sql
+
+    def execute_raw_sql(self, command: str) -> CursorResult:
+        """Run a command against the raw underlying
+        execution engine"""
+        return self.connection.execute(text(command))
+
+    def execute_text(self, command: str) -> List[CursorResult]:
+        """Run a preql text command"""
+        sql = self.parse_text(command)
+        output = []
+        # connection = self.engine.connect()
+        for statement in sql:
+            if isinstance(statement, ProcessedShowStatement):
+                output.append(
+                    generate_result_set(
+                        statement.output_columns,
+                        [
+                            self.generator.compile_statement(x)
+                            for x in statement.output_values
+                            if isinstance(x, ProcessedQuery)
+                        ],
+                    )
+                )
+                continue
+            compiled_sql = self.generator.compile_statement(statement)
+            logger.debug(compiled_sql)
+
+            output.append(self.connection.execute(text(compiled_sql)))
+            # generalize post-run success hooks
+            if isinstance(statement, ProcessedQueryPersist):
+                self.environment.add_datasource(statement.datasource)
+        return output

trilogy/hooks/base_hook.py

@@ -0,0 +1,37 @@
+from trilogy.core.models import (
+    QueryDatasource,
+    CTE,
+    SelectStatement,
+    PersistStatement,
+    MultiSelectStatement,
+    RowsetDerivationStatement,
+)
+from trilogy.core.processing.nodes import StrategyNode
+
+
+class BaseHook:
+    pass
+
+    def process_multiselect_info(self, select: MultiSelectStatement):
+        print("Multiselect with components:")
+        for x in select.selects:
+            self.process_select_info(x)
+
+    def process_select_info(self, select: SelectStatement):
+        print(f"Select statement grain: {str(select.grain)}")
+
+    def process_persist_info(self, persist: PersistStatement):
+        print(f"Persist statement persisting to {persist.address}")
+        self.process_select_info(persist.select)
+
+    def process_rowset_info(self, rowset: RowsetDerivationStatement):
+        print(f"Rowset statement with grain {str(rowset.select.grain)}")
+
+    def process_root_datasource(self, datasource: QueryDatasource):
+        pass
+
+    def process_root_cte(self, cte: CTE):
+        pass
+
+    def process_root_strategy_node(self, node: StrategyNode):
+        pass

trilogy/hooks/graph_hook.py

@@ -0,0 +1,24 @@
+from trilogy.hooks.base_hook import BaseHook
+from networkx import DiGraph
+
+
+class GraphHook(BaseHook):
+    def __init__(self):
+        super().__init__()
+        try:
+            pass
+        except ImportError:
+            raise ImportError("GraphHook requires matplotlib and scipy to be installed")
+
+    def query_graph_built(self, graph: DiGraph):
+        from networkx import draw_kamada_kawai
+        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)
+        draw_kamada_kawai(graph, with_labels=True, connectionstyle="arc3, rad = 0.1")
+        # draw_spring(graph, with_labels=True, connectionstyle='arc3, rad = 0.1')
+        plt.show()

trilogy/hooks/query_debugger.py

@@ -0,0 +1,133 @@
+from typing import Union
+from trilogy.core.models import QueryDatasource, CTE, Datasource, SelectStatement
+
+from trilogy.hooks.base_hook import BaseHook
+from trilogy.constants import logger
+from logging import StreamHandler, DEBUG
+from trilogy.core.processing.nodes import StrategyNode
+
+from trilogy.dialect.bigquery import BigqueryDialect
+
+from enum import Enum
+
+
+class PrintMode(Enum):
+    OFF = False
+    BASIC = True
+    FULL = 3
+
+
+renderer = BigqueryDialect()
+
+
+def print_recursive_resolved(
+    input: Union[QueryDatasource, Datasource], 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")
+    display = [
+        (
+            "  " * depth,
+            input.__class__.__name__,
+            "<",
+            ",".join(extra),
+            ">",
+            # [c.address for c in input.input_concepts],
+            "->",
+            [c.address for c in input.output_concepts],
+        )
+    ]
+    if isinstance(input, QueryDatasource):
+        for child in input.datasources:
+            display += print_recursive_resolved(child, mode=mode, depth=depth + 1)
+    return display
+
+
+def print_recursive_nodes(
+    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 += print_recursive_nodes(child, mode=mode, depth=depth + 1)
+    return display
+
+
+def print_recursive_ctes(input: CTE, 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]
+    print("  " * depth, input.name, "->", input.group_to_grain, "->", select_statement)
+    sql = renderer.render_cte(input).statement
+    for line in sql.split("\n"):
+        logger.debug("  " * (depth) + line)
+    print("-----")
+    if isinstance(input, CTE):
+        for child in input.parent_ctes:
+            print_recursive_ctes(child, depth + 1)
+
+
+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)
+
+    def process_select_info(self, select: SelectStatement):
+        if self.process_datasources != PrintMode.OFF:
+            print(f"grain: {str(select.grain)}")
+
+    def process_root_datasource(self, datasource: QueryDatasource):
+        if self.process_datasources != PrintMode.OFF:
+            printed = print_recursive_resolved(datasource, self.process_datasources)
+            for row in printed:
+                print("".join([str(v) for v in row]))
+
+    def process_root_cte(self, cte: CTE):
+        if self.process_ctes != PrintMode.OFF:
+            print_recursive_ctes(cte, max_depth=self.max_depth)
+
+    def process_root_strategy_node(self, node: StrategyNode):
+        if self.process_nodes != PrintMode.OFF:
+            printed = print_recursive_nodes(node, mode=self.process_nodes)
+            for row in printed:
+                print("".join([str(v) for v in row]))

trilogy/parser.py

@@ -0,0 +1,10 @@
+from typing import Optional
+
+from trilogy.core.models 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)