diff --git a/UPDATING.md b/UPDATING.md index f69f2379118e6..34537a5a2d03a 100644 --- a/UPDATING.md +++ b/UPDATING.md @@ -23,6 +23,8 @@ assists people when migrating to a new version. ## Next +* [9238](https://github.com/apache/incubator-superset/pull/9238): the config option `TIME_GRAIN_FUNCTIONS` has been renamed to `TIME_GRAIN_EXPRESSIONS` to better reflect the content of the dictionary. + * [9218](https://github.com/apache/incubator-superset/pull/9218): SQLite connections have been disabled by default for analytics databases. You can optionally enable SQLite by setting `PREVENT_UNSAFE_DB_CONNECTIONS` to `False`. It is not recommended to change this setting, as arbitrary SQLite connections can lead to security vulnerabilities. diff --git a/superset/config.py b/superset/config.py index 4481a9558e354..443a92d786618 100644 --- a/superset/config.py +++ b/superset/config.py @@ -360,13 +360,14 @@ def _try_json_readsha(filepath, length): # pylint: disable=unused-argument TIME_GRAIN_ADDONS: Dict[str, str] = {} # Implementation of additional time grains per engine. +# The column to be truncated is denoted `{col}` in the expression. # For example: To implement 2 second time grain on clickhouse engine: -# TIME_GRAIN_ADDON_FUNCTIONS = { +# TIME_GRAIN_ADDON_EXPRESSIONS = { # 'clickhouse': { # 'PT2S': 'toDateTime(intDiv(toUInt32(toDateTime({col})), 2)*2)' # } # } -TIME_GRAIN_ADDON_FUNCTIONS: Dict[str, Dict[str, str]] = {} +TIME_GRAIN_ADDON_EXPRESSIONS: Dict[str, Dict[str, str]] = {} # --------------------------------------------------- # List of viz_types not allowed in your environment diff --git a/superset/connectors/sqla/models.py b/superset/connectors/sqla/models.py index 98ba5b2c3df33..f249d842a99cb 100644 --- a/superset/connectors/sqla/models.py +++ b/superset/connectors/sqla/models.py @@ -228,7 +228,9 @@ def get_timestamp_expression( col = literal_column(self.expression) else: col = column(self.column_name) - time_expr = db.db_engine_spec.get_timestamp_expr(col, pdf, time_grain) + time_expr = db.db_engine_spec.get_timestamp_expr( + col, pdf, time_grain, self.type + ) return self.table.make_sqla_column_compatible(time_expr, label) @classmethod diff --git a/superset/db_engine_specs/athena.py b/superset/db_engine_specs/athena.py index aec6c45e4b1b1..bbdba14b3f9ed 100644 --- a/superset/db_engine_specs/athena.py +++ b/superset/db_engine_specs/athena.py @@ -23,7 +23,7 @@ class AthenaEngineSpec(BaseEngineSpec): engine = "awsathena" - _time_grain_functions = { + _time_grain_expressions = { 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 ba74a4266f453..3e5da5d250e6d 100644 --- a/superset/db_engine_specs/base.py +++ b/superset/db_engine_specs/base.py @@ -132,7 +132,8 @@ class BaseEngineSpec: # pylint: disable=too-many-public-methods """Abstract class for database engine specific configurations""" engine = "base" # str as defined in sqlalchemy.engine.engine - _time_grain_functions: Dict[Optional[str], str] = {} + _date_trunc_functions: Dict[str, str] = {} + _time_grain_expressions: Dict[Optional[str], str] = {} time_groupby_inline = False limit_method = LimitMethod.FORCE_LIMIT time_secondary_columns = False @@ -204,7 +205,11 @@ def get_engine( @classmethod def get_timestamp_expr( - cls, col: ColumnClause, pdf: Optional[str], time_grain: Optional[str] + cls, + col: ColumnClause, + pdf: Optional[str], + time_grain: Optional[str], + type_: Optional[str] = None, ) -> TimestampExpression: """ Construct a TimestampExpression to be used in a SQLAlchemy query. @@ -212,14 +217,19 @@ def get_timestamp_expr( :param col: Target column for the TimestampExpression :param pdf: date format (seconds or milliseconds) :param time_grain: time grain, e.g. P1Y for 1 year + :param type_: the source column type :return: TimestampExpression object """ if time_grain: - time_expr = cls.get_time_grain_functions().get(time_grain) + time_expr = cls.get_time_grain_expressions().get(time_grain) if not time_expr: raise NotImplementedError( f"No grain spec for {time_grain} for database {cls.engine}" ) + if type_ and "{func}" in time_expr: + date_trunc_function = cls._date_trunc_functions.get(type_) + if date_trunc_function: + time_expr = time_expr.replace("{func}", date_trunc_function) else: time_expr = "{col}" @@ -240,31 +250,30 @@ def get_time_grains(cls) -> Tuple[TimeGrain, ...]: """ ret_list = [] - time_grain_functions = cls.get_time_grain_functions() time_grains = builtin_time_grains.copy() time_grains.update(config["TIME_GRAIN_ADDONS"]) - for duration, func in time_grain_functions.items(): + for duration, func in cls.get_time_grain_expressions().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]: + def get_time_grain_expressions(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 + :return: All time grain expressions 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["TIME_GRAIN_ADDON_FUNCTIONS"] - time_grain_functions.update(grain_addon_functions.get(cls.engine, {})) + time_grain_expressions = cls._time_grain_expressions.copy() + grain_addon_expressions = config["TIME_GRAIN_ADDON_EXPRESSIONS"] + time_grain_expressions.update(grain_addon_expressions.get(cls.engine, {})) blacklist: List[str] = config["TIME_GRAIN_BLACKLIST"] for key in blacklist: - time_grain_functions.pop(key) - return time_grain_functions + time_grain_expressions.pop(key) + return time_grain_expressions @classmethod def make_select_compatible( diff --git a/superset/db_engine_specs/bigquery.py b/superset/db_engine_specs/bigquery.py index 9bf4b28f3e342..4f5b5d10b117b 100644 --- a/superset/db_engine_specs/bigquery.py +++ b/superset/db_engine_specs/bigquery.py @@ -49,16 +49,23 @@ class BigQueryEngineSpec(BaseEngineSpec): """ arraysize = 5000 - _time_grain_functions = { + _date_trunc_functions = { + "DATE": "DATE_TRUNC", + "DATETIME": "DATETIME_TRUNC", + "TIME": "TIME_TRUNC", + "TIMESTAMP": "TIMESTAMP_TRUNC", + } + + _time_grain_expressions = { None: "{col}", - "PT1S": "TIMESTAMP_TRUNC({col}, SECOND)", - "PT1M": "TIMESTAMP_TRUNC({col}, MINUTE)", - "PT1H": "TIMESTAMP_TRUNC({col}, HOUR)", - "P1D": "TIMESTAMP_TRUNC({col}, DAY)", - "P1W": "TIMESTAMP_TRUNC({col}, WEEK)", - "P1M": "TIMESTAMP_TRUNC({col}, MONTH)", - "P0.25Y": "TIMESTAMP_TRUNC({col}, QUARTER)", - "P1Y": "TIMESTAMP_TRUNC({col}, YEAR)", + "PT1S": "{func}({col}, SECOND)", + "PT1M": "{func}({col}, MINUTE)", + "PT1H": "{func}({col}, HOUR)", + "P1D": "{func}({col}, DAY)", + "P1W": "{func}({col}, WEEK)", + "P1M": "{func}({col}, MONTH)", + "P0.25Y": "{func}({col}, QUARTER)", + "P1Y": "{func}({col}, YEAR)", } @classmethod @@ -68,13 +75,15 @@ def convert_dttm(cls, target_type: str, dttm: datetime) -> Optional[str]: return f"CAST('{dttm.date().isoformat()}' AS DATE)" if tt == "DATETIME": return f"""CAST('{dttm.isoformat(timespec="microseconds")}' AS DATETIME)""" + if tt == "TIME": + return f"""CAST('{dttm.strftime("%H:%M:%S.%f")}' AS TIME)""" if tt == "TIMESTAMP": return f"""CAST('{dttm.isoformat(timespec="microseconds")}' AS TIMESTAMP)""" return None @classmethod def fetch_data(cls, cursor: Any, limit: int) -> List[Tuple]: - data = super(BigQueryEngineSpec, cls).fetch_data(cursor, limit) + data = super().fetch_data(cursor, limit) if data and type(data[0]).__name__ == "Row": data = [r.values() for r in data] # type: ignore return data diff --git a/superset/db_engine_specs/clickhouse.py b/superset/db_engine_specs/clickhouse.py index b9d1ba0c2edb2..3b390530b6950 100644 --- a/superset/db_engine_specs/clickhouse.py +++ b/superset/db_engine_specs/clickhouse.py @@ -28,7 +28,7 @@ class ClickHouseEngineSpec(BaseEngineSpec): # pylint: disable=abstract-method time_secondary_columns = True time_groupby_inline = True - _time_grain_functions = { + _time_grain_expressions = { 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 794f987af5c2d..4087d218647d1 100644 --- a/superset/db_engine_specs/db2.py +++ b/superset/db_engine_specs/db2.py @@ -23,7 +23,7 @@ class Db2EngineSpec(BaseEngineSpec): force_column_alias_quotes = True max_column_name_length = 30 - _time_grain_functions = { + _time_grain_expressions = { None: "{col}", "PT1S": "CAST({col} as TIMESTAMP)" " - MICROSECOND({col}) MICROSECONDS", "PT1M": "CAST({col} as TIMESTAMP)" diff --git a/superset/db_engine_specs/dremio.py b/superset/db_engine_specs/dremio.py index ba570b49a56f1..33f027d3d1e5f 100644 --- a/superset/db_engine_specs/dremio.py +++ b/superset/db_engine_specs/dremio.py @@ -21,7 +21,7 @@ class DremioBaseEngineSpec(BaseEngineSpec): engine = "dremio" - _time_grain_functions = { + _time_grain_expressions = { None: "{col}", "PT1S": "DATE_TRUNC('second', {col})", "PT1M": "DATE_TRUNC('minute', {col})", diff --git a/superset/db_engine_specs/drill.py b/superset/db_engine_specs/drill.py index 73b5912e2f866..8da7cd8bb4bb8 100644 --- a/superset/db_engine_specs/drill.py +++ b/superset/db_engine_specs/drill.py @@ -28,7 +28,7 @@ class DrillEngineSpec(BaseEngineSpec): engine = "drill" - _time_grain_functions = { + _time_grain_expressions = { 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 35b3b47487b68..e08bbdd4a34a6 100644 --- a/superset/db_engine_specs/druid.py +++ b/superset/db_engine_specs/druid.py @@ -31,7 +31,7 @@ class DruidEngineSpec(BaseEngineSpec): # pylint: disable=abstract-method allows_joins = False allows_subqueries = True - _time_grain_functions = { + _time_grain_expressions = { None: "{col}", "PT1S": "FLOOR({col} TO SECOND)", "PT1M": "FLOOR({col} TO MINUTE)", diff --git a/superset/db_engine_specs/elasticsearch.py b/superset/db_engine_specs/elasticsearch.py index 5fdc3de1a64be..7ebc675a9b68a 100644 --- a/superset/db_engine_specs/elasticsearch.py +++ b/superset/db_engine_specs/elasticsearch.py @@ -27,7 +27,7 @@ class ElasticSearchEngineSpec(BaseEngineSpec): # pylint: disable=abstract-metho allows_joins = False allows_subqueries = True - _time_grain_functions = { + _time_grain_expressions = { None: "{col}", "PT1S": "HISTOGRAM({col}, INTERVAL 1 SECOND)", "PT1M": "HISTOGRAM({col}, INTERVAL 1 MINUTE)", diff --git a/superset/db_engine_specs/exasol.py b/superset/db_engine_specs/exasol.py index 8c14581de64f9..480a8c27e358d 100644 --- a/superset/db_engine_specs/exasol.py +++ b/superset/db_engine_specs/exasol.py @@ -26,7 +26,7 @@ class ExasolEngineSpec(BaseEngineSpec): # pylint: disable=abstract-method max_column_name_length = 128 # Exasol's DATE_TRUNC function is PostgresSQL compatible - _time_grain_functions = { + _time_grain_expressions = { None: "{col}", "PT1S": "DATE_TRUNC('second', {col})", "PT1M": "DATE_TRUNC('minute', {col})", diff --git a/superset/db_engine_specs/hana.py b/superset/db_engine_specs/hana.py index 45fc538a2857c..0a80a02677a9d 100644 --- a/superset/db_engine_specs/hana.py +++ b/superset/db_engine_specs/hana.py @@ -27,7 +27,7 @@ class HanaEngineSpec(PostgresBaseEngineSpec): force_column_alias_quotes = True max_column_name_length = 30 - _time_grain_functions = { + _time_grain_expressions = { None: "{col}", "PT1S": "TO_TIMESTAMP(SUBSTRING(TO_TIMESTAMP({col}),0,20))", "PT1M": "TO_TIMESTAMP(SUBSTRING(TO_TIMESTAMP({col}),0,17) || '00')", diff --git a/superset/db_engine_specs/impala.py b/superset/db_engine_specs/impala.py index 31b02301167b8..fd2dce5e324ea 100644 --- a/superset/db_engine_specs/impala.py +++ b/superset/db_engine_specs/impala.py @@ -27,7 +27,7 @@ class ImpalaEngineSpec(BaseEngineSpec): engine = "impala" - _time_grain_functions = { + _time_grain_expressions = { 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 e7879cf8ac5c8..5e8bff5f2f099 100644 --- a/superset/db_engine_specs/kylin.py +++ b/superset/db_engine_specs/kylin.py @@ -25,7 +25,7 @@ class KylinEngineSpec(BaseEngineSpec): # pylint: disable=abstract-method engine = "kylin" - _time_grain_functions = { + _time_grain_expressions = { 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 be91d1438a90b..b7ced5feac45b 100644 --- a/superset/db_engine_specs/mssql.py +++ b/superset/db_engine_specs/mssql.py @@ -29,7 +29,7 @@ class MssqlEngineSpec(BaseEngineSpec): limit_method = LimitMethod.WRAP_SQL max_column_name_length = 128 - _time_grain_functions = { + _time_grain_expressions = { 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 023dd76aa12d9..cf33298e246ab 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_expressions = { 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 e72eef0c1ae53..0488aaedfc3c4 100644 --- a/superset/db_engine_specs/oracle.py +++ b/superset/db_engine_specs/oracle.py @@ -26,7 +26,7 @@ class OracleEngineSpec(BaseEngineSpec): force_column_alias_quotes = True max_column_name_length = 30 - _time_grain_functions = { + _time_grain_expressions = { 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 8a651e73f9f3b..ddd6f83c0d6c5 100644 --- a/superset/db_engine_specs/pinot.py +++ b/superset/db_engine_specs/pinot.py @@ -29,7 +29,7 @@ class PinotEngineSpec(BaseEngineSpec): # pylint: disable=abstract-method allows_column_aliases = False # Pinot does its own conversion below - _time_grain_functions: Dict[Optional[str], str] = { + _time_grain_expressions: Dict[Optional[str], str] = { "PT1S": "1:SECONDS", "PT1M": "1:MINUTES", "PT1H": "1:HOURS", @@ -51,7 +51,11 @@ class PinotEngineSpec(BaseEngineSpec): # pylint: disable=abstract-method @classmethod def get_timestamp_expr( - cls, col: ColumnClause, pdf: Optional[str], time_grain: Optional[str] + cls, + col: ColumnClause, + pdf: Optional[str], + time_grain: Optional[str], + type_: Optional[str] = None, ) -> TimestampExpression: is_epoch = pdf in ("epoch_s", "epoch_ms") @@ -75,7 +79,7 @@ def get_timestamp_expr( else: seconds_or_ms = "MILLISECONDS" if pdf == "epoch_ms" else "SECONDS" tf = f"1:{seconds_or_ms}:EPOCH" - granularity = cls.get_time_grain_functions().get(time_grain) + granularity = cls.get_time_grain_expressions().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 388ae6aa7912f..e99432b23c0d4 100644 --- a/superset/db_engine_specs/postgres.py +++ b/superset/db_engine_specs/postgres.py @@ -38,7 +38,7 @@ class PostgresBaseEngineSpec(BaseEngineSpec): engine = "" - _time_grain_functions = { + _time_grain_expressions = { 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 18bd95858c63f..038b25b8cace9 100644 --- a/superset/db_engine_specs/presto.py +++ b/superset/db_engine_specs/presto.py @@ -100,7 +100,7 @@ def get_children(column: Dict[str, str]) -> List[Dict[str, str]]: class PrestoEngineSpec(BaseEngineSpec): engine = "presto" - _time_grain_functions = { + _time_grain_expressions = { 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 ada9fae8bcff2..f42267c26fc66 100644 --- a/superset/db_engine_specs/snowflake.py +++ b/superset/db_engine_specs/snowflake.py @@ -28,7 +28,7 @@ class SnowflakeEngineSpec(PostgresBaseEngineSpec): force_column_alias_quotes = True max_column_name_length = 256 - _time_grain_functions = { + _time_grain_expressions = { 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 12d422adddf7c..d824c3203e690 100644 --- a/superset/db_engine_specs/sqlite.py +++ b/superset/db_engine_specs/sqlite.py @@ -30,7 +30,7 @@ class SqliteEngineSpec(BaseEngineSpec): engine = "sqlite" - _time_grain_functions = { + _time_grain_expressions = { 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 bbc8475feda2d..0226baef62ee4 100644 --- a/superset/db_engine_specs/teradata.py +++ b/superset/db_engine_specs/teradata.py @@ -24,7 +24,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_expressions = { None: "{col}", "PT1M": "TRUNC(CAST({col} as DATE), 'MI')", "PT1H": "TRUNC(CAST({col} as DATE), 'HH')", diff --git a/tests/db_engine_specs/base_engine_spec_tests.py b/tests/db_engine_specs/base_engine_spec_tests.py index 4ba4e31262eed..1d68e981bc3a0 100644 --- a/tests/db_engine_specs/base_engine_spec_tests.py +++ b/tests/db_engine_specs/base_engine_spec_tests.py @@ -154,13 +154,13 @@ def test_limit_with_non_token_limit(self): def test_time_grain_blacklist(self): with app.app_context(): app.config["TIME_GRAIN_BLACKLIST"] = ["PT1M"] - time_grain_functions = SqliteEngineSpec.get_time_grain_functions() + time_grain_functions = SqliteEngineSpec.get_time_grain_expressions() 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"] = { + app.config["TIME_GRAIN_ADDON_EXPRESSIONS"] = { "sqlite": {"PTXM": "ABC({col})"} } time_grains = SqliteEngineSpec.get_time_grains() @@ -174,7 +174,7 @@ def test_engine_time_grain_validity(self): for engine in engines.values(): if engine is not BaseEngineSpec: # make sure time grain functions have been defined - self.assertGreater(len(engine.get_time_grain_functions()), 0) + self.assertGreater(len(engine.get_time_grain_expressions()), 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) diff --git a/tests/db_engine_specs/bigquery_tests.py b/tests/db_engine_specs/bigquery_tests.py index 9c77970d55115..c9b9878b05047 100644 --- a/tests/db_engine_specs/bigquery_tests.py +++ b/tests/db_engine_specs/bigquery_tests.py @@ -22,35 +22,38 @@ class BigQueryTestCase(DbEngineSpecTestCase): def test_bigquery_sqla_column_label(self): - label = BigQueryEngineSpec.make_label_compatible(column("Col").name) - label_expected = "Col" - self.assertEqual(label, label_expected) - - label = BigQueryEngineSpec.make_label_compatible(column("SUM(x)").name) - label_expected = "SUM_x__5f110" - self.assertEqual(label, label_expected) - - label = BigQueryEngineSpec.make_label_compatible(column("SUM[x]").name) - label_expected = "SUM_x__7ebe1" - self.assertEqual(label, label_expected) - - label = BigQueryEngineSpec.make_label_compatible(column("12345_col").name) - label_expected = "_12345_col_8d390" - self.assertEqual(label, label_expected) + test_cases = { + "Col": "Col", + "SUM(x)": "SUM_x__5f110", + "SUM[x]": "SUM_x__7ebe1", + "12345_col": "_12345_col_8d390", + } + for original, expected in test_cases.items(): + actual = BigQueryEngineSpec.make_label_compatible(column(original).name) + self.assertEqual(actual, expected) def test_convert_dttm(self): dttm = self.get_dttm() - - self.assertEqual( - BigQueryEngineSpec.convert_dttm("DATE", dttm), "CAST('2019-01-02' AS DATE)" - ) - - self.assertEqual( - BigQueryEngineSpec.convert_dttm("DATETIME", dttm), - "CAST('2019-01-02T03:04:05.678900' AS DATETIME)", - ) - - self.assertEqual( - BigQueryEngineSpec.convert_dttm("TIMESTAMP", dttm), - "CAST('2019-01-02T03:04:05.678900' AS TIMESTAMP)", - ) + test_cases = { + "DATE": "CAST('2019-01-02' AS DATE)", + "DATETIME": "CAST('2019-01-02T03:04:05.678900' AS DATETIME)", + "TIMESTAMP": "CAST('2019-01-02T03:04:05.678900' AS TIMESTAMP)", + } + + for target_type, expected in test_cases.items(): + actual = BigQueryEngineSpec.convert_dttm(target_type, dttm) + self.assertEqual(actual, expected) + + def test_timegrain_expressions(self): + col = column("temporal") + test_cases = { + "DATE": "DATE_TRUNC(temporal, HOUR)", + "TIME": "TIME_TRUNC(temporal, HOUR)", + "DATETIME": "DATETIME_TRUNC(temporal, HOUR)", + "TIMESTAMP": "TIMESTAMP_TRUNC(temporal, HOUR)", + } + for type_, expected in test_cases.items(): + actual = BigQueryEngineSpec.get_timestamp_expr( + col=col, pdf=None, time_grain="PT1H", type_=type_ + ) + self.assertEqual(str(actual), expected)