fix: Make cte_alias a property of db engine spec (#22947)

apache · Feb 6, 2023 · 9dfaad7 · 9dfaad7
1 parent aa0a078
commit 9dfaad7
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 8 deletions.
diff --git a/superset/connectors/sqla/models.py b/superset/connectors/sqla/models.py
@@ -96,7 +96,7 @@
     validate_adhoc_subquery,
 )
 from superset.datasets.models import Dataset as NewDataset
-from superset.db_engine_specs.base import BaseEngineSpec, CTE_ALIAS, TimestampExpression
+from superset.db_engine_specs.base import BaseEngineSpec, TimestampExpression
 from superset.exceptions import (
     AdvancedDataTypeResponseError,
     DatasetInvalidPermissionEvaluationException,
@@ -921,7 +921,7 @@ def get_from_clause(
 
         cte = self.db_engine_spec.get_cte_query(from_sql)
         from_clause = (
-            table(CTE_ALIAS)
+            table(self.db_engine_spec.cte_alias)
             if cte
             else TextAsFrom(self.text(from_sql), []).alias(VIRTUAL_TABLE_ALIAS)
         )

diff --git a/superset/db_engine_specs/base.py b/superset/db_engine_specs/base.py
@@ -86,9 +86,6 @@
 logger = logging.getLogger()
 
 
-CTE_ALIAS = "__cte"
-
-
 class TimeGrain(NamedTuple):
     name: str  # TODO: redundant field, remove
     label: str
@@ -346,6 +343,8 @@ class BaseEngineSpec:  # pylint: disable=too-many-public-methods
     # If True, then it will allow  in subquery ,
     # if False it will allow as regular CTE
     allows_cte_in_subquery = True
+    # Define alias for CTE
+    cte_alias = "__cte"
     # Whether allow LIMIT clause in the SQL
     # If True, then the database engine is allowed for LIMIT clause
     # If False, then the database engine is allowed for TOP clause
@@ -889,7 +888,7 @@ def get_cte_query(cls, sql: str) -> Optional[str]:
 
             # extract rest of the SQLs after CTE
             remainder = "".join(str(token) for token in stmt.tokens[idx:]).strip()
-            return f"WITH {token.value},\n{CTE_ALIAS} AS (\n{remainder}\n)"
+            return f"WITH {token.value},\n{cls.cte_alias} AS (\n{remainder}\n)"
 
         return None
 

diff --git a/superset/models/helpers.py b/superset/models/helpers.py
@@ -98,7 +98,6 @@
 config = app.config
 logger = logging.getLogger(__name__)
 
-CTE_ALIAS = "__cte"
 VIRTUAL_TABLE_ALIAS = "virtual_table"
 ADVANCED_DATA_TYPES = config["ADVANCED_DATA_TYPES"]
 
@@ -1049,7 +1048,7 @@ def get_from_clause(
 
         cte = self.db_engine_spec.get_cte_query(from_sql)
         from_clause = (
-            sa.table(CTE_ALIAS)
+            sa.table(self.db_engine_spec.cte_alias)
             if cte
             else TextAsFrom(self.text(from_sql), []).alias(VIRTUAL_TABLE_ALIAS)
         )