Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: method for dynamic allows_alias_in_select #25882

Merged
merged 2 commits into from
Nov 7, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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 # pylint: disable=unused-argument
) -> bool:
"""
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 @@
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):
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 @@
raise NotImplementedError()

@property
def database(self) -> builtins.type["Database"]:
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This was wrong, but was hidden by a bunch of # type: ignore comments. ☹️

def database(self) -> "Database":
raise NotImplementedError()

@property
Expand Down Expand Up @@ -865,7 +865,7 @@
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 @@
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 @@
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 @@
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 @@
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 @@
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
16 changes: 16 additions & 0 deletions tests/unit_tests/db_engine_specs/test_dremio.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
from typing import Optional

import pytest
from pytest_mock import MockerFixture

from tests.unit_tests.db_engine_specs.utils import assert_convert_dttm
from tests.unit_tests.fixtures.common import dttm
Expand All @@ -40,3 +41,18 @@ def test_convert_dttm(
from superset.db_engine_specs.dremio import DremioEngineSpec as spec

assert_convert_dttm(spec, target_type, expected_result, dttm)


def test_get_allows_alias_in_select(mocker: MockerFixture) -> None:
from superset.db_engine_specs.dremio import DremioEngineSpec

database = mocker.MagicMock()

database.get_extra.return_value = {}
assert DremioEngineSpec.get_allows_alias_in_select(database) is True

database.get_extra.return_value = {"version": "24.1.0"}
assert DremioEngineSpec.get_allows_alias_in_select(database) is True

database.get_extra.return_value = {"version": "24.0.0"}
assert DremioEngineSpec.get_allows_alias_in_select(database) is False
Loading