From ac66ae813f0f1a2377d993d1a382187c9050cf5a Mon Sep 17 00:00:00 2001 From: Sam Firke Date: Fri, 20 Sep 2024 13:34:10 -0400 Subject: [PATCH] fix(db_engine_specs): add a few missing time grains to Postgres spec (#30325) --- superset/db_engine_specs/postgres.py | 6 ++ .../db_engine_specs/test_postgres.py | 82 +++++++++++++------ 2 files changed, 65 insertions(+), 23 deletions(-) diff --git a/superset/db_engine_specs/postgres.py b/superset/db_engine_specs/postgres.py index 0a638f65fe136..70373927d521b 100644 --- a/superset/db_engine_specs/postgres.py +++ b/superset/db_engine_specs/postgres.py @@ -103,7 +103,13 @@ class PostgresBaseEngineSpec(BaseEngineSpec): _time_grain_expressions = { None: "{col}", TimeGrain.SECOND: "DATE_TRUNC('second', {col})", + TimeGrain.FIVE_SECONDS: "DATE_TRUNC('minute', {col}) + INTERVAL '5 seconds' * FLOOR(EXTRACT(SECOND FROM {col}) / 5)", + TimeGrain.THIRTY_SECONDS: "DATE_TRUNC('minute', {col}) + INTERVAL '30 seconds' * FLOOR(EXTRACT(SECOND FROM {col}) / 30)", TimeGrain.MINUTE: "DATE_TRUNC('minute', {col})", + TimeGrain.FIVE_MINUTES: "DATE_TRUNC('hour', {col}) + INTERVAL '5 minutes' * FLOOR(EXTRACT(MINUTE FROM {col}) / 5)", + TimeGrain.TEN_MINUTES: "DATE_TRUNC('hour', {col}) + INTERVAL '10 minutes' * FLOOR(EXTRACT(MINUTE FROM {col}) / 10)", + TimeGrain.FIFTEEN_MINUTES: "DATE_TRUNC('hour', {col}) + INTERVAL '15 minutes' * FLOOR(EXTRACT(MINUTE FROM {col}) / 15)", + TimeGrain.THIRTY_MINUTES: "DATE_TRUNC('hour', {col}) + INTERVAL '30 minutes' * FLOOR(EXTRACT(MINUTE FROM {col}) / 30)", TimeGrain.HOUR: "DATE_TRUNC('hour', {col})", TimeGrain.DAY: "DATE_TRUNC('day', {col})", TimeGrain.WEEK: "DATE_TRUNC('week', {col})", diff --git a/tests/unit_tests/db_engine_specs/test_postgres.py b/tests/unit_tests/db_engine_specs/test_postgres.py index da5a4ccf80600..08e2034adbe52 100644 --- a/tests/unit_tests/db_engine_specs/test_postgres.py +++ b/tests/unit_tests/db_engine_specs/test_postgres.py @@ -20,10 +20,11 @@ import pytest from pytest_mock import MockerFixture -from sqlalchemy import types +from sqlalchemy import column, types from sqlalchemy.dialects.postgresql import DOUBLE_PRECISION, ENUM, JSON from sqlalchemy.engine.url import make_url +from superset.db_engine_specs.postgres import PostgresEngineSpec as spec from superset.exceptions import SupersetSecurityException from superset.utils.core import GenericDataType from tests.unit_tests.db_engine_specs.utils import ( @@ -53,8 +54,6 @@ def test_convert_dttm( expected_result: Optional[str], dttm: datetime, # noqa: F811 ) -> None: - from superset.db_engine_specs.postgres import PostgresEngineSpec as spec - assert_convert_dttm(spec, target_type, expected_result, dttm) @@ -91,8 +90,6 @@ def test_get_column_spec( generic_type: GenericDataType, is_dttm: bool, ) -> None: - from superset.db_engine_specs.postgres import PostgresEngineSpec as spec - assert_column_spec(spec, native_type, sqla_type, attrs, generic_type, is_dttm) @@ -100,17 +97,16 @@ def test_get_schema_from_engine_params() -> None: """ Test the ``get_schema_from_engine_params`` method. """ - from superset.db_engine_specs.postgres import PostgresEngineSpec assert ( - PostgresEngineSpec.get_schema_from_engine_params( + spec.get_schema_from_engine_params( make_url("postgresql://user:password@host/db1"), {} ) is None ) assert ( - PostgresEngineSpec.get_schema_from_engine_params( + spec.get_schema_from_engine_params( make_url("postgresql://user:password@host/db1"), {"options": "-csearch_path=secret"}, ) @@ -118,7 +114,7 @@ def test_get_schema_from_engine_params() -> None: ) assert ( - PostgresEngineSpec.get_schema_from_engine_params( + spec.get_schema_from_engine_params( make_url("postgresql://user:password@host/db1"), {"options": "-c search_path = secret -cfoo=bar -c debug"}, ) @@ -126,7 +122,7 @@ def test_get_schema_from_engine_params() -> None: ) with pytest.raises(Exception) as excinfo: - PostgresEngineSpec.get_schema_from_engine_params( + spec.get_schema_from_engine_params( make_url("postgresql://user:password@host/db1"), {"options": "-csearch_path=secret,public"}, ) @@ -141,28 +137,23 @@ def test_get_prequeries(mocker: MockerFixture) -> None: """ Test the ``get_prequeries`` method. """ - from superset.db_engine_specs.postgres import PostgresEngineSpec - database = mocker.MagicMock() - assert PostgresEngineSpec.get_prequeries(database) == [] - assert PostgresEngineSpec.get_prequeries(database, schema="test") == [ - 'set search_path = "test"' - ] + assert spec.get_prequeries(database) == [] + assert spec.get_prequeries(database, schema="test") == ['set search_path = "test"'] def test_get_default_schema_for_query(mocker: MockerFixture) -> None: """ Test the ``get_default_schema_for_query`` method. """ - from superset.db_engine_specs.postgres import PostgresEngineSpec database = mocker.MagicMock() query = mocker.MagicMock() query.sql = "SELECT * FROM some_table" query.schema = "foo" - assert PostgresEngineSpec.get_default_schema_for_query(database, query) == "foo" + assert spec.get_default_schema_for_query(database, query) == "foo" query.sql = """ set @@ -172,7 +163,7 @@ def test_get_default_schema_for_query(mocker: MockerFixture) -> None: SELECT * FROM some_table; """ with pytest.raises(SupersetSecurityException) as excinfo: - PostgresEngineSpec.get_default_schema_for_query(database, query) + spec.get_default_schema_for_query(database, query) assert ( str(excinfo.value) == "Users are not allowed to set a search path for security reasons." @@ -185,9 +176,8 @@ def test_adjust_engine_params() -> None: The method can be used to adjust the catalog (database) dynamically. """ - from superset.db_engine_specs.postgres import PostgresEngineSpec - adjusted = PostgresEngineSpec.adjust_engine_params( + adjusted = spec.adjust_engine_params( make_url("postgresql://user:password@host:5432/dev"), {}, catalog="prod", @@ -199,11 +189,57 @@ def test_get_default_catalog() -> None: """ Test `get_default_catalog`. """ - from superset.db_engine_specs.postgres import PostgresEngineSpec from superset.models.core import Database database = Database( database_name="postgres", sqlalchemy_uri="postgresql://user:password@host:5432/dev", ) - assert PostgresEngineSpec.get_default_catalog(database) == "dev" + assert spec.get_default_catalog(database) == "dev" + + +@pytest.mark.parametrize( + "time_grain,expected_result", + [ + ("PT1S", "DATE_TRUNC('second', col)"), + ( + "PT5S", + "DATE_TRUNC('minute', col) + INTERVAL '5 seconds' * FLOOR(EXTRACT(SECOND FROM col) / 5)", + ), + ( + "PT30S", + "DATE_TRUNC('minute', col) + INTERVAL '30 seconds' * FLOOR(EXTRACT(SECOND FROM col) / 30)", + ), + ("PT1M", "DATE_TRUNC('minute', col)"), + ( + "PT5M", + "DATE_TRUNC('hour', col) + INTERVAL '5 minutes' * FLOOR(EXTRACT(MINUTE FROM col) / 5)", + ), + ( + "PT10M", + "DATE_TRUNC('hour', col) + INTERVAL '10 minutes' * FLOOR(EXTRACT(MINUTE FROM col) / 10)", + ), + ( + "PT15M", + "DATE_TRUNC('hour', col) + INTERVAL '15 minutes' * FLOOR(EXTRACT(MINUTE FROM col) / 15)", + ), + ( + "PT30M", + "DATE_TRUNC('hour', col) + INTERVAL '30 minutes' * FLOOR(EXTRACT(MINUTE FROM col) / 30)", + ), + ("PT1H", "DATE_TRUNC('hour', col)"), + ("P1D", "DATE_TRUNC('day', col)"), + ("P1W", "DATE_TRUNC('week', col)"), + ("P1M", "DATE_TRUNC('month', col)"), + ("P3M", "DATE_TRUNC('quarter', col)"), + ("P1Y", "DATE_TRUNC('year', col)"), + ], +) +def test_timegrain_expressions(time_grain: str, expected_result: str) -> None: + """ + DB Eng Specs (postgres): Test time grain expressions + """ + actual = str( + spec.get_timestamp_expr(col=column("col"), pdf=None, time_grain=time_grain) + ) + assert actual == expected_result