Skip to content

Commit

Permalink
fix(db_engine_specs): add a few missing time grains to Postgres spec (a…
Browse files Browse the repository at this point in the history 
…pache#30325)
  • Loading branch information
@sfirke
sfirke authored Sep 20, 2024
1 parent 038ef32 commit ac66ae8
Show file tree
Hide file tree
Showing 2 changed files with 65 additions and 23 deletions.
6 changes: 6 additions & 0 deletions superset/db_engine_specs/postgres.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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})",
Expand Down
82 changes: 59 additions & 23 deletions tests/unit_tests/db_engine_specs/test_postgres.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down Expand Up @@ -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)


Expand Down Expand Up @@ -91,42 +90,39 @@ 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)


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

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

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

0 comments on commit ac66ae8

Please sign in to comment.