From 9a3d729ec4b2856617306d493da76fac7d2ed2fb Mon Sep 17 00:00:00 2001 From: Yongjie Zhao Date: Fri, 22 Oct 2021 16:23:24 +0800 Subject: [PATCH 1/4] fix(bigquery): calculated column cannot orderby in BigQuery --- superset/connectors/sqla/models.py | 7 +++++++ superset/db_engine_specs/base.py | 6 ++++++ superset/db_engine_specs/bigquery.py | 2 ++ 3 files changed, 15 insertions(+) diff --git a/superset/connectors/sqla/models.py b/superset/connectors/sqla/models.py index a080507075325..5ed1c1876cecb 100644 --- a/superset/connectors/sqla/models.py +++ b/superset/connectors/sqla/models.py @@ -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_orderby_calculated_column + 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)) diff --git a/superset/db_engine_specs/base.py b/superset/db_engine_specs/base.py index ed3851efe6ce7..f9c203211b862 100644 --- a/superset/db_engine_specs/base.py +++ b/superset/db_engine_specs/base.py @@ -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 safe for most database + # But for backward compatibility, set to False by default + allows_hidden_orderby_calculated_column = False + force_column_alias_quotes = False arraysize = 0 max_column_name_length = 0 diff --git a/superset/db_engine_specs/bigquery.py b/superset/db_engine_specs/bigquery.py index 5e1ff739cd2c8..70057de544446 100644 --- a/superset/db_engine_specs/bigquery.py +++ b/superset/db_engine_specs/bigquery.py @@ -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_orderby_calculated_column = True + """ https://www.python.org/dev/peps/pep-0249/#arraysize raw_connections bypass the pybigquery query execution context and deal with From d5afc18b8d0ecc2868c036fc56be51d210f50134 Mon Sep 17 00:00:00 2001 From: Yongjie Zhao Date: Fri, 22 Oct 2021 16:27:33 +0800 Subject: [PATCH 2/4] typo --- superset/db_engine_specs/base.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/superset/db_engine_specs/base.py b/superset/db_engine_specs/base.py index f9c203211b862..a2903be0672a2 100644 --- a/superset/db_engine_specs/base.py +++ b/superset/db_engine_specs/base.py @@ -291,8 +291,8 @@ class BaseEngineSpec: # pylint: disable=too-many-public-methods # Whether ORDER BY clause can use sql caculated expression # if True, use alias of select column for `order by` - # the True is safe for most database - # But for backward compatibility, set to False by default + # the True is safely for most database + # But for backward compatibility, False by default allows_hidden_orderby_calculated_column = False force_column_alias_quotes = False From 96d4fb87d64fd0adfaf892bfef51085ec741bf56 Mon Sep 17 00:00:00 2001 From: Yongjie Zhao Date: Fri, 22 Oct 2021 20:58:41 +0800 Subject: [PATCH 3/4] add ut --- .../db_engine_specs/base_engine_spec_tests.py | 34 +++++++++++++++++++ .../db_engine_specs/bigquery_tests.py | 32 +++++++++++++++++ 2 files changed, 66 insertions(+) diff --git a/tests/integration_tests/db_engine_specs/base_engine_spec_tests.py b/tests/integration_tests/db_engine_specs/base_engine_spec_tests.py index 275df7ba6545b..6ee8c4c847359 100644 --- a/tests/integration_tests/db_engine_specs/base_engine_spec_tests.py +++ b/tests/integration_tests/db_engine_specs/base_engine_spec_tests.py @@ -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, @@ -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 @@ -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: diff --git a/tests/integration_tests/db_engine_specs/bigquery_tests.py b/tests/integration_tests/db_engine_specs/bigquery_tests.py index dbed69607245f..ba6d0916fcd13 100644 --- a/tests/integration_tests/db_engine_specs/bigquery_tests.py +++ b/tests/integration_tests/db_engine_specs/bigquery_tests.py @@ -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): @@ -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 From 78fce8b1355c1c0c2a31dc3dc3d2c35f7347c807 Mon Sep 17 00:00:00 2001 From: Yongjie Zhao Date: Fri, 22 Oct 2021 21:12:25 +0800 Subject: [PATCH 4/4] fix lint --- superset/connectors/sqla/models.py | 2 +- superset/db_engine_specs/base.py | 2 +- superset/db_engine_specs/bigquery.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/superset/connectors/sqla/models.py b/superset/connectors/sqla/models.py index 5ed1c1876cecb..0bf5a7940ae9a 100644 --- a/superset/connectors/sqla/models.py +++ b/superset/connectors/sqla/models.py @@ -1303,7 +1303,7 @@ def get_sqla_query( # pylint: disable=too-many-arguments,too-many-locals,too-ma if ( db_engine_spec.allows_alias_in_select - and db_engine_spec.allows_hidden_orderby_calculated_column + 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) diff --git a/superset/db_engine_specs/base.py b/superset/db_engine_specs/base.py index a2903be0672a2..8444fc09e1bdb 100644 --- a/superset/db_engine_specs/base.py +++ b/superset/db_engine_specs/base.py @@ -293,7 +293,7 @@ class BaseEngineSpec: # pylint: disable=too-many-public-methods # 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_orderby_calculated_column = False + allows_hidden_cc_in_orderby = False force_column_alias_quotes = False arraysize = 0 diff --git a/superset/db_engine_specs/bigquery.py b/superset/db_engine_specs/bigquery.py index 70057de544446..c8785c4b549e8 100644 --- a/superset/db_engine_specs/bigquery.py +++ b/superset/db_engine_specs/bigquery.py @@ -99,7 +99,7 @@ class BigQueryEngineSpec(BaseEngineSpec): # same cursor, so we need to run all statements at once run_multiple_statements_as_one = True - allows_hidden_orderby_calculated_column = True + allows_hidden_cc_in_orderby = True """ https://www.python.org/dev/peps/pep-0249/#arraysize