From 086238fb10de2a5e5b8c760c50c77eadc29484e4 Mon Sep 17 00:00:00 2001 From: Elizabeth Thompson Date: Thu, 25 Mar 2021 18:00:52 -0700 Subject: [PATCH] feat: sort time grain configs (#13720) * sort time grain configs * Fix lint Co-authored-by: Beto Dealmeida --- superset/db_engine_specs/base.py | 75 ++++++++++++++++++- .../db_engine_specs/base_engine_spec_tests.py | 68 +++++++++++++++++ 2 files changed, 142 insertions(+), 1 deletion(-) diff --git a/superset/db_engine_specs/base.py b/superset/db_engine_specs/base.py index 1685ec7dea..b0545e625c 100644 --- a/superset/db_engine_specs/base.py +++ b/superset/db_engine_specs/base.py @@ -351,6 +351,68 @@ class BaseEngineSpec: # pylint: disable=too-many-public-methods ret_list.append(TimeGrain(name, _(name), func, duration)) return tuple(ret_list) + @classmethod + def _sort_time_grains( + cls, val: Tuple[Optional[str], str], index: int + ) -> Union[float, int, str]: + """ + Return an ordered time-based value of a portion of a time grain + for sorting + Values are expected to be either None or start with P or PT + Have a numerical value in the middle and end with + a value for the time interval + It can also start or end with epoch start time denoting a range + i.e, week beginning or ending with a day + """ + pos = { + "FIRST": 0, + "SECOND": 1, + "THIRD": 2, + "LAST": 3, + } + + if val[0] is None: + return pos["FIRST"] + + prog = re.compile(r"(.*\/)?(P|PT)([0-9\.]+)(S|M|H|D|W|M|Y)(\/.*)?") + result = prog.match(val[0]) + + # for any time grains that don't match the format, put them at the end + if result is None: + return pos["LAST"] + + second_minute_hour = ["S", "M", "H"] + day_week_month_year = ["D", "W", "M", "Y"] + is_less_than_day = result.group(2) == "PT" + interval = result.group(4) + epoch_time_start_string = result.group(1) or result.group(5) + has_starting_or_ending = bool(len(epoch_time_start_string or "")) + + def sort_day_week() -> int: + if has_starting_or_ending: + return pos["LAST"] + if is_less_than_day: + return pos["SECOND"] + return pos["THIRD"] + + def sort_interval() -> float: + if is_less_than_day: + return second_minute_hour.index(interval) + return day_week_month_year.index(interval) + + # 0: all "PT" values should come before "P" values (i.e, PT10M) + # 1: order values within the above arrays ("D" before "W") + # 2: sort by numeric value (PT10M before PT15M) + # 3: sort by any week starting/ending values + plist = { + 0: sort_day_week(), + 1: pos["SECOND"] if is_less_than_day else pos["THIRD"], + 2: sort_interval(), + 3: float(result.group(3)), + } + + return plist.get(index, 0) + @classmethod def get_time_grain_expressions(cls) -> Dict[Optional[str], str]: """ @@ -366,7 +428,18 @@ class BaseEngineSpec: # pylint: disable=too-many-public-methods denylist: List[str] = config["TIME_GRAIN_DENYLIST"] for key in denylist: time_grain_expressions.pop(key) - return time_grain_expressions + + return dict( + sorted( + time_grain_expressions.items(), + key=lambda x: ( + cls._sort_time_grains(x, 0), + cls._sort_time_grains(x, 1), + cls._sort_time_grains(x, 2), + cls._sort_time_grains(x, 3), + ), + ) + ) @classmethod def make_select_compatible( diff --git a/tests/db_engine_specs/base_engine_spec_tests.py b/tests/db_engine_specs/base_engine_spec_tests.py index 960a8d1c46..8c1c25d6a9 100644 --- a/tests/db_engine_specs/base_engine_spec_tests.py +++ b/tests/db_engine_specs/base_engine_spec_tests.py @@ -25,6 +25,7 @@ from superset.db_engine_specs.base import ( builtin_time_grains, LimitMethod, ) +from superset.db_engine_specs.mysql import MySQLEngineSpec from superset.db_engine_specs.sqlite import SqliteEngineSpec from superset.sql_parse import ParsedQuery from superset.utils.core import get_example_database @@ -197,6 +198,73 @@ class TestDbEngineSpecs(TestDbEngineSpec): intersection = time_grains.intersection(defined_grains) self.assertSetEqual(defined_grains, intersection, engine) + def test_get_time_grain_with_config(self): + """ Should concatenate from configs and then sort in the proper order """ + app.config["TIME_GRAIN_ADDON_EXPRESSIONS"] = { + "mysql": { + "PT2H": "foo", + "PT4H": "foo", + "PT6H": "foo", + "PT8H": "foo", + "PT10H": "foo", + "PT12H": "foo", + "PT1S": "foo", + } + } + time_grains = MySQLEngineSpec.get_time_grain_expressions() + self.assertEqual( + list(time_grains.keys()), + [ + None, + "PT1S", + "PT1M", + "PT1H", + "PT2H", + "PT4H", + "PT6H", + "PT8H", + "PT10H", + "PT12H", + "P1D", + "P1W", + "P1M", + "P0.25Y", + "P1Y", + "1969-12-29T00:00:00Z/P1W", + ], + ) + app.config["TIME_GRAIN_ADDON_EXPRESSIONS"] = {} + + def test_get_time_grain_expressions(self): + time_grains = MySQLEngineSpec.get_time_grain_expressions() + self.assertEqual( + list(time_grains.keys()), + [ + None, + "PT1S", + "PT1M", + "PT1H", + "P1D", + "P1W", + "P1M", + "P0.25Y", + "P1Y", + "1969-12-29T00:00:00Z/P1W", + ], + ) + + def test_get_time_grain_with_unkown_values(self): + """ Should concatenate from configs and then sort in the proper order + putting unknown patterns at the end """ + app.config["TIME_GRAIN_ADDON_EXPRESSIONS"] = { + "mysql": {"PT2H": "foo", "weird": "foo", "PT12H": "foo",} + } + time_grains = MySQLEngineSpec.get_time_grain_expressions() + self.assertEqual( + list(time_grains)[-1], "weird", + ) + app.config["TIME_GRAIN_ADDON_EXPRESSIONS"] = {} + def test_get_table_names(self): inspector = mock.Mock() inspector.get_table_names = mock.Mock(return_value=["schema.table", "table_2"])