fix(db_engine_specs): add a few missing time grains to Postgres spec (#30325)
This commit is contained in:
Sam Firke
GitHub
parent
038ef32454
commit
ac66ae813f
|
|
@ -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})",
|
||||
|
|
|
|||
|
|
@ -20,10 +20,11 @@ from typing import Any, Optional
|
|||
|
||||
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 @@ search_path -- another one
|
|||
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
|
||||
|
|
|
|||
Loading…
Reference in New Issue