pytrilogy 0.0.1.105__py3-none-any.whl → 0.0.1.106__py3-none-any.whl

trilogy/core/query_processor.py

@@ -29,6 +29,7 @@ from trilogy.hooks.base_hook import BaseHook
 from trilogy.constants import logger
 from random import shuffle
 from trilogy.core.ergonomics import CTE_NAMES
+from trilogy.core.optimization import optimize_ctes
 from math import ceil
 
 LOGGER_PREFIX = "[QUERY BUILD]"
@@ -186,7 +187,7 @@ def datasource_to_ctes(
             source_map = {k: "" for k in query_datasource.source_map}
         else:
             source_map = {
-                k: "" if not v else source.full_name
+                k: "" if not v else source.identifier
                 for k, v in query_datasource.source_map.items()
             }
     human_id = generate_cte_name(query_datasource.full_name, name_map)
@@ -315,7 +316,9 @@ 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]
-    final_ctes: List[CTE] = list(seen.values())
+    deduped_ctes: List[CTE] = list(seen.values())
+
+    final_ctes = optimize_ctes(deduped_ctes, root_cte, statement)
 
     return ProcessedQuery(
         order_by=statement.order_by,

trilogy/dialect/base.py

@@ -170,6 +170,9 @@ GENERIC_SQL_TEMPLATE = Template(
     """{%- if ctes %}
 WITH {% for cte in ctes %}
 {{cte.name}} as ({{cte.statement}}){% if not loop.last %},{% endif %}{% endfor %}{% endif %}
+{%- if full_select -%}
+{{full_select}}
+{%- else -%}
 SELECT
 {%- if limit is not none %}
 TOP {{ limit }}{% endif %}
@@ -184,8 +187,8 @@ TOP {{ limit }}{% endif %}
 \t{{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 %}
+    {{ order }}{% if not loop.last %},{% endif %}{% endfor %}
+{% endif %}{% endif %}
 """
 )
 
@@ -218,15 +221,19 @@ 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, address: str, rendered: str):
+def safe_get_cte_value(coalesce, cte: CTE, c: Concept, quote_char: str):
+    address = c.address
     raw = cte.source_map.get(address, None)
+
     if not raw:
         return INVALID_REFERENCE_STRING("Missing source reference")
     if isinstance(raw, str):
-        return f"{raw}.{rendered}"
+        rendered = cte.get_alias(c, raw)
+        return f"{raw}.{quote_char}{rendered}{quote_char}"
     if isinstance(raw, list) and len(raw) == 1:
-        return f"{raw[0]}.{rendered}"
-    return coalesce([f"{x}.{rendered}" for x in raw])
+        rendered = cte.get_alias(c, raw[0])
+        return f"{raw[0]}.{quote_char}{rendered}{quote_char}"
+    return coalesce([f"{x}.{quote_char}{cte.get_alias(c, x)}{quote_char}" for x in raw])
 
 
 class BaseDialect:
@@ -238,21 +245,13 @@ class BaseDialect:
     DATATYPE_MAP = DATATYPE_MAP
     UNNEST_MODE = UnnestMode.CROSS_APPLY
 
-    def render_order_item(self, order_item: OrderItem, ctes: List[CTE]) -> str:
-        matched_ctes = [
-            cte
-            for cte in ctes
-            if order_item.expr.address in [a.address for a in cte.output_columns]
-        ]
-        if not matched_ctes:
-            all_outputs = set()
-            for cte in ctes:
-                all_outputs.update([a.address for a in cte.output_columns])
-            raise ValueError(
-                f"No source found for concept {order_item.expr}, have {all_outputs}"
-            )
-        selected = matched_ctes[0]
-        return f"{selected.name}.{self.QUOTE_CHARACTER}{order_item.expr.safe_address}{self.QUOTE_CHARACTER} {order_item.order.value}"
+    def render_order_item(
+        self, order_item: OrderItem, cte: CTE, final: bool = False
+    ) -> str:
+        if final:
+            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) -> str:
         # only recurse while it's in sources of the current cte
@@ -310,13 +309,14 @@ class BaseDialect:
             logger.debug(
                 f"{LOGGER_PREFIX} [{c.address}] Rendering basic lookup from {cte.source_map.get(c.address, INVALID_REFERENCE_STRING('Missing source reference'))}"
             )
+
             raw_content = cte.get_alias(c)
             if isinstance(raw_content, RawColumnExpr):
                 rval = raw_content.text
             elif isinstance(raw_content, Function):
                 rval = self.render_expr(raw_content, cte=cte)
             else:
-                rval = f"{safe_get_cte_value(self.FUNCTION_MAP[FunctionType.COALESCE], cte, c.address, rendered=safe_quote(raw_content, self.QUOTE_CHARACTER))}"
+                rval = f"{safe_get_cte_value(self.FUNCTION_MAP[FunctionType.COALESCE], cte, c, self.QUOTE_CHARACTER)}"
         if alias:
             return (
                 f"{rval} as"
@@ -456,7 +456,7 @@ class BaseDialect:
                     else None
                 ),
                 grain=cte.grain,
-                limit=None,
+                limit=cte.limit,
                 # some joins may not need to be rendered
                 joins=[
                     j
@@ -475,9 +475,11 @@ class BaseDialect:
                 where=(
                     self.render_expr(cte.condition, cte) if cte.condition else None
                 ),  # source_map=cte_output_map)
-                # where=self.render_expr(where_assignment[cte.name], cte)
-                # if cte.name in where_assignment
-                # else None,
+                order_by=(
+                    [self.render_order_item(i, cte) for i in cte.order_by.items]
+                    if cte.order_by
+                    else None
+                ),
                 group_by=(
                     list(
                         set(
@@ -522,7 +524,8 @@ class BaseDialect:
         )
 
     def generate_ctes(
-        self, query: ProcessedQuery, where_assignment: Dict[str, Conditional]
+        self,
+        query: ProcessedQuery,
     ):
         return [self.render_cte(cte) for cte in query.ctes]
 
@@ -649,35 +652,54 @@ class BaseDialect:
                     " filtered concept instead."
                 )
 
-        compiled_ctes = self.generate_ctes(query, {})
+        compiled_ctes = self.generate_ctes(query)
 
         # restort selections by the order they were written in
         sorted_select: List[str] = []
         for output_c in output_addresses:
             sorted_select.append(select_columns[output_c])
-        final = self.SQL_TEMPLATE.render(
-            output=(
-                query.output_to if isinstance(query, ProcessedQueryPersist) else None
-            ),
-            select_columns=sorted_select,
-            base=query.base.name,
-            joins=[
-                render_join(join, self.QUOTE_CHARACTER, None) for join in query.joins
-            ],
-            ctes=compiled_ctes,
-            limit=query.limit,
-            # move up to CTEs
-            where=(
-                self.render_expr(query.where_clause.conditional, cte_map=cte_output_map)
-                if query.where_clause and output_where
-                else None
-            ),
-            order_by=(
-                [self.render_order_item(i, [query.base]) for i in query.order_by.items]
-                if query.order_by
-                else None
-            ),
-        )
+        if not query.base.requires_nesting:
+            final = self.SQL_TEMPLATE.render(
+                output=(
+                    query.output_to
+                    if isinstance(query, ProcessedQueryPersist)
+                    else None
+                ),
+                full_select=compiled_ctes[-1].statement,
+                ctes=compiled_ctes[:-1],
+            )
+        else:
+            final = self.SQL_TEMPLATE.render(
+                output=(
+                    query.output_to
+                    if isinstance(query, ProcessedQueryPersist)
+                    else None
+                ),
+                select_columns=sorted_select,
+                base=query.base.name,
+                joins=[
+                    render_join(join, self.QUOTE_CHARACTER, None)
+                    for join in query.joins
+                ],
+                ctes=compiled_ctes,
+                limit=query.limit,
+                # move up to CTEs
+                where=(
+                    self.render_expr(
+                        query.where_clause.conditional, cte_map=cte_output_map
+                    )
+                    if query.where_clause and output_where
+                    else None
+                ),
+                order_by=(
+                    [
+                        self.render_order_item(i, query.base, final=True)
+                        for i in query.order_by.items
+                    ]
+                    if query.order_by
+                    else None
+                ),
+            )
 
         if CONFIG.strict_mode and INVALID_REFERENCE_STRING(1) in final:
             raise ValueError(

trilogy/dialect/bigquery.py

@@ -43,8 +43,11 @@ CREATE OR REPLACE TABLE {{ output.address.location }} AS
 {% endif %}{%- if ctes %}
 WITH {% for cte in ctes %}
 {{cte.name}} as ({{cte.statement}}){% if not loop.last %},{% endif %}{% endfor %}{% endif %}
-SELECT
+{%- if full_select -%}
+{{full_select}}
+{%- else -%}
 
+SELECT
 {%- for select in select_columns %}
     {{ select }}{% if not loop.last %},{% endif %}{% endfor %}
 {% if base %}FROM
@@ -59,10 +62,9 @@ SELECT
     {{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 %}
+    {{ order }}{% if not loop.last %},{% endif %}{% endfor %}{% endif %}
 {%- if limit is not none %}
-LIMIT {{ limit }}{% endif %}
+LIMIT {{ limit }}{% endif %}{% endif %}
 """
 )
 MAX_IDENTIFIER_LENGTH = 50

trilogy/dialect/common.py

@@ -1,4 +1,4 @@
-from trilogy.core.models import Join, InstantiatedUnnestJoin, CTE, Concept
+from trilogy.core.models import Join, InstantiatedUnnestJoin, CTE, Concept, Datasource
 from trilogy.core.enums import UnnestMode, Modifier
 from typing import Optional, Callable
 
@@ -21,18 +21,20 @@ def render_join(
         if unnest_mode == UnnestMode.DIRECT:
             return None
         if not render_func:
-            raise ValueError("must provide a render func to build an unnest joins")
+            raise ValueError("must provide a render function to build an unnest joins")
         if not cte:
             raise ValueError("must provide a cte to build an unnest joins")
         if unnest_mode == UnnestMode.CROSS_JOIN:
             return f"CROSS JOIN {render_func(join.concept, cte, False)} as {quote_character}{join.concept.safe_address}{quote_character}"
 
         return f"FULL JOIN {render_func(join.concept, cte, False)} as unnest_wrapper({quote_character}{join.concept.safe_address}{quote_character})"
-
+    left_name = join.left_name
+    right_name = join.right_name
+    right_base = join.right_ref
     base_joinkeys = [
         null_wrapper(
-            f"{join.left_cte.name}.{quote_character}{key.concept.safe_address}{quote_character}",
-            f"{join.right_cte.name}.{quote_character}{key.concept.safe_address}{quote_character}",
+            f"{left_name}.{quote_character}{join.left_cte.get_alias(key.concept) if isinstance(join.left_cte, Datasource) else key.concept.safe_address}{quote_character}",
+            f"{right_name}.{quote_character}{join.right_cte.get_alias(key.concept) if isinstance(join.right_cte, Datasource) else key.concept.safe_address}{quote_character}",
             key.concept,
         )
         for key in join.joinkeys
@@ -40,4 +42,4 @@ def render_join(
     if not base_joinkeys:
         base_joinkeys = ["1=1"]
     joinkeys = " AND ".join(base_joinkeys)
-    return f"{join.jointype.value.upper()} JOIN {join.right_cte.name} on {joinkeys}"
+    return f"{join.jointype.value.upper()} JOIN {right_base} on {joinkeys}"

trilogy/dialect/config.py

@@ -1,5 +1,27 @@
 class DialectConfig:
-    pass
+
+    def __init__(self):
+        pass
+
+    def connection_string(self) -> str:
+        raise NotImplementedError
+
+    @property
+    def connect_args(self) -> dict:
+        return {}
+
+
+class BigQueryConfig(DialectConfig):
+    def __init__(self, project: str, client):
+        self.project = project
+        self.client = client
+
+    def connection_string(self) -> str:
+        return f"bigquery://{self.project}?user_supplied_client=True"
+
+    @property
+    def connect_args(self) -> dict:
+        return {"client": self.client}
 
 
 class DuckDBConfig(DialectConfig):
@@ -53,3 +75,49 @@ class SnowflakeConfig(DialectConfig):
 
     def connection_string(self) -> str:
         return f"snowflake://{self.username}:{self.password}@{self.account}"
+
+
+class PrestoConfig(DialectConfig):
+    def __init__(
+        self,
+        host: str,
+        port: int,
+        username: str,
+        password: str,
+        catalog: str,
+        schema: str | None = None,
+    ):
+        self.host = host
+        self.port = port
+        self.username = username
+        self.password = password
+        self.catalog = catalog
+        self.schema = schema
+
+    def connection_string(self) -> str:
+        if self.schema:
+            return f"presto://{self.username}:{self.password}@{self.host}:{self.port}/{self.catalog}/{self.schema}"
+        return f"presto://{self.username}:{self.password}@{self.host}:{self.port}/{self.catalog}"
+
+
+class TrinoConfig(DialectConfig):
+    def __init__(
+        self,
+        host: str,
+        port: int,
+        username: str,
+        password: str,
+        catalog: str,
+        schema: str | None = None,
+    ):
+        self.host = host
+        self.port = port
+        self.username = username
+        self.password = password
+        self.catalog = catalog
+        self.schema = schema
+
+    def connection_string(self) -> str:
+        if self.schema:
+            return f"trino://{self.username}:{self.password}@{self.host}:{self.port}/{self.catalog}/{self.schema}"
+        return f"trino://{self.username}:{self.password}@{self.host}:{self.port}/{self.catalog}"

trilogy/dialect/duckdb.py

@@ -47,8 +47,10 @@ CREATE OR REPLACE TABLE {{ output.address.location }} AS
 {% endif %}{%- if ctes %}
 WITH {% for cte in ctes %}
 {{cte.name}} as ({{cte.statement}}){% if not loop.last %},{% endif %}{% endfor %}{% endif %}
-SELECT
+{% if full_select -%}{{full_select}}
+{% else -%}
 
+SELECT
 {%- for select in select_columns %}
     {{ select }}{% if not loop.last %},{% endif %}{% endfor %}
 {% if base %}FROM
@@ -63,10 +65,9 @@ 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 %}
+    {{ order }}{% if not loop.last %},{% endif %}{% endfor %}{% endif %}
 {%- if limit is not none %}
-LIMIT ({{ limit }}){% endif %}
+LIMIT ({{ limit }}){% endif %}{% endif %}
 """
 )
 

trilogy/dialect/enums.py

@@ -1,5 +1,5 @@
 from enum import Enum
-from typing import List, TYPE_CHECKING, Optional
+from typing import List, TYPE_CHECKING, Optional, Callable
 
 if TYPE_CHECKING:
     from trilogy.hooks.base_hook import BaseHook
@@ -9,6 +9,20 @@ from trilogy.dialect.config import DialectConfig
 from trilogy.constants import logger
 
 
+def default_factory(conf: DialectConfig, config_type):
+    from sqlalchemy import create_engine
+
+    if not isinstance(conf, config_type):
+        raise TypeError(
+            f"Invalid dialect configuration for type {type(config_type).__name__}"
+        )
+    if conf.connect_args:
+        return create_engine(
+            conf.connection_string(), future=True, connect_args=conf.connect_args
+        )
+    return create_engine(conf.connection_string(), future=True)
+
+
 class Dialects(Enum):
     BIGQUERY = "bigquery"
     SQL_SERVER = "sql_server"
@@ -24,38 +38,32 @@ class Dialects(Enum):
             return cls.DUCK_DB
         return super()._missing_(value)
 
-    def default_engine(self, conf=None):
+    def default_engine(self, conf=None, _engine_factory: Callable = default_factory):
         if self == Dialects.BIGQUERY:
-            from sqlalchemy import create_engine
             from google.auth import default
             from google.cloud import bigquery
+            from trilogy.dialect.config import BigQueryConfig
 
             credentials, project = default()
             client = bigquery.Client(credentials=credentials, project=project)
-            return create_engine(
-                f"bigquery://{project}?user_supplied_client=True",
-                connect_args={"client": client},
+            conf = conf or BigQueryConfig(project=project, client=client)
+            return _engine_factory(
+                conf,
+                BigQueryConfig,
             )
         elif self == Dialects.SQL_SERVER:
-            from sqlalchemy import create_engine
 
             raise NotImplementedError()
         elif self == Dialects.DUCK_DB:
-            from sqlalchemy import create_engine
             from trilogy.dialect.config import DuckDBConfig
 
             if not conf:
                 conf = DuckDBConfig()
-            if not isinstance(conf, DuckDBConfig):
-                raise TypeError("Invalid dialect configuration for type duck_db")
-            return create_engine(conf.connection_string(), future=True)
+            return _engine_factory(conf, DuckDBConfig)
         elif self == Dialects.SNOWFLAKE:
-            from sqlalchemy import create_engine
             from trilogy.dialect.config import SnowflakeConfig
 
-            if not isinstance(conf, SnowflakeConfig):
-                raise TypeError("Invalid dialect configuration for type snowflake")
-            return create_engine(conf.connection_string(), future=True)
+            return _engine_factory(conf, SnowflakeConfig)
         elif self == Dialects.POSTGRES:
             logger.warn(
                 "WARN: Using experimental postgres dialect. Most functionality will not work."
@@ -67,13 +75,17 @@ class Dialects(Enum):
                 raise ImportError(
                     "postgres driver not installed. python -m pip install pypreql[postgres]"
                 )
-            from sqlalchemy import create_engine
             from trilogy.dialect.config import PostgresConfig
 
-            if not isinstance(conf, PostgresConfig):
-                raise TypeError("Invalid dialect configuration for type postgres")
+            return _engine_factory(conf, PostgresConfig)
+        elif self == Dialects.PRESTO:
+            from trilogy.dialect.config import PrestoConfig
+
+            return _engine_factory(conf, PrestoConfig)
+        elif self == Dialects.TRINO:
+            from trilogy.dialect.config import TrinoConfig
 
-            return create_engine(conf.connection_string(), future=True)
+            return _engine_factory(conf, TrinoConfig)
         else:
             raise ValueError(
                 f"Unsupported dialect {self} for default engine creation; create one explicitly."
@@ -84,9 +96,18 @@ class Dialects(Enum):
         environment: Optional["Environment"] = None,
         hooks: List["BaseHook"] | None = None,
         conf: DialectConfig | None = None,
+        _engine_factory: Callable | None = None,
     ) -> "Executor":
         from trilogy import Executor, Environment
 
+        if _engine_factory is not None:
+            return Executor(
+                engine=self.default_engine(conf=conf, _engine_factory=_engine_factory),
+                environment=environment or Environment(),
+                dialect=self,
+                hooks=hooks,
+            )
+
         return Executor(
             engine=self.default_engine(conf=conf),
             environment=environment or Environment(),

trilogy/dialect/postgres.py

@@ -49,8 +49,10 @@ CREATE TABLE {{ output.address.location }} AS
 {% endif %}{%- if ctes %}
 WITH {% for cte in ctes %}
 {{cte.name}} as ({{cte.statement}}){% if not loop.last %},{% endif %}{% endfor %}{% endif %}
+{%- if full_select -%}
+{{full_select}}
+{%- else -%}
 SELECT
-
 {%- for select in select_columns %}
     {{ select }}{% if not loop.last %},{% endif %}{% endfor %}
 {% if base %}FROM
@@ -68,7 +70,7 @@ ORDER BY {% for order in order_by %}
     {{ order }}{% if not loop.last %},{% endif %}
 {% endfor %}{% endif %}
 {%- if limit is not none %}
-LIMIT {{ limit }}{% endif %}
+LIMIT {{ limit }}{% endif %}{% endif %}
 """
 )
 

trilogy/dialect/presto.py

@@ -42,8 +42,11 @@ CREATE OR REPLACE TABLE {{ output.address }} AS
 {% endif %}{%- if ctes %}
 WITH {% for cte in ctes %}
 {{cte.name}} as ({{cte.statement}}){% if not loop.last %},{% endif %}{% endfor %}{% endif %}
-SELECT
 
+SELECT
+{%- if full_select -%}
+{{full_select}}
+{%- else -%}
 {%- for select in select_columns %}
     {{ select }}{% if not loop.last %},{% endif %}{% endfor %}
 {% if base %}FROM
@@ -58,10 +61,9 @@ SELECT
     {{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 %}
+    {{ order }}{% if not loop.last %},{% endif %}{% endfor %}{% endif %}
 {%- if limit is not none %}
-LIMIT {{ limit }}{% endif %}
+LIMIT {{ limit }}{% endif %}{% endif %}
 """
 )
 MAX_IDENTIFIER_LENGTH = 50

trilogy/dialect/snowflake.py

@@ -45,8 +45,11 @@ CREATE OR REPLACE TABLE {{ output.address.location }} AS
 {% endif %}{%- if ctes %}
 WITH {% for cte in ctes %}
 {{cte.name}} as ({{cte.statement}}){% if not loop.last %},{% endif %}{% endfor %}{% endif %}
-SELECT
+{%- if full_select -%}
+{{full_select}}
+{%- else -%}
 
+SELECT
 {%- for select in select_columns %}
     {{ select }}{% if not loop.last %},{% endif %}{% endfor %}
 {% if base %}FROM
@@ -61,10 +64,9 @@ SELECT
     {{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 %}
+    {{ order }}{% if not loop.last %},{% endif %}{% endfor %}{% endif %}
 {%- if limit is not none %}
-LIMIT {{ limit }}{% endif %}
+LIMIT {{ limit }}{% endif %}{% endif %}
 """
 )
 MAX_IDENTIFIER_LENGTH = 50

trilogy/dialect/sql_server.py

@@ -40,6 +40,9 @@ TSQL_TEMPLATE = Template(
     """{%- if ctes %}
 WITH {% for cte in ctes %}
 {{cte.name}} as ({{cte.statement}}){% if not loop.last %},{% endif %}{% endfor %}{% endif %}
+{%- if full_select -%}
+{{full_select}}
+{%- else -%}
 SELECT
 {%- if limit is not none %}
 TOP {{ limit }}{% endif %}
@@ -60,7 +63,7 @@ GROUP BY {% for group in group_by %}
 {%- if order_by %}
 ORDER BY {% for order in order_by %}
     {{ order }}{% if not loop.last %},{% endif %}
-{% endfor %}{% endif %}
+{% endfor %}{% endif %}{% endif %}
 """
 )
 

trilogy/executor.py

@@ -99,7 +99,14 @@ class Executor(object):
         raise NotImplementedError("Cannot execute type {}".format(type(query)))
 
     @execute_query.register
-    def _(self, query: SelectStatement | PersistStatement) -> CursorResult:
+    def _(self, query: SelectStatement) -> CursorResult:
+        sql = self.generator.generate_queries(
+            self.environment, [query], hooks=self.hooks
+        )
+        return self.execute_query(sql[0])
+
+    @execute_query.register
+    def _(self, query: PersistStatement) -> CursorResult:
         sql = self.generator.generate_queries(
             self.environment, [query], hooks=self.hooks
         )
@@ -117,16 +124,22 @@ class Executor(object):
         )
 
     @execute_query.register
-    def _(self, query: ProcessedQuery | ProcessedQueryPersist) -> CursorResult:
+    def _(self, query: ProcessedQuery) -> CursorResult:
+        sql = self.generator.compile_statement(query)
+        # connection = self.engine.connect()
+        output = self.connection.execute(text(sql))
+        return output
+
+    @execute_query.register
+    def _(self, query: 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)
+        self.environment.add_datasource(query.datasource)
         return output
 
     @singledispatchmethod
-    def generate_sql(self, command: ProcessedQuery | str) -> list[str]:
+    def generate_sql(self, command) -> list[str]:
         raise NotImplementedError(
             "Cannot generate sql for type {}".format(type(command))
         )

trilogy/parsing/parse_engine.py

@@ -864,7 +864,7 @@ class ParseToObjects(Transformer):
             elif isinstance(val, Grain):
                 grain = val
             elif isinstance(val, Query):
-                address = Address(location=f"({val.text})")
+                address = Address(location=f"({val.text})", is_query=True)
         if not address:
             raise ValueError(
                 "Malformed datasource, missing address or query declaration"