Support timeFormat extraction function column for Druid in Filter (#8728)

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)

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"
     )