[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
This commit is contained in:
Ville Brofeldt 2019-09-08 08:34:40 +02:00 committed by GitHub
parent 9dfa0a3f8e
commit 3250c5ac94
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
19 changed files with 67 additions and 66 deletions

View File

@ -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))",

View File

@ -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(

View File

@ -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)",

View File

@ -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)",

View File

@ -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)"

View File

@ -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')",

View File

@ -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)",

View File

@ -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')",

View File

@ -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)",

View File

@ -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)",

View File

@ -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"

View File

@ -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')",

View File

@ -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

View File

@ -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})",

View File

@ -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))",

View File

@ -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})",

View File

@ -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}))",

View File

@ -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')",

View File

@ -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)