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

trilogy/core/query_processor.py

@@ -1,43 +1,41 @@
-from typing import List, Optional, Set, Union, Dict, Tuple
+from collections import defaultdict
+from math import ceil
+from typing import Dict, List, Optional, Set, Tuple, Union
 
+from trilogy.constants import CONFIG, logger
+from trilogy.core.constants import CONSTANT_DATASET
+from trilogy.core.enums import BooleanOperator, SourceType
 from trilogy.core.env_processor import generate_graph
+from trilogy.core.ergonomics import generate_cte_names
 from trilogy.core.graph_models import ReferenceGraph
-from trilogy.core.constants import CONSTANT_DATASET
-from trilogy.core.processing.concept_strategies_v3 import source_query_concepts
-from trilogy.core.enums import BooleanOperator
-from trilogy.constants import CONFIG
-from trilogy.core.processing.nodes import SelectNode, StrategyNode, History
 from trilogy.core.models import (
+    CTE,
+    BaseJoin,
     Concept,
-    Environment,
-    PersistStatement,
     ConceptDeclarationStatement,
-    SelectStatement,
-    MultiSelectStatement,
-    CTE,
+    Conditional,
+    CopyStatement,
+    CTEConceptPair,
+    Datasource,
+    Environment,
+    InstantiatedUnnestJoin,
     Join,
-    UnnestJoin,
     MaterializedDataset,
+    MultiSelectStatement,
+    PersistStatement,
+    ProcessedCopyStatement,
     ProcessedQuery,
     ProcessedQueryPersist,
     QueryDatasource,
-    Datasource,
-    BaseJoin,
-    InstantiatedUnnestJoin,
-    Conditional,
-    ProcessedCopyStatement,
-    CopyStatement,
-    CTEConceptPair,
+    SelectStatement,
+    UnionCTE,
+    UnnestJoin,
 )
-
-from trilogy.utility import unique
-
-from trilogy.hooks.base_hook import BaseHook
-from trilogy.constants import logger
-from trilogy.core.ergonomics import generate_cte_names
 from trilogy.core.optimization import optimize_ctes
-from math import ceil
-from collections import defaultdict
+from trilogy.core.processing.concept_strategies_v3 import source_query_concepts
+from trilogy.core.processing.nodes import History, SelectNode, StrategyNode
+from trilogy.hooks.base_hook import BaseHook
+from trilogy.utility import unique
 
 LOGGER_PREFIX = "[QUERY BUILD]"
 
@@ -103,7 +101,7 @@ def base_join_to_join(
 
 
 def generate_source_map(
-    query_datasource: QueryDatasource, all_new_ctes: List[CTE]
+    query_datasource: QueryDatasource, all_new_ctes: List[CTE | UnionCTE]
 ) -> Tuple[Dict[str, list[str]], Dict[str, list[str]]]:
     source_map: Dict[str, list[str]] = defaultdict(list)
     # now populate anything derived in this level
@@ -246,24 +244,44 @@ def resolve_cte_base_name_and_alias_v2(
     return None, None
 
 
-def datasource_to_ctes(
+def datasource_to_cte(
     query_datasource: QueryDatasource, name_map: dict[str, str]
-) -> List[CTE]:
-    output: List[CTE] = []
-    parents: list[CTE] = []
+) -> CTE | UnionCTE:
+    parents: list[CTE | UnionCTE] = []
+    if query_datasource.source_type == SourceType.UNION:
+        direct_parents: list[CTE | UnionCTE] = []
+        for child in query_datasource.datasources:
+            assert isinstance(child, QueryDatasource)
+            child_cte = datasource_to_cte(child, name_map=name_map)
+            direct_parents.append(child_cte)
+            parents += child_cte.parent_ctes
+        human_id = generate_cte_name(query_datasource.identifier, name_map)
+        final = UnionCTE(
+            name=human_id,
+            source=query_datasource,
+            parent_ctes=parents,
+            internal_ctes=direct_parents,
+            output_columns=[
+                c.with_grain(query_datasource.grain)
+                for c in query_datasource.output_concepts
+            ],
+            grain=direct_parents[0].grain,
+        )
+        return final
+
     if len(query_datasource.datasources) > 1 or any(
         [isinstance(x, QueryDatasource) for x in query_datasource.datasources]
     ):
-        all_new_ctes: List[CTE] = []
+        all_new_ctes: List[CTE | UnionCTE] = []
         for datasource in query_datasource.datasources:
             if isinstance(datasource, QueryDatasource):
                 sub_datasource = datasource
             else:
                 sub_datasource = datasource_to_query_datasource(datasource)
 
-            sub_cte = datasource_to_ctes(sub_datasource, name_map)
-            parents += sub_cte
-            all_new_ctes += sub_cte
+            sub_cte = datasource_to_cte(sub_datasource, name_map)
+            parents.append(sub_cte)
+            all_new_ctes.append(sub_cte)
         source_map, existence_map = generate_source_map(query_datasource, all_new_ctes)
 
     else:
@@ -284,7 +302,10 @@ def datasource_to_ctes(
 
     human_id = generate_cte_name(query_datasource.identifier, name_map)
 
-    final_joins = [base_join_to_join(join, parents) for join in query_datasource.joins]
+    final_joins = [
+        base_join_to_join(join, [x for x in parents if isinstance(x, CTE)])
+        for join in query_datasource.joins
+    ]
 
     base_name, base_alias = resolve_cte_base_name_and_alias_v2(
         human_id, query_datasource, source_map, final_joins
@@ -326,8 +347,7 @@ def datasource_to_ctes(
                 f"Missing {x.address} in {cte.source_map}, source map {cte.source.source_map.keys()} "
             )
 
-    output.append(cte)
-    return output
+    return cte
 
 
 def get_query_node(
@@ -336,6 +356,9 @@ def get_query_node(
     graph: Optional[ReferenceGraph] = None,
     history: History | None = None,
 ) -> StrategyNode:
+    environment = environment.duplicate()
+    for k, v in statement.local_concepts.items():
+        environment.concepts[k] = v
     graph = graph or generate_graph(environment)
     logger.info(
         f"{LOGGER_PREFIX} getting source datasource for query with filtering {statement.where_clause_category} and output {[str(c) for c in statement.output_components]}"
@@ -393,7 +416,7 @@ def get_query_datasources(
     return final_qds
 
 
-def flatten_ctes(input: CTE) -> list[CTE]:
+def flatten_ctes(input: CTE | UnionCTE) -> list[CTE | UnionCTE]:
     output = [input]
     for cte in input.parent_ctes:
         output += flatten_ctes(cte)
@@ -456,18 +479,18 @@ def process_query(
     hooks: List[BaseHook] | None = None,
 ) -> ProcessedQuery:
     hooks = hooks or []
-    statement.refresh_bindings(environment)
     graph = generate_graph(environment)
+
     root_datasource = get_query_datasources(
         environment=environment, graph=graph, statement=statement, hooks=hooks
     )
     for hook in hooks:
         hook.process_root_datasource(root_datasource)
     # this should always return 1 - TODO, refactor
-    root_cte = datasource_to_ctes(root_datasource, environment.cte_name_map)[0]
+    root_cte = datasource_to_cte(root_datasource, environment.cte_name_map)
     for hook in hooks:
         hook.process_root_cte(root_cte)
-    raw_ctes: List[CTE] = list(reversed(flatten_ctes(root_cte)))
+    raw_ctes: List[CTE | UnionCTE] = list(reversed(flatten_ctes(root_cte)))
     seen = dict()
     # we can have duplicate CTEs at this point
     # so merge them together
@@ -479,7 +502,7 @@ def process_query(
             seen[cte.name] = seen[cte.name] + cte
     for cte in raw_ctes:
         cte.parent_ctes = [seen[x.name] for x in cte.parent_ctes]
-    deduped_ctes: List[CTE] = list(seen.values())
+    deduped_ctes: List[CTE | UnionCTE] = list(seen.values())
     root_cte.order_by = statement.order_by
     root_cte.limit = statement.limit
     root_cte.hidden_concepts = [x for x in statement.hidden_components]
@@ -497,4 +520,5 @@ def process_query(
         # we no longer do any joins at final level, this should always happen in parent CTEs
         joins=[],
         hidden_columns=[x for x in statement.hidden_components],
+        local_concepts=statement.local_concepts,
     )

trilogy/dialect/base.py

@@ -1,66 +1,67 @@
-from typing import List, Union, Optional, Dict, Any, Sequence, Callable
+from typing import Any, Callable, Dict, List, Optional, Sequence, Union
 
 from jinja2 import Template
 
-from trilogy.core.processing.utility import (
-    is_scalar_condition,
-    decompose_condition,
-    sort_select_output,
-)
-from trilogy.constants import CONFIG, logger, MagicConstants
-from trilogy.core.internal import DEFAULT_CONCEPTS
+from trilogy.constants import CONFIG, MagicConstants, logger
 from trilogy.core.enums import (
+    DatePart,
     FunctionType,
+    UnnestMode,
     WindowType,
-    DatePart,
 )
+from trilogy.core.internal import DEFAULT_CONCEPTS
 from trilogy.core.models import (
-    ListType,
-    DataType,
-    Concept,
     CTE,
-    ProcessedQuery,
-    ProcessedQueryPersist,
-    ProcessedShowStatement,
+    AggregateWrapper,
+    CaseElse,
+    CaseWhen,
+    Comparison,
     CompiledCTE,
+    Concept,
+    ConceptDeclarationStatement,
     Conditional,
-    Comparison,
-    SubselectComparison,
-    OrderItem,
-    WindowItem,
+    CopyStatement,
+    Datasource,
+    DataType,
+    Environment,
     FilterItem,
     Function,
-    AggregateWrapper,
-    Parenthetical,
-    CaseWhen,
-    CaseElse,
-    SelectStatement,
-    PersistStatement,
-    Environment,
-    RawColumnExpr,
+    ImportStatement,
+    ListType,
     ListWrapper,
-    TupleWrapper,
+    MapType,
     MapWrapper,
-    ShowStatement,
-    RowsetItem,
+    MergeStatementV2,
     MultiSelectStatement,
-    RowsetDerivationStatement,
-    ConceptDeclarationStatement,
-    ImportStatement,
-    RawSQLStatement,
-    ProcessedRawSQLStatement,
     NumericType,
-    MapType,
-    StructType,
-    MergeStatementV2,
-    Datasource,
-    CopyStatement,
+    OrderItem,
+    Parenthetical,
+    PersistStatement,
     ProcessedCopyStatement,
+    ProcessedQuery,
+    ProcessedQueryPersist,
+    ProcessedRawSQLStatement,
+    ProcessedShowStatement,
+    RawColumnExpr,
+    RawSQLStatement,
+    RowsetDerivationStatement,
+    RowsetItem,
+    SelectStatement,
+    ShowStatement,
+    StructType,
+    SubselectComparison,
+    TupleWrapper,
+    UnionCTE,
+    WindowItem,
 )
-from trilogy.core.query_processor import process_query, process_persist, process_copy
+from trilogy.core.processing.utility import (
+    decompose_condition,
+    is_scalar_condition,
+    sort_select_output,
+)
+from trilogy.core.query_processor import process_copy, process_persist, process_query
 from trilogy.dialect.common import render_join, render_unnest
 from trilogy.hooks.base_hook import BaseHook
-from trilogy.core.enums import UnnestMode
 
 LOGGER_PREFIX = "[RENDERING]"
 
@@ -130,7 +131,7 @@ FUNCTION_MAP = {
     FunctionType.SPLIT: lambda x: f"split({x[0]}, {x[1]})",
     FunctionType.IS_NULL: lambda x: f"isnull({x[0]})",
     FunctionType.BOOL: lambda x: f"CASE WHEN {x[0]} THEN TRUE ELSE FALSE END",
-    # complex
+    # Complex
     FunctionType.INDEX_ACCESS: lambda x: f"{x[0]}[{x[1]}]",
     FunctionType.MAP_ACCESS: lambda x: f"{x[0]}[{x[1]}][1]",
     FunctionType.UNNEST: lambda x: f"unnest({x[0]})",
@@ -230,7 +231,7 @@ def safe_quote(string: str, quote_char: str):
     return ".".join([f"{quote_char}{string}{quote_char}" for string in components])
 
 
-def safe_get_cte_value(coalesce, cte: CTE, c: Concept, quote_char: str):
+def safe_get_cte_value(coalesce, cte: CTE | UnionCTE, c: Concept, quote_char: str):
     address = c.address
     raw = cte.source_map.get(address, None)
 
@@ -255,15 +256,26 @@ class BaseDialect:
     UNNEST_MODE = UnnestMode.CROSS_APPLY
 
     def render_order_item(
-        self, order_item: OrderItem, cte: CTE, final: bool = False
+        self,
+        order_item: OrderItem,
+        cte: CTE | UnionCTE,
+        final: bool = False,
+        alias: bool = True,
     ) -> str:
         if final:
+            if not alias:
+                return f"{self.QUOTE_CHARACTER}{order_item.expr.safe_address}{self.QUOTE_CHARACTER} {order_item.order.value}"
+
             return f"{cte.name}.{self.QUOTE_CHARACTER}{order_item.expr.safe_address}{self.QUOTE_CHARACTER} {order_item.order.value}"
 
         return f"{self.render_concept_sql(order_item.expr, cte=cte, alias=False)} {order_item.order.value}"
 
     def render_concept_sql(
-        self, c: Concept, cte: CTE, alias: bool = True, raise_invalid: bool = False
+        self,
+        c: Concept,
+        cte: CTE | UnionCTE,
+        alias: bool = True,
+        raise_invalid: bool = False,
     ) -> str:
         result = None
         if c.pseudonyms:
@@ -290,7 +302,7 @@ class BaseDialect:
         return result
 
     def _render_concept_sql(
-        self, c: Concept, cte: CTE, raise_invalid: bool = False
+        self, c: Concept, cte: CTE | UnionCTE, raise_invalid: bool = False
     ) -> str:
         # only recurse while it's in sources of the current cte
         logger.debug(
@@ -348,6 +360,20 @@ class BaseDialect:
                         " target grain"
                     )
                     rval = f"{self.FUNCTION_GRAIN_MATCH_MAP[c.lineage.function.operator](args)}"
+            elif (
+                isinstance(c.lineage, Function)
+                and c.lineage.operator == FunctionType.UNION
+            ):
+                local_matched = [
+                    x
+                    for x in c.lineage.arguments
+                    if isinstance(x, Concept) and x.address in cte.output_columns
+                ]
+                if not local_matched:
+                    raise SyntaxError(
+                        "Could not find appropriate source element for union"
+                    )
+                rval = self.render_expr(local_matched[0], cte)
             elif (
                 isinstance(c.lineage, Function)
                 and c.lineage.operator == FunctionType.CONSTANT
@@ -447,13 +473,11 @@ class BaseDialect:
             FilterItem,
             # FilterItem
         ],
-        cte: Optional[CTE] = None,
-        cte_map: Optional[Dict[str, CTE]] = None,
+        cte: Optional[CTE | UnionCTE] = None,
+        cte_map: Optional[Dict[str, CTE | UnionCTE]] = None,
         raise_invalid: bool = False,
     ) -> str:
-
         if isinstance(e, SubselectComparison):
-
             if isinstance(e.right, Concept):
                 # we won't always have an existnce map
                 # so fall back to the normal map
@@ -559,6 +583,13 @@ class BaseDialect:
         elif isinstance(e, FilterItem):
             return f"CASE WHEN {self.render_expr(e.where.conditional,cte=cte, cte_map=cte_map, raise_invalid=raise_invalid)} THEN {self.render_expr(e.content, cte, cte_map=cte_map, raise_invalid=raise_invalid)} ELSE NULL END"
         elif isinstance(e, Concept):
+            if (
+                isinstance(e.lineage, Function)
+                and e.lineage.operator == FunctionType.CONSTANT
+                and CONFIG.rendering.parameters is True
+                and e.datatype.data_type != DataType.MAP
+            ):
+                return f":{e.safe_address}"
             if cte:
                 return self.render_concept_sql(
                     e, cte, alias=False, raise_invalid=raise_invalid
@@ -592,7 +623,18 @@ class BaseDialect:
         else:
             raise ValueError(f"Unable to render type {type(e)} {e}")
 
-    def render_cte(self, cte: CTE, auto_sort: bool = True) -> CompiledCTE:
+    def render_cte(self, cte: CTE | UnionCTE, auto_sort: bool = True) -> CompiledCTE:
+        if isinstance(cte, UnionCTE):
+            base_statement = f"\n{cte.operator}\n".join(
+                [self.render_cte(child).statement for child in cte.internal_ctes]
+            )
+            if cte.order_by:
+                ordering = [
+                    self.render_order_item(i, cte, final=True, alias=False)
+                    for i in cte.order_by.items
+                ]
+                base_statement += "\nORDER BY " + ",".join(ordering)
+            return CompiledCTE(name=cte.name, statement=base_statement)
         if self.UNNEST_MODE in (
             UnnestMode.CROSS_APPLY,
             UnnestMode.CROSS_JOIN,
@@ -863,7 +905,7 @@ class BaseDialect:
         if CONFIG.strict_mode and INVALID_REFERENCE_STRING(1) in final:
             raise ValueError(
                 f"Invalid reference string found in query: {final}, this should never"
-                " occur. Please report this issue."
+                " occur. Please create a GitHub issue to report this."
             )
         logger.info(f"{LOGGER_PREFIX} Compiled query: {final}")
         return final

trilogy/dialect/bigquery.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, 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]] = {}
 
 FUNCTION_MAP = {

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