diff --git a/superset/db_engine_specs/elasticsearch.py b/superset/db_engine_specs/elasticsearch.py index 4f18397d8c716..163bc640a6ad2 100644 --- a/superset/db_engine_specs/elasticsearch.py +++ b/superset/db_engine_specs/elasticsearch.py @@ -40,15 +40,19 @@ class ElasticSearchEngineSpec(BaseEngineSpec): # pylint: disable=abstract-metho allows_subqueries = True allows_sql_comments = False + _date_trunc_functions = { + "DATETIME": "DATE_TRUNC", + } + _time_grain_expressions = { None: "{col}", - TimeGrain.SECOND: "HISTOGRAM({col}, INTERVAL 1 SECOND)", - TimeGrain.MINUTE: "HISTOGRAM({col}, INTERVAL 1 MINUTE)", - TimeGrain.HOUR: "HISTOGRAM({col}, INTERVAL 1 HOUR)", - TimeGrain.DAY: "HISTOGRAM({col}, INTERVAL 1 DAY)", - TimeGrain.WEEK: "DATE_TRUNC('week', {col})", - TimeGrain.MONTH: "HISTOGRAM({col}, INTERVAL 1 MONTH)", - TimeGrain.YEAR: "HISTOGRAM({col}, INTERVAL 1 YEAR)", + TimeGrain.SECOND: "{func}('second', {col})", + TimeGrain.MINUTE: "{func}('minute', {col})", + TimeGrain.HOUR: "{func}('hour', {col})", + TimeGrain.DAY: "{func}('day', {col})", + TimeGrain.WEEK: "{func}('week', {col})", + TimeGrain.MONTH: "{func}('month', {col})", + TimeGrain.YEAR: "{func}('year', {col})", } type_code_map: dict[int, str] = {} # loaded from get_datatype only if needed diff --git a/tests/integration_tests/db_engine_specs/elasticsearch_tests.py b/tests/integration_tests/db_engine_specs/elasticsearch_tests.py index 7140e10ee13cb..8b07b2ebdd785 100644 --- a/tests/integration_tests/db_engine_specs/elasticsearch_tests.py +++ b/tests/integration_tests/db_engine_specs/elasticsearch_tests.py @@ -14,27 +14,30 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +from parameterized import parameterized from sqlalchemy import column +from superset.constants import TimeGrain from superset.db_engine_specs.elasticsearch import ElasticSearchEngineSpec from tests.integration_tests.db_engine_specs.base_tests import TestDbEngineSpec class TestElasticsearchDbEngineSpec(TestDbEngineSpec): - def test_time_grain_week_expression(self): + @parameterized.expand( + [ + [TimeGrain.SECOND, "DATE_TRUNC('second', ts)"], + [TimeGrain.MINUTE, "DATE_TRUNC('minute', ts)"], + [TimeGrain.HOUR, "DATE_TRUNC('hour', ts)"], + [TimeGrain.DAY, "DATE_TRUNC('day', ts)"], + [TimeGrain.WEEK, "DATE_TRUNC('week', ts)"], + [TimeGrain.MONTH, "DATE_TRUNC('month', ts)"], + [TimeGrain.YEAR, "DATE_TRUNC('year', ts)"], + ] + ) + def test_time_grain_expressions(self, time_grain, expected_time_grain_expression): col = column("ts") - col.type = "datetime" - expected_time_grain_expression = "DATE_TRUNC('week', ts)" + col.type = "DATETIME" actual = ElasticSearchEngineSpec.get_timestamp_expr( - col=col, pdf=None, time_grain="P1W" - ) - self.assertEqual(str(actual), expected_time_grain_expression) - - def test_time_grain_hour_expression(self): - col = column("ts") - col.type = "datetime" - expected_time_grain_expression = "HISTOGRAM(ts, INTERVAL 1 HOUR)" - actual = ElasticSearchEngineSpec.get_timestamp_expr( - col=col, pdf=None, time_grain="PT1H" + col=col, pdf=None, time_grain=time_grain ) self.assertEqual(str(actual), expected_time_grain_expression)