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

fix(bigquery): calculated column cannot orderby in BigQuery #17196

Merged
merged 4 commits into from
Oct 22, 2021
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.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
7 changes: 7 additions & 0 deletions superset/connectors/sqla/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -1300,6 +1300,13 @@ def get_sqla_query( # pylint: disable=too-many-arguments,too-many-locals,too-ma
# if engine does not allow using SELECT alias in ORDER BY
# revert to the underlying column
col = col.element

if (
db_engine_spec.allows_alias_in_select
and db_engine_spec.allows_hidden_cc_in_orderby
and col.name in [select_col.name for select_col in select_exprs]
):
col = literal_column(col.name)
direction = asc if ascending else desc
qry = qry.order_by(direction(col))

Expand Down
6 changes: 6 additions & 0 deletions superset/db_engine_specs/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -289,6 +289,12 @@ class BaseEngineSpec: # pylint: disable=too-many-public-methods
# if TRUE, then it doesn't have to.
allows_hidden_ordeby_agg = True

# Whether ORDER BY clause can use sql caculated expression
# if True, use alias of select column for `order by`
# the True is safely for most database
# But for backward compatibility, False by default
allows_hidden_cc_in_orderby = False

force_column_alias_quotes = False
arraysize = 0
max_column_name_length = 0
Expand Down
2 changes: 2 additions & 0 deletions superset/db_engine_specs/bigquery.py
Original file line number Diff line number Diff line change
Expand Up @@ -99,6 +99,8 @@ class BigQueryEngineSpec(BaseEngineSpec):
# same cursor, so we need to run all statements at once
run_multiple_statements_as_one = True

allows_hidden_cc_in_orderby = True

"""
https://www.python.org/dev/peps/pep-0249/#arraysize
raw_connections bypass the pybigquery query execution context and deal with
Expand Down
34 changes: 34 additions & 0 deletions tests/integration_tests/db_engine_specs/base_engine_spec_tests.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@

import pytest

from superset.connectors.sqla.models import TableColumn
from superset.db_engine_specs import get_engine_specs
from superset.db_engine_specs.base import (
BaseEngineSpec,
Expand All @@ -34,6 +35,7 @@
from tests.integration_tests.db_engine_specs.base_tests import TestDbEngineSpec
from tests.integration_tests.test_app import app

from ..fixtures.birth_names_dashboard import load_birth_names_dashboard_with_slices
from ..fixtures.energy_dashboard import load_energy_table_with_slice
from ..fixtures.pyodbcRow import Row

Expand Down Expand Up @@ -273,6 +275,38 @@ def test_pyodbc_rows_to_tuples_passthrough(self):
result = BaseEngineSpec.pyodbc_rows_to_tuples(data)
self.assertListEqual(result, data)

@mock.patch("superset.models.core.Database.db_engine_spec", BaseEngineSpec)
@pytest.mark.usefixtures("load_birth_names_dashboard_with_slices")
def test_calculated_column_in_order_by_base_engine_spec(self):
table = self.get_table(name="birth_names")
TableColumn(
column_name="gender_cc",
type="VARCHAR(255)",
table=table,
expression="""
case
when gender=true then "male"
else "female"
end
""",
)

table.database.sqlalchemy_uri = "sqlite://"
query_obj = {
"groupby": ["gender_cc"],
"is_timeseries": False,
"filter": [],
"orderby": [["gender_cc", True]],
}
sql = table.get_query_str(query_obj)
assert (
"""ORDER BY case
when gender=true then "male"
else "female"
end ASC;"""
in sql
)


def test_is_readonly():
def is_readonly(sql: str) -> bool:
Expand Down
32 changes: 32 additions & 0 deletions tests/integration_tests/db_engine_specs/bigquery_tests.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,14 +17,19 @@
import sys
import unittest.mock as mock

import pytest
from pandas import DataFrame
from sqlalchemy import column

from superset.connectors.sqla.models import TableColumn
from superset.db_engine_specs.base import BaseEngineSpec
from superset.db_engine_specs.bigquery import BigQueryEngineSpec
from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
from superset.sql_parse import Table
from tests.integration_tests.db_engine_specs.base_tests import TestDbEngineSpec
from tests.integration_tests.fixtures.birth_names_dashboard import (
load_birth_names_dashboard_with_slices,
)


class TestBigQueryDbEngineSpec(TestDbEngineSpec):
Expand Down Expand Up @@ -333,3 +338,30 @@ def test_extract_errors(self):
},
)
]

@mock.patch("superset.models.core.Database.db_engine_spec", BigQueryEngineSpec)
@mock.patch("pybigquery._helpers.create_bigquery_client", mock.Mock)
@pytest.mark.usefixtures("load_birth_names_dashboard_with_slices")
def test_calculated_column_in_order_by(self):
table = self.get_table(name="birth_names")
TableColumn(
column_name="gender_cc",
type="VARCHAR(255)",
table=table,
expression="""
case
when gender=true then "male"
else "female"
end
""",
)

table.database.sqlalchemy_uri = "bigquery://"
query_obj = {
"groupby": ["gender_cc"],
"is_timeseries": False,
"filter": [],
"orderby": [["gender_cc", True]],
}
sql = table.get_query_str(query_obj)
assert "ORDER BY gender_cc ASC" in sql