pytrilogy 0.0.2.47__py3-none-any.whl → 0.0.2.49__py3-none-any.whl

trilogy/dialect/common.py

@@ -1,13 +1,14 @@
+from typing import Callable
+
+from trilogy.core.enums import Modifier, UnnestMode
 from trilogy.core.models import (
-    Join,
-    InstantiatedUnnestJoin,
     CTE,
     Concept,
     Function,
+    InstantiatedUnnestJoin,
+    Join,
     RawColumnExpr,
 )
-from trilogy.core.enums import UnnestMode, Modifier
-from typing import Callable
 
 
 def null_wrapper(lval: str, rval: str, modifiers: list[Modifier]) -> str:

trilogy/dialect/config.py

@@ -1,5 +1,4 @@
 class DialectConfig:
-
     def __init__(self):
         pass
 
@@ -101,7 +100,6 @@ class PrestoConfig(DialectConfig):
 
 
 class TrinoConfig(PrestoConfig):
-
     def connection_string(self) -> str:
         if self.schema:
             return f"trino://{self.username}:{self.password}@{self.host}:{self.port}/{self.catalog}/{self.schema}"

trilogy/dialect/duckdb.py

@@ -1,8 +1,8 @@
-from typing import Mapping, Callable, Any
+from typing import Any, Callable, Mapping
 
 from jinja2 import Template
 
-from trilogy.core.enums import FunctionType, WindowType, UnnestMode
+from trilogy.core.enums import FunctionType, UnnestMode, WindowType
 from trilogy.dialect.base import BaseDialect
 
 WINDOW_FUNCTION_MAP: Mapping[WindowType, Callable[[Any, Any, Any], str]] = {}

trilogy/dialect/enums.py

@@ -1,12 +1,12 @@
 from enum import Enum
-from typing import List, TYPE_CHECKING, Optional, Callable
+from typing import TYPE_CHECKING, Callable, List, Optional
 
 if TYPE_CHECKING:
+    from trilogy import Environment, Executor
     from trilogy.hooks.base_hook import BaseHook
-    from trilogy import Executor, Environment
 
-from trilogy.dialect.config import DialectConfig
 from trilogy.constants import logger
+from trilogy.dialect.config import DialectConfig
 
 
 def default_factory(conf: DialectConfig, config_type):
@@ -42,6 +42,7 @@ class Dialects(Enum):
         if self == Dialects.BIGQUERY:
             from google.auth import default
             from google.cloud import bigquery
+
             from trilogy.dialect.config import BigQueryConfig
 
             credentials, project = default()
@@ -52,7 +53,6 @@ class Dialects(Enum):
                 BigQueryConfig,
             )
         elif self == Dialects.SQL_SERVER:
-
             raise NotImplementedError()
         elif self == Dialects.DUCK_DB:
             from trilogy.dialect.config import DuckDBConfig
@@ -98,7 +98,7 @@ class Dialects(Enum):
         conf: DialectConfig | None = None,
         _engine_factory: Callable | None = None,
     ) -> "Executor":
-        from trilogy import Executor, Environment
+        from trilogy import Environment, Executor
 
         if _engine_factory is not None:
             return Executor(

trilogy/dialect/postgres.py

@@ -1,8 +1,8 @@
-from typing import Mapping, Callable, Any
+from typing import Any, Callable, Mapping
 
 from jinja2 import Template
 
-from trilogy.core.enums import FunctionType, WindowType, DatePart
+from trilogy.core.enums import DatePart, FunctionType, WindowType
 from trilogy.dialect.base import BaseDialect
 
 

trilogy/dialect/presto.py

@@ -1,11 +1,10 @@
-from typing import Mapping, Callable, Any
+from typing import Any, Callable, Mapping
 
 from jinja2 import Template
 
-from trilogy.core.enums import FunctionType, WindowType
-from trilogy.dialect.base import BaseDialect
+from trilogy.core.enums import FunctionType, UnnestMode, WindowType
 from trilogy.core.models import DataType
-from trilogy.core.enums import UnnestMode
+from trilogy.dialect.base import BaseDialect
 
 WINDOW_FUNCTION_MAP: Mapping[WindowType, Callable[[Any, Any, Any], str]] = {}
 

trilogy/dialect/snowflake.py

@@ -1,8 +1,8 @@
-from typing import Mapping, Callable, Any
+from typing import Any, Callable, Mapping
 
 from jinja2 import Template
 
-from trilogy.core.enums import FunctionType, WindowType, UnnestMode
+from trilogy.core.enums import FunctionType, UnnestMode, WindowType
 from trilogy.dialect.base import BaseDialect
 
 ENV_SNOWFLAKE_PW = "PREQL_SNOWFLAKE_PW"

trilogy/dialect/sql_server.py

@@ -1,17 +1,16 @@
-from typing import Mapping, Callable, Any
+from typing import Any, Callable, Mapping
 
 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,
     ProcessedRawSQLStatement,
+    ProcessedShowStatement,
 )
 from trilogy.dialect.base import BaseDialect
+from trilogy.utility import string_to_hash
 
 WINDOW_FUNCTION_MAP: Mapping[WindowType, Callable[[Any, Any, Any], str]] = {}
 

trilogy/engine.py

@@ -1,6 +1,7 @@
-from sqlalchemy.engine import Engine, Connection, CursorResult
 from typing import Protocol
 
+from sqlalchemy.engine import Connection, CursorResult, Engine
+
 
 class EngineResult(Protocol):
     pass

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))}"