diff --git a/superset/connectors/druid/models.py b/superset/connectors/druid/models.py index f204a90c9..b908056aa 100644 --- a/superset/connectors/druid/models.py +++ b/superset/connectors/druid/models.py @@ -35,7 +35,11 @@ import pandas try: from pydruid.client import PyDruid from pydruid.utils.aggregators import count - from pydruid.utils.dimensions import MapLookupExtraction, RegexExtraction + from pydruid.utils.dimensions import ( + MapLookupExtraction, + RegexExtraction, + RegisteredLookupExtraction, + ) from pydruid.utils.filters import Dimension, Filter from pydruid.utils.having import Aggregation from pydruid.utils.postaggregator import ( @@ -1402,6 +1406,8 @@ class DruidDatasource(Model, BaseDatasource): ) elif ext_type == "regex": extraction_fn = RegexExtraction(fn["expr"]) + elif ext_type == "registeredLookup": + extraction_fn = RegisteredLookupExtraction(fn.get("lookup")) 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 c4d6ab92c..460595dd1 100644 --- a/tests/druid_func_tests.py +++ b/tests/druid_func_tests.py @@ -19,7 +19,11 @@ import unittest from unittest.mock import Mock try: - from pydruid.utils.dimensions import MapLookupExtraction, RegexExtraction + from pydruid.utils.dimensions import ( + MapLookupExtraction, + RegexExtraction, + RegisteredLookupExtraction, + ) import pydruid.utils.postaggregator as postaggs except ImportError: pass @@ -110,6 +114,27 @@ class DruidFuncTestCase(unittest.TestCase): f_ext_fn = f.extraction_function self.assertEqual(dim_ext_fn["expr"], f_ext_fn._expr) + @unittest.skipUnless( + SupersetTestCase.is_module_installed("pydruid"), "pydruid not installed" + ) + def test_get_filters_extraction_fn_registered_lookup_extraction(self): + filters = [{"col": "country", "val": ["Spain"], "op": "in"}] + dimension_spec = { + "type": "extraction", + "dimension": "country_name", + "outputName": "country", + "outputType": "STRING", + "extractionFn": {"type": "registeredLookup", "lookup": "country_name"}, + } + spec_json = json.dumps(dimension_spec) + col = DruidColumn(column_name="country", dimension_spec_json=spec_json) + column_dict = {"country": col} + f = DruidDatasource.get_filters(filters, [], column_dict) + assert isinstance(f.extraction_function, RegisteredLookupExtraction) + dim_ext_fn = dimension_spec["extractionFn"] + 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" )