pytrilogy 0.0.2.46__py3-none-any.whl → 0.0.2.48__py3-none-any.whl

trilogy/executor.py

@@ -1,39 +1,40 @@
-from typing import List, Optional, Any, Generator, Protocol
+from dataclasses import dataclass
 from functools import singledispatchmethod
+from pathlib import Path
+from typing import Any, Generator, List, Optional, Protocol
+
 from sqlalchemy import text
-from sqlalchemy.engine import Engine, CursorResult
+from sqlalchemy.engine import CursorResult, Engine
 
 from trilogy.constants import logger
+from trilogy.core.enums import Granularity, IOType
 from trilogy.core.models import (
+    Concept,
+    ConceptDeclarationStatement,
+    CopyStatement,
+    Datasource,
     Environment,
+    Function,
+    FunctionType,
+    ImportStatement,
+    ListWrapper,
+    MapWrapper,
+    MergeStatementV2,
+    MultiSelectStatement,
+    PersistStatement,
+    ProcessedCopyStatement,
     ProcessedQuery,
-    ProcessedShowStatement,
     ProcessedQueryPersist,
     ProcessedRawSQLStatement,
-    ProcessedCopyStatement,
+    ProcessedShowStatement,
     RawSQLStatement,
-    MultiSelectStatement,
     SelectStatement,
-    PersistStatement,
     ShowStatement,
-    Concept,
-    ConceptDeclarationStatement,
-    Datasource,
-    CopyStatement,
-    ImportStatement,
-    MergeStatementV2,
-    Function,
-    FunctionType,
-    MapWrapper,
-    ListWrapper,
 )
 from trilogy.dialect.base import BaseDialect
 from trilogy.dialect.enums import Dialects
-from trilogy.core.enums import IOType, Granularity
-from trilogy.parser import parse_text
 from trilogy.hooks.base_hook import BaseHook
-from pathlib import Path
-from dataclasses import dataclass
+from trilogy.parser import parse_text
 
 
 class ResultProtocol(Protocol):
@@ -103,7 +104,6 @@ class Executor(object):
 
             self.generator = PostgresDialect()
         elif self.dialect == Dialects.SNOWFLAKE:
-
             from trilogy.dialect.snowflake import SnowflakeDialect
 
             self.generator = SnowflakeDialect()
@@ -146,7 +146,6 @@ class Executor(object):
 
     @execute_query.register
     def _(self, query: Datasource) -> CursorResult:
-
         return MockResult(
             [
                 {
@@ -232,22 +231,23 @@ class Executor(object):
     @execute_query.register
     def _(self, query: ProcessedQuery) -> CursorResult:
         sql = self.generator.compile_statement(query)
-        output = self.execute_raw_sql(sql)
+        output = self.execute_raw_sql(sql, local_concepts=query.local_concepts)
         return output
 
     @execute_query.register
     def _(self, query: ProcessedQueryPersist) -> CursorResult:
-
         sql = self.generator.compile_statement(query)
 
-        output = self.execute_raw_sql(sql)
+        output = self.execute_raw_sql(sql, local_concepts=query.local_concepts)
         self.environment.add_datasource(query.datasource)
         return output
 
     @execute_query.register
     def _(self, query: ProcessedCopyStatement) -> CursorResult:
         sql = self.generator.compile_statement(query)
-        output: CursorResult = self.execute_raw_sql(sql)
+        output: CursorResult = self.execute_raw_sql(
+            sql, local_concepts=query.local_concepts
+        )
         if query.target_type == IOType.CSV:
             import csv
 
@@ -355,7 +355,6 @@ class Executor(object):
         | ProcessedRawSQLStatement
         | ProcessedCopyStatement
     ]:
-
         return list(self.parse_text_generator(command, persist=persist))
 
     def parse_text_generator(self, command: str, persist: bool = False) -> Generator[
@@ -395,12 +394,20 @@ class Executor(object):
             if persist and isinstance(x, ProcessedQueryPersist):
                 self.environment.add_datasource(x.datasource)
 
-    def _hydrate_param(self, param: str) -> Any:
+    def _hydrate_param(
+        self, param: str, local_concepts: dict[str, Concept] | None = None
+    ) -> Any:
         matched = [
             v
             for v in self.environment.concepts.values()
             if v.safe_address == param or v.address == param
         ]
+        if local_concepts and not matched:
+            matched = [
+                v
+                for v in local_concepts.values()
+                if v.safe_address == param or v.address == param
+            ]
         if not matched:
             raise SyntaxError(f"No concept found for parameter {param}")
 
@@ -424,7 +431,10 @@ class Executor(object):
         return results[0]
 
     def execute_raw_sql(
-        self, command: str, variables: dict | None = None
+        self,
+        command: str,
+        variables: dict | None = None,
+        local_concepts: dict[str, Concept] | None = None,
     ) -> CursorResult:
         """Run a command against the raw underlying
         execution engine"""
@@ -435,7 +445,10 @@ class Executor(object):
         else:
             params = q.compile().params
             if params:
-                final_params = {x: self._hydrate_param(x) for x in params}
+                final_params = {
+                    x: self._hydrate_param(x, local_concepts=local_concepts)
+                    for x in params
+                }
 
         if final_params:
             return self.connection.execute(text(command), final_params)

trilogy/hooks/base_hook.py

@@ -1,10 +1,11 @@
 from trilogy.core.models import (
-    QueryDatasource,
     CTE,
-    SelectStatement,
-    PersistStatement,
     MultiSelectStatement,
+    PersistStatement,
+    QueryDatasource,
     RowsetDerivationStatement,
+    SelectStatement,
+    UnionCTE,
 )
 from trilogy.core.processing.nodes import StrategyNode
 
@@ -30,7 +31,7 @@ class BaseHook:
     def process_root_datasource(self, datasource: QueryDatasource):
         pass
 
-    def process_root_cte(self, cte: CTE):
+    def process_root_cte(self, cte: CTE | UnionCTE):
         pass
 
     def process_root_strategy_node(self, node: StrategyNode):

trilogy/hooks/graph_hook.py

@@ -1,6 +1,7 @@
-from trilogy.hooks.base_hook import BaseHook
 import networkx as nx
 
+from trilogy.hooks.base_hook import BaseHook
+
 
 class GraphHook(BaseHook):
     def __init__(self):

trilogy/hooks/query_debugger.py

@@ -1,14 +1,18 @@
+from enum import Enum
+from logging import DEBUG, StreamHandler
 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.models import (
+    CTE,
+    Datasource,
+    QueryDatasource,
+    SelectStatement,
+    UnionCTE,
+)
 from trilogy.core.processing.nodes import StrategyNode
-
 from trilogy.dialect.bigquery import BigqueryDialect
-
-from enum import Enum
+from trilogy.hooks.base_hook import BaseHook
 
 
 class PrintMode(Enum):
@@ -85,7 +89,9 @@ def print_recursive_nodes(
     return display
 
 
-def print_recursive_ctes(input: CTE, depth: int = 0, max_depth: int | None = None):
+def print_recursive_ctes(
+    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]
@@ -96,6 +102,10 @@ def print_recursive_ctes(input: CTE, depth: int = 0, max_depth: int | None = Non
     if isinstance(input, CTE):
         for child in input.parent_ctes:
             print_recursive_ctes(child, depth + 1)
+    elif isinstance(input, UnionCTE):
+        for child in input.parent_ctes:
+            for parent in child.parent_ctes:
+                print_recursive_ctes(parent, depth + 1)
 
 
 class DebuggingHook(BaseHook):
@@ -128,7 +138,7 @@ class DebuggingHook(BaseHook):
             for row in printed:
                 print("".join([str(v) for v in row]))
 
-    def process_root_cte(self, cte: CTE):
+    def process_root_cte(self, cte: CTE | UnionCTE):
         if self.process_ctes != PrintMode.OFF:
             print_recursive_ctes(cte, max_depth=self.max_depth)
 

trilogy/parsing/common.py

@@ -1,32 +1,28 @@
+from typing import List, Tuple
+
+from trilogy.constants import (
+    VIRTUAL_CONCEPT_PREFIX,
+)
+from trilogy.core.enums import FunctionType, Modifier, PurposeLineage, WindowType
+from trilogy.core.functions import arg_to_datatype, function_args_to_output_purpose
 from trilogy.core.models import (
     AggregateWrapper,
     Concept,
+    DataType,
+    Environment,
+    FilterItem,
     Function,
+    FunctionClass,
     Grain,
-    Purpose,
-    Metadata,
-    FilterItem,
     ListWrapper,
     MapWrapper,
-    WindowItem,
     Meta,
+    Metadata,
     Parenthetical,
-    FunctionClass,
-    Environment,
-    DataType,
-)
-from typing import List, Tuple
-from trilogy.core.functions import (
-    function_args_to_output_purpose,
-    FunctionType,
-    arg_to_datatype,
-)
-from trilogy.utility import unique, string_to_hash
-from trilogy.core.enums import PurposeLineage
-from trilogy.constants import (
-    VIRTUAL_CONCEPT_PREFIX,
+    Purpose,
+    WindowItem,
 )
-from trilogy.core.enums import Modifier, WindowType
+from trilogy.utility import string_to_hash, unique
 
 
 def get_upstream_modifiers(keys: List[Concept]) -> list[Modifier]:
@@ -304,7 +300,6 @@ def arbitrary_to_concept(
     metadata: Metadata | None = None,
     purpose: Purpose | None = None,
 ) -> Concept:
-
     if isinstance(parent, AggregateWrapper):
         if not name:
             name = f"{VIRTUAL_CONCEPT_PREFIX}_agg_{parent.function.operator.value}_{string_to_hash(str(parent))}"