diff --git a/superset/db_engine_specs/base.py b/superset/db_engine_specs/base.py index f355e4ef8cea8..542e73739a636 100644 --- a/superset/db_engine_specs/base.py +++ b/superset/db_engine_specs/base.py @@ -398,6 +398,19 @@ class BaseEngineSpec: # pylint: disable=too-many-public-methods # Can the catalog be changed on a per-query basis? supports_dynamic_catalog = False + @classmethod + def get_allows_alias_in_select( + cls, database: Database + ) -> bool: # pylint: disable=unused-argument + """ + Method for dynamic `allows_alias_in_select`. + + In Dremio this atribute is version-dependent, so Superset needs to inspect the + database configuration in order to determine it. This method allows engine-specs + to define dynamic values for the attribute. + """ + return cls.allows_alias_in_select + @classmethod def supports_url(cls, url: URL) -> bool: """ diff --git a/superset/db_engine_specs/dremio.py b/superset/db_engine_specs/dremio.py index c96159f1b8aa4..746576d3f30da 100644 --- a/superset/db_engine_specs/dremio.py +++ b/superset/db_engine_specs/dremio.py @@ -14,14 +14,25 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. + +from __future__ import annotations + from datetime import datetime -from typing import Any, Optional +from typing import Any, TYPE_CHECKING +from packaging.version import Version from sqlalchemy import types from superset.constants import TimeGrain from superset.db_engine_specs.base import BaseEngineSpec +if TYPE_CHECKING: + from superset.models.core import Database + + +# See https://github.com/apache/superset/pull/25657 +FIXED_ALIAS_IN_SELECT_VERSION = Version("24.1.0") + class DremioEngineSpec(BaseEngineSpec): engine = "dremio" @@ -43,10 +54,25 @@ class DremioEngineSpec(BaseEngineSpec): def epoch_to_dttm(cls) -> str: return "TO_DATE({col})" + @classmethod + def get_allows_alias_in_select(cls, database: Database) -> bool: + """ + Dremio supports aliases in SELECT statements since version 24.1.0. + + If no version is specified in the DB extra, we assume the Dremio version is post + 24.1.0. This way, as we move forward people don't have to specify a version when + setting up their databases. + """ + version = database.get_extra().get("version") + if version and Version(version) < FIXED_ALIAS_IN_SELECT_VERSION: + return False + + return True + @classmethod def convert_dttm( - cls, target_type: str, dttm: datetime, db_extra: Optional[dict[str, Any]] = None - ) -> Optional[str]: + cls, target_type: str, dttm: datetime, db_extra: dict[str, Any] | None = None + ) -> str | None: sqla_type = cls.get_sqla_column_type(target_type) if isinstance(sqla_type, types.Date): diff --git a/superset/models/core.py b/superset/models/core.py index f6e4b972b48db..ab3600b2be5aa 100755 --- a/superset/models/core.py +++ b/superset/models/core.py @@ -965,7 +965,7 @@ def make_sqla_column_compatible( """ label_expected = label or sqla_col.name # add quotes to tables - if self.db_engine_spec.allows_alias_in_select: + if self.db_engine_spec.get_allows_alias_in_select(self): label = self.db_engine_spec.make_label_compatible(label_expected) sqla_col = sqla_col.label(label) sqla_col.key = label_expected diff --git a/superset/models/helpers.py b/superset/models/helpers.py index 83ec9ba37c4fa..cfa09154ed748 100644 --- a/superset/models/helpers.py +++ b/superset/models/helpers.py @@ -765,7 +765,7 @@ def db_engine_spec(self) -> builtins.type["BaseEngineSpec"]: raise NotImplementedError() @property - def database(self) -> builtins.type["Database"]: + def database(self) -> "Database": raise NotImplementedError() @property @@ -865,7 +865,7 @@ def make_sqla_column_compatible( label_expected = label or sqla_col.name db_engine_spec = self.db_engine_spec # add quotes to tables - if db_engine_spec.allows_alias_in_select: + if db_engine_spec.get_allows_alias_in_select(self.database): label = db_engine_spec.make_label_compatible(label_expected) sqla_col = sqla_col.label(label) sqla_col.key = label_expected @@ -900,7 +900,7 @@ def get_query_str_extended( self, query_obj: QueryObjectDict, mutate: bool = True ) -> QueryStringExtended: sqlaq = self.get_sqla_query(**query_obj) - sql = self.database.compile_sqla_query(sqlaq.sqla_query) # type: ignore + sql = self.database.compile_sqla_query(sqlaq.sqla_query) sql = self._apply_cte(sql, sqlaq.cte) sql = sqlparse.format(sql, reindent=True) if mutate: @@ -939,7 +939,7 @@ def _normalize_prequery_result_type( value = value.item() column_ = columns_by_name[dimension] - db_extra: dict[str, Any] = self.database.get_extra() # type: ignore + db_extra: dict[str, Any] = self.database.get_extra() if isinstance(column_, dict): if ( @@ -1024,9 +1024,7 @@ def assign_column_label(df: pd.DataFrame) -> Optional[pd.DataFrame]: return df try: - df = self.database.get_df( - sql, self.schema, mutator=assign_column_label # type: ignore - ) + df = self.database.get_df(sql, self.schema, mutator=assign_column_label) except Exception as ex: # pylint: disable=broad-except df = pd.DataFrame() status = QueryStatus.FAILED @@ -1361,7 +1359,7 @@ def values_for_column(self, column_name: str, limit: int = 10000) -> list[Any]: if limit: qry = qry.limit(limit) - with self.database.get_sqla_engine_with_context() as engine: # type: ignore + with self.database.get_sqla_engine_with_context() as engine: sql = qry.compile(engine, compile_kwargs={"literal_binds": True}) sql = self._apply_cte(sql, cte) sql = self.mutate_query_from_config(sql) @@ -1958,7 +1956,7 @@ def get_sqla_query( # pylint: disable=too-many-arguments,too-many-locals,too-ma col = col.element if ( - db_engine_spec.allows_alias_in_select + db_engine_spec.get_allows_alias_in_select(self.database) and db_engine_spec.allows_hidden_cc_in_orderby and col.name in [select_col.name for select_col in select_exprs] ):