pytrilogy 0.0.1.110__py3-none-any.whl → 0.0.1.111__py3-none-any.whl

trilogy/core/query_processor.py

@@ -1,4 +1,4 @@
-from typing import List, Optional, Set, Union, Dict
+from typing import List, Optional, Set, Union, Dict, Tuple
 
 from trilogy.core.env_processor import generate_graph
 from trilogy.core.graph_models import ReferenceGraph
@@ -84,19 +84,31 @@ def base_join_to_join(
 
 def generate_source_map(
     query_datasource: QueryDatasource, all_new_ctes: List[CTE]
-) -> Dict[str, str | list[str]]:
+) -> Tuple[Dict[str, list[str]], Dict[str, list[str]]]:
     source_map: Dict[str, list[str]] = defaultdict(list)
     # now populate anything derived in this level
     for qdk, qdv in query_datasource.source_map.items():
+        unnest = [x for x in qdv if isinstance(x, UnnestJoin)]
+        for x in unnest:
+            source_map[qdk] = []
         if (
             qdk not in source_map
             and len(qdv) == 1
             and isinstance(list(qdv)[0], UnnestJoin)
         ):
             source_map[qdk] = []
+        basic = [x for x in qdv if isinstance(x, Datasource)]
+        for base in basic:
+            source_map[qdk].append(base.name)
 
-        else:
-            for cte in all_new_ctes:
+        ctes = [x for x in qdv if isinstance(x, QueryDatasource)]
+        if ctes:
+            names = set([x.name for x in ctes])
+            matches = [cte for cte in all_new_ctes if cte.source.name in names]
+
+            if not matches and names:
+                raise SyntaxError(query_datasource.source_map)
+            for cte in matches:
                 output_address = [
                     x.address
                     for x in cte.output_columns
@@ -105,21 +117,27 @@ def generate_source_map(
                 if qdk in output_address:
                     source_map[qdk].append(cte.name)
             # now do a pass that accepts partials
-            # TODO: move this into a second loop by first creationg all sub sourcdes
+            # TODO: move this into a second loop by first creationg all sub sources
             # then loop through this
-            for cte in all_new_ctes:
-                output_address = [x.address for x in cte.output_columns]
-                if qdk in output_address:
-                    if qdk not in source_map:
-                        source_map[qdk] = [cte.name]
-        if qdk not in source_map and not qdv:
-            # set source to empty, as it must be derived in this element
-            source_map[qdk] = []
+            for cte in matches:
+                if qdk not in source_map:
+                    source_map[qdk] = [cte.name]
         if qdk not in source_map:
-            raise ValueError(
-                f"Missing {qdk} in {source_map}, source map {query_datasource.source_map.keys()} "
-            )
-    return {k: "" if not v else v for k, v in source_map.items()}
+            if not qdv:
+                source_map[qdk] = []
+            elif CONFIG.validate_missing:
+                raise ValueError(
+                    f"Missing {qdk} in {source_map}, source map {query_datasource.source_map} "
+                )
+
+    # existence lookups use a separate map
+    # as they cannot be referenced in row resolution
+    existence_source_map: Dict[str, list[str]] = defaultdict(list)
+    for ek, ev in query_datasource.existence_source_map.items():
+        names = set([x.name for x in ev])
+        ematches = [cte.name for cte in all_new_ctes if cte.source.name in names]
+        existence_source_map[ek] = ematches
+    return {k: [] if not v else v for k, v in source_map.items()}, existence_source_map
 
 
 def datasource_to_query_datasource(datasource: Datasource) -> QueryDatasource:
@@ -158,6 +176,52 @@ def generate_cte_name(full_name: str, name_map: dict[str, str]) -> str:
         return full_name.replace("<", "").replace(">", "").replace(",", "_")
 
 
+def resolve_cte_base_name_and_alias(
+    name: str,
+    source: QueryDatasource,
+    parents: List[CTE],
+    joins: List[Join | InstantiatedUnnestJoin],
+) -> Tuple[str | None, str | None]:
+
+    valid_joins: List[Join] = [join for join in joins if isinstance(join, Join)]
+    relevant_parent_sources = set()
+    for k, v in source.source_map.items():
+        if v:
+            relevant_parent_sources.update(v)
+    eligible = [x for x in source.datasources if x in relevant_parent_sources]
+    if (
+        len(eligible) == 1
+        and isinstance(eligible[0], Datasource)
+        and not eligible[0].name == CONSTANT_DATASET
+    ):
+        ds = eligible[0]
+        return ds.safe_location, ds.identifier
+
+    # if we have multiple joined CTEs, pick the base
+    # as the root
+    elif len(eligible) == 1 and len(parents) == 1:
+        return parents[0].name, parents[0].name
+    elif valid_joins and len(valid_joins) > 0:
+        candidates = [x.left_cte.name for x in valid_joins]
+        disallowed = [x.right_cte.name for x in valid_joins]
+        try:
+            cte = [y for y in candidates if y not in disallowed][0]
+            return cte, cte
+        except IndexError:
+            raise SyntaxError(
+                f"Invalid join configuration {candidates} {disallowed} with all parents {[x.base_name for x in parents]}"
+            )
+    elif eligible:
+        matched = [x for x in parents if x.source.name == eligible[0].name]
+        if matched:
+            return matched[0].name, matched[0].name
+
+    logger.info(
+        f"Could not determine CTE base name for {name} with relevant sources {relevant_parent_sources}"
+    )
+    return None, None
+
+
 def datasource_to_ctes(
     query_datasource: QueryDatasource, name_map: dict[str, str]
 ) -> List[CTE]:
@@ -176,7 +240,8 @@ def datasource_to_ctes(
             sub_cte = datasource_to_ctes(sub_datasource, name_map)
             parents += sub_cte
             all_new_ctes += sub_cte
-        source_map = generate_source_map(query_datasource, all_new_ctes)
+        source_map, existence_map = generate_source_map(query_datasource, all_new_ctes)
+
     else:
         # source is the first datasource of the query datasource
         source = query_datasource.datasources[0]
@@ -184,13 +249,27 @@ def datasource_to_ctes(
         # render properly on initial access; since they have
         # no actual source
         if source.full_name == DEFAULT_NAMESPACE + "_" + CONSTANT_DATASET:
-            source_map = {k: "" for k in query_datasource.source_map}
+            source_map = {k: [] for k in query_datasource.source_map}
+            existence_map = source_map
         else:
             source_map = {
-                k: "" if not v else source.identifier
+                k: [] if not v else [source.identifier]
                 for k, v in query_datasource.source_map.items()
             }
+            existence_map = source_map
+
     human_id = generate_cte_name(query_datasource.full_name, name_map)
+    logger.info(
+        f"Finished building source map for {human_id} with {len(parents)} parents, have {source_map}, parent had non-empty keys {[k for k, v in query_datasource.source_map.items() if v]} "
+    )
+    final_joins = [
+        x
+        for x in [base_join_to_join(join, parents) for join in query_datasource.joins]
+        if x
+    ]
+    base_name, base_alias = resolve_cte_base_name_and_alias(
+        human_id, query_datasource, parents, final_joins
+    )
     cte = CTE(
         name=human_id,
         source=query_datasource,
@@ -200,14 +279,9 @@ def datasource_to_ctes(
             for c in query_datasource.output_concepts
         ],
         source_map=source_map,
+        existence_source_map=existence_map,
         # related columns include all referenced columns, such as filtering
-        joins=[
-            x
-            for x in [
-                base_join_to_join(join, parents) for join in query_datasource.joins
-            ]
-            if x
-        ],
+        joins=final_joins,
         grain=query_datasource.grain,
         group_to_grain=query_datasource.group_required,
         # we restrict parent_ctes to one level
@@ -217,11 +291,13 @@ def datasource_to_ctes(
         partial_concepts=query_datasource.partial_concepts,
         join_derived_concepts=query_datasource.join_derived_concepts,
         hidden_concepts=query_datasource.hidden_concepts,
+        base_name_override=base_name,
+        base_alias_override=base_alias,
     )
     if cte.grain != query_datasource.grain:
         raise ValueError("Grain was corrupted in CTE generation")
     for x in cte.output_columns:
-        if x.address not in cte.source_map:
+        if x.address not in cte.source_map and CONFIG.validate_missing:
             raise ValueError(
                 f"Missing {x.address} in {cte.source_map}, source map {cte.source.source_map.keys()} "
             )
@@ -243,12 +319,32 @@ def get_query_datasources(
     if not statement.output_components:
         raise ValueError(f"Statement has no output components {statement}")
     ds = source_query_concepts(
-        statement.output_components, environment=environment, g=graph
+        statement.output_components,
+        environment=environment,
+        g=graph,
     )
     if hooks:
         for hook in hooks:
             hook.process_root_strategy_node(ds)
     final_qds = ds.resolve()
+
+    # we if we have a where clause doing an existence check
+    # treat that as separate subquery
+    if (where := statement.where_clause) and where.existence_arguments:
+        for subselect in where.existence_arguments:
+            if not subselect:
+                continue
+            logger.info(
+                f"{LOGGER_PREFIX} fetching existance clause inputs {[str(c) for c in subselect]}"
+            )
+            eds = source_query_concepts([*subselect], environment=environment, g=graph)
+
+            final_eds = eds.resolve()
+            first_parent = final_qds
+            first_parent.datasources.append(final_eds)
+            for x in final_eds.output_concepts:
+                if x.address not in first_parent.existence_source_map:
+                    first_parent.existence_source_map[x.address] = {final_eds}
     return final_qds
 
 

trilogy/dialect/base.py

@@ -193,27 +193,6 @@ ORDER BY {% for order in order_by %}
 )
 
 
-def check_lineage(c: Concept, cte: CTE) -> bool:
-    checks = []
-    if not c.lineage:
-        return True
-    for sub_c in c.lineage.concept_arguments:
-        if not isinstance(sub_c, Concept):
-            continue
-        if sub_c.address in cte.source_map or (
-            sub_c.lineage and check_lineage(sub_c, cte)
-        ):
-            checks.append(True)
-        else:
-            logger.debug(
-                f"{LOGGER_PREFIX} [{sub_c.address}] not found in source map for"
-                f" {cte.name}, have cte keys {[c for c in cte.source_map.keys()]} and"
-                f" datasource keys {[c for c in cte.source.source_map.keys()]}"
-            )
-            checks.append(False)
-    return all(checks)
-
-
 def safe_quote(string: str, quote_char: str):
     # split dotted identifiers
     # TODO: evaluate if we need smarter parsing for strings that could actually include .
@@ -259,7 +238,7 @@ class BaseDialect:
             f"{LOGGER_PREFIX} [{c.address}] Starting rendering loop on cte: {cte.name}"
         )
 
-        if c.lineage and cte.source_map.get(c.address, "") == "":
+        if c.lineage and cte.source_map.get(c.address, []) == []:
             logger.debug(
                 f"{LOGGER_PREFIX} [{c.address}] rendering concept with lineage that is not already existing"
             )
@@ -273,7 +252,11 @@ class BaseDialect:
                 ]
                 rval = f"{self.WINDOW_FUNCTION_MAP[c.lineage.type](concept = self.render_concept_sql(c.lineage.content, cte=cte, alias=False), window=','.join(rendered_over_components), sort=','.join(rendered_order_components))}"  # noqa: E501
             elif isinstance(c.lineage, FilterItem):
-                rval = f"CASE WHEN {self.render_expr(c.lineage.where.conditional, cte=cte)} THEN {self.render_concept_sql(c.lineage.content, cte=cte, alias=False)} ELSE NULL END"
+                # for cases when we've optimized this
+                if len(cte.output_columns) == 1:
+                    rval = self.render_expr(c.lineage.content, cte=cte)
+                else:
+                    rval = f"CASE WHEN {self.render_expr(c.lineage.where.conditional, cte=cte)} THEN {self.render_concept_sql(c.lineage.content, cte=cte, alias=False)} ELSE NULL END"
             elif isinstance(c.lineage, RowsetItem):
                 rval = f"{self.render_concept_sql(c.lineage.content, cte=cte, alias=False)}"
             elif isinstance(c.lineage, MultiSelectStatement):
@@ -356,17 +339,28 @@ class BaseDialect:
         cte: Optional[CTE] = None,
         cte_map: Optional[Dict[str, CTE]] = None,
     ) -> str:
-        # if isinstance(e, Concept):
-        #     cte = cte or cte_map.get(e.address, None)
 
         if isinstance(e, SubselectComparison):
-            assert cte, "Subselects must be rendered with a CTE in context"
+
             if isinstance(e.right, Concept):
-                return f"{self.render_expr(e.left, cte=cte, cte_map=cte_map)} {e.operator.value} (select {self.render_expr(e.right, cte=cte, cte_map=cte_map)} from {cte.source_map[e.right.address][0]})"
+                # we won't always have an existnce map
+                # so fall back to the normal map
+                lookup_cte = cte
+                if cte_map and not lookup_cte:
+                    lookup_cte = cte_map.get(e.right.address)
+                assert lookup_cte, "Subselects must be rendered with a CTE in context"
+                if e.right.address not in lookup_cte.existence_source_map:
+                    lookup = lookup_cte.source_map[e.right.address]
+                else:
+                    lookup = lookup_cte.existence_source_map[e.right.address]
+
+                return f"{self.render_expr(e.left, cte=cte, cte_map=cte_map)} {e.operator.value} (select {lookup[0]}.{self.QUOTE_CHARACTER}{e.right.safe_address}{self.QUOTE_CHARACTER} from {lookup[0]})"
+            elif isinstance(e.right, (ListWrapper, Parenthetical)):
+                return f"{self.render_expr(e.left, cte=cte, cte_map=cte_map)} {e.operator.value} {self.render_expr(e.right, cte=cte, cte_map=cte_map)}"
+            elif isinstance(e.right, (str, int, bool, float, list)):
+                return f"{self.render_expr(e.left, cte=cte, cte_map=cte_map)} {e.operator.value} ({self.render_expr(e.right, cte=cte, cte_map=cte_map)})"
             else:
-                raise NotImplementedError(
-                    f"Subselects must be a concept, got {e.right}"
-                )
+                return f"{self.render_expr(e.left, cte=cte, cte_map=cte_map)} {e.operator.value} ({self.render_expr(e.right, cte=cte, cte_map=cte_map)})"
         elif isinstance(e, Comparison):
             return f"{self.render_expr(e.left, cte=cte, cte_map=cte_map)} {e.operator.value} {self.render_expr(e.right, cte=cte, cte_map=cte_map)}"
         elif isinstance(e, Conditional):
@@ -449,15 +443,15 @@ class BaseDialect:
                 for c in cte.output_columns
                 if c.address not in [y.address for y in cte.hidden_concepts]
             ]
+        if cte.base_name == cte.base_alias:
+            source = cte.base_name
+        else:
+            source = f"{cte.base_name} as {cte.base_alias}"
         return CompiledCTE(
             name=cte.name,
             statement=self.SQL_TEMPLATE.render(
                 select_columns=select_columns,
-                base=(
-                    f"{cte.base_name} as {cte.base_alias}"
-                    if cte.render_from_clause
-                    else None
-                ),
+                base=(f"{source}" if cte.render_from_clause else None),
                 grain=cte.grain,
                 limit=cte.limit,
                 # some joins may not need to be rendered
@@ -513,7 +507,7 @@ class BaseDialect:
                                         c
                                         for c in cte.output_columns
                                         if c.purpose == Purpose.CONSTANT
-                                        and cte.source_map[c.address] != ""
+                                        and cte.source_map[c.address] != []
                                     ],
                                     "address",
                                 )
@@ -639,7 +633,7 @@ class BaseDialect:
             filter = set(
                 [
                     str(x.address)
-                    for x in query.where_clause.concept_arguments
+                    for x in query.where_clause.row_arguments
                     if not x.derivation == PurposeLineage.CONSTANT
                 ]
             )
@@ -650,10 +644,21 @@ class BaseDialect:
 
             if not found:
                 raise NotImplementedError(
-                    f"Cannot generate query with filtering on grain {filter} that is"
-                    f" not a subset of the query output grain {query_output}. Use a"
-                    " filtered concept instead."
+                    f"Cannot generate query with filtering on row arguments {filter} that is"
+                    f" not a subset of the query output grain {query_output}. Try a"
+                    " filtered concept instead, or include it in the select clause"
                 )
+            for ex_set in query.where_clause.existence_arguments:
+                for c in ex_set:
+                    if c.address not in cte_output_map:
+                        cts = [
+                            ct
+                            for ct in query.ctes
+                            if ct.name in query.base.existence_source_map[c.address]
+                        ]
+                        if not cts:
+                            raise ValueError(query.base.existence_source_map[c.address])
+                        cte_output_map[c.address] = cts[0]
 
         compiled_ctes = self.generate_ctes(query)
 

trilogy/executor.py

@@ -1,4 +1,4 @@
-from typing import List, Optional, Any
+from typing import List, Optional, Any, Generator
 from functools import singledispatchmethod
 from sqlalchemy import text
 from sqlalchemy.engine import Engine, CursorResult
@@ -222,6 +222,35 @@ class Executor(object):
             sql.append(x)
         return sql
 
+    def parse_text_generator(
+        self, command: str, persist: bool = False
+    ) -> Generator[
+        ProcessedQuery | ProcessedQueryPersist | ProcessedShowStatement, None, None
+    ]:
+        """Process a preql text command"""
+        _, parsed = parse_text(command, self.environment)
+        generatable = [
+            x
+            for x in parsed
+            if isinstance(
+                x,
+                (
+                    SelectStatement,
+                    PersistStatement,
+                    MultiSelectStatement,
+                    ShowStatement,
+                ),
+            )
+        ]
+        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)
+            yield x
+
     def execute_raw_sql(self, command: str) -> CursorResult:
         """Run a command against the raw underlying
         execution engine"""
@@ -229,10 +258,9 @@ class Executor(object):
 
     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:
+        for statement in self.parse_text_generator(command):
             if isinstance(statement, ProcessedShowStatement):
                 output.append(
                     generate_result_set(

trilogy/parsing/parse_engine.py

@@ -101,6 +101,7 @@ from trilogy.core.models import (
     ConceptDerivation,
     RowsetDerivationStatement,
     LooseConceptList,
+    list_to_wrapper,
 )
 from trilogy.parsing.exceptions import ParseError
 from trilogy.utility import string_to_hash
@@ -113,6 +114,7 @@ from trilogy.parsing.common import (
     arbitrary_to_concept,
 )
 
+
 CONSTANT_TYPES = (int, float, str, bool, ListWrapper)
 
 grammar = r"""
@@ -138,7 +140,7 @@ grammar = r"""
     //<customer_id,country>.property local_alias STRING
     concept_property_declaration: PROPERTY (prop_ident | IDENTIFIER) data_type concept_nullable_modifier? metadata?
     //metric post_length <- len(post_text);
-    concept_derivation:  (PURPOSE | AUTO | PROPERTY ) IDENTIFIER "<" "-" expr
+    concept_derivation:  (PURPOSE | AUTO | PROPERTY ) (prop_ident | IDENTIFIER) "<" "-" expr
 
     rowset_derivation_statement: ("rowset"i IDENTIFIER "<" "-" (multi_select_statement | select_statement)) | ("with"i IDENTIFIER "as"i (multi_select_statement | select_statement))
     
@@ -179,13 +181,11 @@ grammar = r"""
     // multiple_selects
     multi_select_statement: select_statement ("merge" select_statement)+ "align"i align_clause  where? comment* order_by? comment* limit? comment*
 
-
     align_item: IDENTIFIER ":" IDENTIFIER ("," IDENTIFIER)* ","?
 
     align_clause: align_item ("," align_item)*  ","?
 
     // merge statemment
-
     merge_statement: "merge" IDENTIFIER ("," IDENTIFIER)* ","? comment*
 
     // FUNCTION blocks
@@ -193,7 +193,6 @@ grammar = r"""
     function_binding_item: IDENTIFIER ":" data_type
     function_binding_list: (function_binding_item ",")* function_binding_item ","?
     raw_function: "bind" "sql"  IDENTIFIER  "("  function_binding_list ")" "-" ">" data_type "as"i MULTILINE_STRING
-
     
     // user_id where state = Mexico
     filter_item: "filter"i IDENTIFIER where
@@ -249,9 +248,9 @@ grammar = r"""
 
     COMPARISON_OPERATOR: (/is[\s]+not/ | "is" |"=" | ">" | "<" | ">=" | "<=" | "!="  )
     
-    comparison: (expr COMPARISON_OPERATOR expr) | (expr array_comparison expr_tuple)
+    comparison: (expr COMPARISON_OPERATOR expr) 
 
-    subselect_comparison: expr array_comparison expr
+    subselect_comparison: expr array_comparison expr | (expr array_comparison expr_tuple)
     
     expr_tuple: "("  (expr ",")* expr ","?  ")"
 
@@ -267,7 +266,6 @@ grammar = r"""
     
     // functions
     
-    //math TODO: add syntactic sugar
     fadd: ("add"i "(" expr "," expr ")" ) | ( expr "+" expr )
     fsub: ("subtract"i "(" expr "," expr ")" ) | ( expr "-" expr )
     fmul: ("multiply"i "(" expr "," expr ")" ) | ( expr "*" expr )
@@ -738,10 +736,17 @@ class ParseToObjects(Transformer):
         purpose = args[0]
         if purpose == Purpose.AUTO:
             purpose = None
-        name = args[1]
-        lookup, namespace, name, parent_concept = parse_concept_reference(
-            name, self.environment, purpose
-        )
+        raw_name = args[1]
+        if isinstance(raw_name, str):
+            lookup, namespace, name, parent_concept = parse_concept_reference(
+                raw_name, self.environment, purpose
+            )
+        else:
+            keys, name = raw_name
+            if "." in name:
+                namespace, name = name.rsplit(".", 1)
+            else:
+                namespace = self.environment.namespace or DEFAULT_NAMESPACE
         source_value = args[2]
         # we need to strip off every parenthetical to see what is being assigned.
         while isinstance(source_value, Parenthetical):
@@ -971,7 +976,26 @@ class ParseToObjects(Transformer):
         return Ordering(args.lower())
 
     def order_list(self, args):
-        return [OrderItem(expr=x, order=y) for x, y in zip(args[::2], args[1::2])]
+
+        def handle_order_item(x, namespace: str):
+            if not isinstance(x, Concept):
+                x = arbitrary_to_concept(
+                    x,
+                    namespace=namespace,
+                    name=f"{VIRTUAL_CONCEPT_PREFIX}_{string_to_hash(str(x))}",
+                )
+            return x
+
+        return [
+            OrderItem(
+                expr=handle_order_item(
+                    x,
+                    self.environment.namespace,
+                ),
+                order=y,
+            )
+            for x, y in zip(args[::2], args[1::2])
+        ]
 
     def order_by(self, args):
         return OrderBy(items=args[0])
@@ -1207,26 +1231,34 @@ class ParseToObjects(Transformer):
         return float(args[0])
 
     def array_lit(self, args):
-        types = [arg_to_datatype(arg) for arg in args]
-        assert len(set(types)) == 1
-        return ListWrapper(args, type=types[0])
+        return list_to_wrapper(args)
 
     def literal(self, args):
         return args[0]
 
     def comparison(self, args) -> Comparison:
+        if args[1] == ComparisonOperator.IN:
+            raise SyntaxError
         return Comparison(left=args[0], right=args[2], operator=args[1])
 
     @v_args(meta=True)
     def subselect_comparison(self, meta: Meta, args) -> SubselectComparison:
         right = args[2]
-        if not isinstance(right, Concept):
+
+        while isinstance(right, Parenthetical) and isinstance(
+            right.content,
+            (Concept, Function, FilterItem, WindowItem, AggregateWrapper, ListWrapper),
+        ):
+            right = right.content
+        if isinstance(
+            right, (Function, FilterItem, WindowItem, AggregateWrapper, ListWrapper)
+        ):
             right = arbitrary_to_concept(
                 right,
                 namespace=self.environment.namespace,
                 name=f"{VIRTUAL_CONCEPT_PREFIX}_{string_to_hash(str(right))}",
             )
-            self.environment.add_concept(right)
+            self.environment.add_concept(right, meta=meta)
         return SubselectComparison(
             left=args[0],
             right=right,