[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:
parent
9dfa0a3f8e
commit
3250c5ac94
|
|
@ -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))",
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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)",
|
||||
|
|
|
|||
|
|
@ -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)",
|
||||
|
|
|
|||
|
|
@ -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)"
|
||||
|
|
|
|||
|
|
@ -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')",
|
||||
|
|
|
|||
|
|
@ -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)",
|
||||
|
|
|
|||
|
|
@ -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')",
|
||||
|
|
|
|||
|
|
@ -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)",
|
||||
|
|
|
|||
|
|
@ -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)",
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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')",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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})",
|
||||
|
|
|
|||
|
|
@ -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))",
|
||||
|
|
|
|||
|
|
@ -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})",
|
||||
|
|
|
|||
|
|
@ -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}))",
|
||||
|
|
|
|||
|
|
@ -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')",
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Reference in New Issue