From 3250c5ac948cae8edd7deb2bf94ef4154608b705 Mon Sep 17 00:00:00 2001 From: Ville Brofeldt <33317356+villebro@users.noreply.github.com> Date: Sun, 8 Sep 2019 08:34:40 +0200 Subject: [PATCH] [bugfix] fix timegrain addon regression (#8165) * Fix regression in time grain addons * Revert privatization of time_grain_functions * Fix test * Rename variable * Fix test * Fix typing error * Refactor and add tests * Add TODO --- superset/db_engine_specs/athena.py | 2 +- superset/db_engine_specs/base.py | 62 ++++++++++++-------------- superset/db_engine_specs/bigquery.py | 2 +- superset/db_engine_specs/clickhouse.py | 2 +- superset/db_engine_specs/db2.py | 2 +- superset/db_engine_specs/drill.py | 2 +- superset/db_engine_specs/druid.py | 2 +- superset/db_engine_specs/impala.py | 2 +- superset/db_engine_specs/kylin.py | 2 +- superset/db_engine_specs/mssql.py | 2 +- superset/db_engine_specs/mysql.py | 2 +- superset/db_engine_specs/oracle.py | 2 +- superset/db_engine_specs/pinot.py | 4 +- superset/db_engine_specs/postgres.py | 2 +- superset/db_engine_specs/presto.py | 2 +- superset/db_engine_specs/snowflake.py | 2 +- superset/db_engine_specs/sqlite.py | 2 +- superset/db_engine_specs/teradata.py | 2 +- tests/db_engine_specs_test.py | 35 ++++++++------- 19 files changed, 67 insertions(+), 66 deletions(-) diff --git a/superset/db_engine_specs/athena.py b/superset/db_engine_specs/athena.py index 861516d43..e516664b7 100644 --- a/superset/db_engine_specs/athena.py +++ b/superset/db_engine_specs/athena.py @@ -23,7 +23,7 @@ from superset.db_engine_specs.base import BaseEngineSpec class AthenaEngineSpec(BaseEngineSpec): engine = "awsathena" - time_grain_functions = { + _time_grain_functions = { None: "{col}", "PT1S": "date_trunc('second', CAST({col} AS TIMESTAMP))", "PT1M": "date_trunc('minute', CAST({col} AS TIMESTAMP))", diff --git a/superset/db_engine_specs/base.py b/superset/db_engine_specs/base.py index 1a5bb8c65..7df093fdb 100644 --- a/superset/db_engine_specs/base.py +++ b/superset/db_engine_specs/base.py @@ -109,34 +109,11 @@ class LimitMethod(object): FORCE_LIMIT = "force_limit" -def _create_time_grains_tuple( - time_grains: Dict[Optional[str], str], - time_grain_functions: Dict[Optional[str], str], - blacklist: List[str], -) -> Tuple[TimeGrain, ...]: - """ - function for creating a tuple of time grains based on time grains provided by - the engine and any potential additional or blacklisted grains in the config file. - - :param time_grains: all time grains supported by the engine + config files - :param time_grain_functions: mapping between time grain id and sql expression - :param blacklist: list of time grain ids to be excluded - :return: final collection of time grains - """ - ret_list = [] - blacklist = blacklist if blacklist else [] - for duration, func in time_grain_functions.items(): - if duration in time_grains and duration not in blacklist: - name = time_grains[duration] - ret_list.append(TimeGrain(name, _(name), func, duration)) - return tuple(ret_list) - - class BaseEngineSpec: """Abstract class for database engine specific configurations""" engine = "base" # str as defined in sqlalchemy.engine.engine - time_grain_functions: Dict[Optional[str], str] = {} + _time_grain_functions: Dict[Optional[str], str] = {} time_groupby_inline = False limit_method = LimitMethod.FORCE_LIMIT time_secondary_columns = False @@ -161,7 +138,7 @@ class BaseEngineSpec: :return: TimestampExpression object """ if time_grain: - time_expr = cls.time_grain_functions.get(time_grain) + time_expr = cls.get_time_grain_functions().get(time_grain) if not time_expr: raise NotImplementedError( f"No grain spec for {time_grain} for database {cls.engine}" @@ -180,18 +157,37 @@ class BaseEngineSpec: @classmethod def get_time_grains(cls) -> Tuple[TimeGrain, ...]: """ - Generate a tuple of time grains based on time grains provided by the engine - and any potential additional or blacklisted grains in the config file. + Generate a tuple of supported time grains. :return: All time grains supported by the engine """ - blacklist: List[str] = config.get("TIME_GRAIN_BLACKLIST", []) - supported_grains = builtin_time_grains.copy() - supported_grains.update(config.get("TIME_GRAIN_ADDONS", {})) - grain_functions = cls.time_grain_functions.copy() + + ret_list = [] + time_grain_functions = cls.get_time_grain_functions() + time_grains = builtin_time_grains.copy() + time_grains.update(config.get("TIME_GRAIN_ADDONS", {})) + for duration, func in time_grain_functions.items(): + if duration in time_grains: + name = time_grains[duration] + ret_list.append(TimeGrain(name, _(name), func, duration)) + return tuple(ret_list) + + @classmethod + def get_time_grain_functions(cls) -> Dict[Optional[str], str]: + """ + Return a dict of all supported time grains including any potential added grains + but excluding any potentially blacklisted grains in the config file. + + :return: All time grain functions supported by the engine + """ + # TODO: use @memoize decorator or similar to avoid recomputation on every call + time_grain_functions = cls._time_grain_functions.copy() grain_addon_functions = config.get("TIME_GRAIN_ADDON_FUNCTIONS", {}) - grain_functions.update(grain_addon_functions.get(cls.engine, {})) - return _create_time_grains_tuple(supported_grains, grain_functions, blacklist) + time_grain_functions.update(grain_addon_functions.get(cls.engine, {})) + blacklist: List[str] = config.get("TIME_GRAIN_BLACKLIST", []) + for key in blacklist: + time_grain_functions.pop(key) + return time_grain_functions @classmethod def make_select_compatible( diff --git a/superset/db_engine_specs/bigquery.py b/superset/db_engine_specs/bigquery.py index addf4fd11..2e0476c95 100644 --- a/superset/db_engine_specs/bigquery.py +++ b/superset/db_engine_specs/bigquery.py @@ -45,7 +45,7 @@ class BigQueryEngineSpec(BaseEngineSpec): """ arraysize = 5000 - time_grain_functions = { + _time_grain_functions = { None: "{col}", "PT1S": "TIMESTAMP_TRUNC({col}, SECOND)", "PT1M": "TIMESTAMP_TRUNC({col}, MINUTE)", diff --git a/superset/db_engine_specs/clickhouse.py b/superset/db_engine_specs/clickhouse.py index 1095e1fc1..e72f8752e 100644 --- a/superset/db_engine_specs/clickhouse.py +++ b/superset/db_engine_specs/clickhouse.py @@ -28,7 +28,7 @@ class ClickHouseEngineSpec(BaseEngineSpec): time_secondary_columns = True time_groupby_inline = True - time_grain_functions = { + _time_grain_functions = { None: "{col}", "PT1M": "toStartOfMinute(toDateTime({col}))", "PT5M": "toDateTime(intDiv(toUInt32(toDateTime({col})), 300)*300)", diff --git a/superset/db_engine_specs/db2.py b/superset/db_engine_specs/db2.py index 30822e3d7..93cb76f71 100644 --- a/superset/db_engine_specs/db2.py +++ b/superset/db_engine_specs/db2.py @@ -26,7 +26,7 @@ class Db2EngineSpec(BaseEngineSpec): force_column_alias_quotes = True max_column_name_length = 30 - time_grain_functions = { + _time_grain_functions = { None: "{col}", "PT1S": "CAST({col} as TIMESTAMP)" " - MICROSECOND({col}) MICROSECONDS", "PT1M": "CAST({col} as TIMESTAMP)" diff --git a/superset/db_engine_specs/drill.py b/superset/db_engine_specs/drill.py index d82a25ff6..9ebe87703 100644 --- a/superset/db_engine_specs/drill.py +++ b/superset/db_engine_specs/drill.py @@ -26,7 +26,7 @@ class DrillEngineSpec(BaseEngineSpec): engine = "drill" - time_grain_functions = { + _time_grain_functions = { None: "{col}", "PT1S": "NEARESTDATE({col}, 'SECOND')", "PT1M": "NEARESTDATE({col}, 'MINUTE')", diff --git a/superset/db_engine_specs/druid.py b/superset/db_engine_specs/druid.py index 8cd0c4cb9..78a2a6461 100644 --- a/superset/db_engine_specs/druid.py +++ b/superset/db_engine_specs/druid.py @@ -25,7 +25,7 @@ class DruidEngineSpec(BaseEngineSpec): allows_joins = False allows_subqueries = True - time_grain_functions = { + _time_grain_functions = { None: "{col}", "PT1S": "FLOOR({col} TO SECOND)", "PT1M": "FLOOR({col} TO MINUTE)", diff --git a/superset/db_engine_specs/impala.py b/superset/db_engine_specs/impala.py index f2fa57a39..4feb3fc0d 100644 --- a/superset/db_engine_specs/impala.py +++ b/superset/db_engine_specs/impala.py @@ -28,7 +28,7 @@ class ImpalaEngineSpec(BaseEngineSpec): engine = "impala" - time_grain_functions = { + _time_grain_functions = { None: "{col}", "PT1M": "TRUNC({col}, 'MI')", "PT1H": "TRUNC({col}, 'HH')", diff --git a/superset/db_engine_specs/kylin.py b/superset/db_engine_specs/kylin.py index b3ebe81b0..edde71ded 100644 --- a/superset/db_engine_specs/kylin.py +++ b/superset/db_engine_specs/kylin.py @@ -25,7 +25,7 @@ class KylinEngineSpec(BaseEngineSpec): engine = "kylin" - time_grain_functions = { + _time_grain_functions = { None: "{col}", "PT1S": "CAST(FLOOR(CAST({col} AS TIMESTAMP) TO SECOND) AS TIMESTAMP)", "PT1M": "CAST(FLOOR(CAST({col} AS TIMESTAMP) TO MINUTE) AS TIMESTAMP)", diff --git a/superset/db_engine_specs/mssql.py b/superset/db_engine_specs/mssql.py index 0c4685118..4e5a4fe89 100644 --- a/superset/db_engine_specs/mssql.py +++ b/superset/db_engine_specs/mssql.py @@ -31,7 +31,7 @@ class MssqlEngineSpec(BaseEngineSpec): limit_method = LimitMethod.WRAP_SQL max_column_name_length = 128 - time_grain_functions = { + _time_grain_functions = { None: "{col}", "PT1S": "DATEADD(second, DATEDIFF(second, '2000-01-01', {col}), '2000-01-01')", "PT1M": "DATEADD(minute, DATEDIFF(minute, 0, {col}), 0)", diff --git a/superset/db_engine_specs/mysql.py b/superset/db_engine_specs/mysql.py index 3a48032e7..a8964a411 100644 --- a/superset/db_engine_specs/mysql.py +++ b/superset/db_engine_specs/mysql.py @@ -29,7 +29,7 @@ class MySQLEngineSpec(BaseEngineSpec): engine = "mysql" max_column_name_length = 64 - time_grain_functions = { + _time_grain_functions = { None: "{col}", "PT1S": "DATE_ADD(DATE({col}), " "INTERVAL (HOUR({col})*60*60 + MINUTE({col})*60" diff --git a/superset/db_engine_specs/oracle.py b/superset/db_engine_specs/oracle.py index eaa815c5b..3b42f9d43 100644 --- a/superset/db_engine_specs/oracle.py +++ b/superset/db_engine_specs/oracle.py @@ -27,7 +27,7 @@ class OracleEngineSpec(PostgresBaseEngineSpec): force_column_alias_quotes = True max_column_name_length = 30 - time_grain_functions = { + _time_grain_functions = { None: "{col}", "PT1S": "CAST({col} as DATE)", "PT1M": "TRUNC(CAST({col} as DATE), 'MI')", diff --git a/superset/db_engine_specs/pinot.py b/superset/db_engine_specs/pinot.py index f4bb84333..af3fc2b25 100644 --- a/superset/db_engine_specs/pinot.py +++ b/superset/db_engine_specs/pinot.py @@ -29,7 +29,7 @@ class PinotEngineSpec(BaseEngineSpec): allows_column_aliases = False # Pinot does its own conversion below - time_grain_functions: Dict[Optional[str], str] = { + _time_grain_functions: Dict[Optional[str], str] = { "PT1S": "1:SECONDS", "PT1M": "1:MINUTES", "PT1H": "1:HOURS", @@ -52,7 +52,7 @@ class PinotEngineSpec(BaseEngineSpec): # We are not really converting any time units, just bucketing them. seconds_or_ms = "MILLISECONDS" if pdf == "epoch_ms" else "SECONDS" tf = f"1:{seconds_or_ms}:EPOCH" - granularity = cls.time_grain_functions.get(time_grain) + granularity = cls.get_time_grain_functions().get(time_grain) if not granularity: raise NotImplementedError("No pinot grain spec for " + str(time_grain)) # In pinot the output is a string since there is no timestamp column like pg diff --git a/superset/db_engine_specs/postgres.py b/superset/db_engine_specs/postgres.py index 6bfe12da7..4716b0758 100644 --- a/superset/db_engine_specs/postgres.py +++ b/superset/db_engine_specs/postgres.py @@ -28,7 +28,7 @@ class PostgresBaseEngineSpec(BaseEngineSpec): engine = "" - time_grain_functions = { + _time_grain_functions = { None: "{col}", "PT1S": "DATE_TRUNC('second', {col})", "PT1M": "DATE_TRUNC('minute', {col})", diff --git a/superset/db_engine_specs/presto.py b/superset/db_engine_specs/presto.py index 87aac9b9c..e648f65a0 100644 --- a/superset/db_engine_specs/presto.py +++ b/superset/db_engine_specs/presto.py @@ -43,7 +43,7 @@ QueryStatus = utils.QueryStatus class PrestoEngineSpec(BaseEngineSpec): engine = "presto" - time_grain_functions = { + _time_grain_functions = { None: "{col}", "PT1S": "date_trunc('second', CAST({col} AS TIMESTAMP))", "PT1M": "date_trunc('minute', CAST({col} AS TIMESTAMP))", diff --git a/superset/db_engine_specs/snowflake.py b/superset/db_engine_specs/snowflake.py index 07d266d49..8a1edc774 100644 --- a/superset/db_engine_specs/snowflake.py +++ b/superset/db_engine_specs/snowflake.py @@ -25,7 +25,7 @@ class SnowflakeEngineSpec(PostgresBaseEngineSpec): force_column_alias_quotes = True max_column_name_length = 256 - time_grain_functions = { + _time_grain_functions = { None: "{col}", "PT1S": "DATE_TRUNC('SECOND', {col})", "PT1M": "DATE_TRUNC('MINUTE', {col})", diff --git a/superset/db_engine_specs/sqlite.py b/superset/db_engine_specs/sqlite.py index 58a2c6712..28c884305 100644 --- a/superset/db_engine_specs/sqlite.py +++ b/superset/db_engine_specs/sqlite.py @@ -27,7 +27,7 @@ from superset.utils import core as utils class SqliteEngineSpec(BaseEngineSpec): engine = "sqlite" - time_grain_functions = { + _time_grain_functions = { None: "{col}", "PT1S": "DATETIME(STRFTIME('%Y-%m-%dT%H:%M:%S', {col}))", "PT1M": "DATETIME(STRFTIME('%Y-%m-%dT%H:%M:00', {col}))", diff --git a/superset/db_engine_specs/teradata.py b/superset/db_engine_specs/teradata.py index ceba88c68..fc6304908 100644 --- a/superset/db_engine_specs/teradata.py +++ b/superset/db_engine_specs/teradata.py @@ -25,7 +25,7 @@ class TeradataEngineSpec(BaseEngineSpec): limit_method = LimitMethod.WRAP_SQL max_column_name_length = 30 # since 14.10 this is 128 - time_grain_functions = { + _time_grain_functions = { None: "{col}", "PT1M": "TRUNC(CAST({col} as DATE), 'MI')", "PT1H": "TRUNC(CAST({col} as DATE), 'HH')", diff --git a/tests/db_engine_specs_test.py b/tests/db_engine_specs_test.py index 3419d25e1..2824205de 100644 --- a/tests/db_engine_specs_test.py +++ b/tests/db_engine_specs_test.py @@ -24,12 +24,9 @@ from sqlalchemy.engine.result import RowProxy from sqlalchemy.sql import select from sqlalchemy.types import String, UnicodeText +from superset import app from superset.db_engine_specs import engines -from superset.db_engine_specs.base import ( - _create_time_grains_tuple, - BaseEngineSpec, - builtin_time_grains, -) +from superset.db_engine_specs.base import BaseEngineSpec, builtin_time_grains from superset.db_engine_specs.bigquery import BigQueryEngineSpec from superset.db_engine_specs.hive import HiveEngineSpec from superset.db_engine_specs.mssql import MssqlEngineSpec @@ -38,6 +35,7 @@ from superset.db_engine_specs.oracle import OracleEngineSpec from superset.db_engine_specs.pinot import PinotEngineSpec from superset.db_engine_specs.postgres import PostgresEngineSpec from superset.db_engine_specs.presto import PrestoEngineSpec +from superset.db_engine_specs.sqlite import SqliteEngineSpec from superset.models.core import Database from superset.utils.core import get_example_database from .base_tests import SupersetTestCase @@ -315,14 +313,21 @@ class DbEngineSpecsTestCase(SupersetTestCase): ) def test_time_grain_blacklist(self): - blacklist = ["PT1M"] - time_grains = {"PT1S": "second", "PT1M": "minute"} - time_grain_functions = {"PT1S": "{col}", "PT1M": "{col}"} - time_grains = _create_time_grains_tuple( - time_grains, time_grain_functions, blacklist - ) - self.assertEqual(1, len(time_grains)) - self.assertEqual("PT1S", time_grains[0].duration) + with app.app_context(): + app.config["TIME_GRAIN_BLACKLIST"] = ["PT1M"] + time_grain_functions = SqliteEngineSpec.get_time_grain_functions() + self.assertNotIn("PT1M", time_grain_functions) + + def test_time_grain_addons(self): + with app.app_context(): + app.config["TIME_GRAIN_ADDONS"] = {"PTXM": "x seconds"} + app.config["TIME_GRAIN_ADDON_FUNCTIONS"] = { + "sqlite": {"PTXM": "ABC({col})"} + } + time_grains = SqliteEngineSpec.get_time_grains() + time_grain_addon = time_grains[-1] + self.assertEqual("PTXM", time_grain_addon.duration) + self.assertEqual("x seconds", time_grain_addon.label) def test_engine_time_grain_validity(self): time_grains = set(builtin_time_grains.keys()) @@ -330,8 +335,8 @@ class DbEngineSpecsTestCase(SupersetTestCase): for engine in engines.values(): if engine is not BaseEngineSpec: # make sure time grain functions have been defined - self.assertGreater(len(engine.time_grain_functions), 0) - # make sure that all defined time grains are supported + self.assertGreater(len(engine.get_time_grain_functions()), 0) + # make sure all defined time grains are supported defined_grains = {grain.duration for grain in engine.get_time_grains()} intersection = time_grains.intersection(defined_grains) self.assertSetEqual(defined_grains, intersection, engine)