diff --git a/superset/connectors/druid/models.py b/superset/connectors/druid/models.py index b42d66e3e..57d9d041c 100644 --- a/superset/connectors/druid/models.py +++ b/superset/connectors/druid/models.py @@ -63,6 +63,7 @@ try: MapLookupExtraction, RegexExtraction, RegisteredLookupExtraction, + TimeFormatExtraction, ) from pydruid.utils.filters import Bound, Dimension, Filter from pydruid.utils.having import Aggregation, Having @@ -1440,6 +1441,10 @@ class DruidDatasource(Model, BaseDatasource): extraction_fn = RegexExtraction(fn["expr"]) elif ext_type == "registeredLookup": extraction_fn = RegisteredLookupExtraction(fn.get("lookup")) + elif ext_type == "timeFormat": + extraction_fn = TimeFormatExtraction( + fn.get("format"), fn.get("locale"), fn.get("timeZone") + ) else: raise Exception(_("Unsupported extraction function: " + ext_type)) return (col, extraction_fn) diff --git a/tests/druid_func_tests.py b/tests/druid_func_tests.py index c13a65297..c1ac6a908 100644 --- a/tests/druid_func_tests.py +++ b/tests/druid_func_tests.py @@ -29,6 +29,7 @@ try: MapLookupExtraction, RegexExtraction, RegisteredLookupExtraction, + TimeFormatExtraction, ) import pydruid.utils.postaggregator as postaggs except ImportError: @@ -136,6 +137,33 @@ class DruidFuncTestCase(SupersetTestCase): self.assertEqual(dim_ext_fn["type"], f.extraction_function.extraction_type) self.assertEqual(dim_ext_fn["lookup"], f.extraction_function._lookup) + @unittest.skipUnless( + SupersetTestCase.is_module_installed("pydruid"), "pydruid not installed" + ) + def test_get_filters_extraction_fn_time_format(self): + filters = [{"col": "dayOfMonth", "val": ["1", "20"], "op": "in"}] + dimension_spec = { + "type": "extraction", + "dimension": "__time", + "outputName": "dayOfMonth", + "extractionFn": { + "type": "timeFormat", + "format": "d", + "timeZone": "Asia/Kolkata", + "locale": "en", + }, + } + spec_json = json.dumps(dimension_spec) + col = DruidColumn(column_name="dayOfMonth", dimension_spec_json=spec_json) + column_dict = {"dayOfMonth": col} + f = DruidDatasource.get_filters(filters, [], column_dict) + assert isinstance(f.extraction_function, TimeFormatExtraction) + dim_ext_fn = dimension_spec["extractionFn"] + self.assertEqual(dim_ext_fn["type"], f.extraction_function.extraction_type) + self.assertEqual(dim_ext_fn["format"], f.extraction_function._format) + self.assertEqual(dim_ext_fn["timeZone"], f.extraction_function._time_zone) + self.assertEqual(dim_ext_fn["locale"], f.extraction_function._locale) + @unittest.skipUnless( SupersetTestCase.is_module_installed("pydruid"), "pydruid not installed" )