Skip to content

Commit

Permalink
feat: method for dynamic allows_alias_in_select
Browse files Browse the repository at this point in the history
  • Loading branch information
betodealmeida committed Nov 6, 2023
1 parent d6fde3c commit 491a2ce
Show file tree
Hide file tree
Showing 4 changed files with 50 additions and 13 deletions.
13 changes: 13 additions & 0 deletions superset/db_engine_specs/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
"""
Expand Down
32 changes: 29 additions & 3 deletions superset/db_engine_specs/dremio.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Check warning on line 30 in superset/db_engine_specs/dremio.py

View check run for this annotation

Codecov / codecov/patch

superset/db_engine_specs/dremio.py#L30

Added line #L30 was not covered by tests


# See https://github.com/apache/superset/pull/25657
FIXED_ALIAS_IN_SELECT_VERSION = Version("24.1.0")


class DremioEngineSpec(BaseEngineSpec):
engine = "dremio"
Expand All @@ -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

Check warning on line 68 in superset/db_engine_specs/dremio.py

View check run for this annotation

Codecov / codecov/patch

superset/db_engine_specs/dremio.py#L66-L68

Added lines #L66 - L68 were not covered by tests

return True

Check warning on line 70 in superset/db_engine_specs/dremio.py

View check run for this annotation

Codecov / codecov/patch

superset/db_engine_specs/dremio.py#L70

Added line #L70 was not covered by tests

@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):
Expand Down
2 changes: 1 addition & 1 deletion superset/models/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
16 changes: 7 additions & 9 deletions superset/models/helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)

Check warning on line 903 in superset/models/helpers.py

View check run for this annotation

Codecov / codecov/patch

superset/models/helpers.py#L903

Added line #L903 was not covered by tests
sql = self._apply_cte(sql, sqlaq.cte)
sql = sqlparse.format(sql, reindent=True)
if mutate:
Expand Down Expand Up @@ -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()

Check warning on line 942 in superset/models/helpers.py

View check run for this annotation

Codecov / codecov/patch

superset/models/helpers.py#L942

Added line #L942 was not covered by tests

if isinstance(column_, dict):
if (
Expand Down Expand Up @@ -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)

Check warning on line 1027 in superset/models/helpers.py

View check run for this annotation

Codecov / codecov/patch

superset/models/helpers.py#L1027

Added line #L1027 was not covered by tests
except Exception as ex: # pylint: disable=broad-except
df = pd.DataFrame()
status = QueryStatus.FAILED
Expand Down Expand Up @@ -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:

Check warning on line 1362 in superset/models/helpers.py

View check run for this annotation

Codecov / codecov/patch

superset/models/helpers.py#L1362

Added line #L1362 was not covered by tests
sql = qry.compile(engine, compile_kwargs={"literal_binds": True})
sql = self._apply_cte(sql, cte)
sql = self.mutate_query_from_config(sql)
Expand Down Expand Up @@ -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]
):
Expand Down

0 comments on commit 491a2ce

Please sign in to comment.