From 08e89e4cabbecf5748f31b503e384d75bfed01a5 Mon Sep 17 00:00:00 2001
From: SongYunSeop <dbstjq91@gmail.com>
Date: Wed, 22 Apr 2020 11:40:28 +0900
Subject: [PATCH] Enable extractionFn on Search Filter and Like Filter

---
 pydruid/utils/filters.py    |  2 ++
 tests/utils/test_filters.py | 37 ++++++++++++++++++++++++++++++++++++-
 2 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/pydruid/utils/filters.py b/pydruid/utils/filters.py
index 3028c645..3e1c5d2a 100644
--- a/pydruid/utils/filters.py
+++ b/pydruid/utils/filters.py
@@ -32,6 +32,8 @@ class Filter:
         "bound",
         "interval",
         "extraction",
+        "like",
+        "search",
     )
 
     def __init__(self, extraction_function=None, ordering="lexicographic", **args):
diff --git a/tests/utils/test_filters.py b/tests/utils/test_filters.py
index 2b419543..046afc34 100644
--- a/tests/utils/test_filters.py
+++ b/tests/utils/test_filters.py
@@ -275,9 +275,44 @@ def test_search_filter(self):
                     'query': {'type': 'contains', 'caseSensitive': 'true', 'value': 'val'}}
         assert actual == expected
 
+    def test_search_filter_with_extraction_function(self):
+        extraction_fn = dimensions.RegexExtraction("([a-b])")
+        actual = filters.Filter.build_filter(
+            filters.Filter(
+                type="search",
+                dimension="dim",
+                value="val",
+                extraction_function=extraction_fn,
+            )
+        )
+        expected = {
+            "type": "search",
+            "dimension": "dim",
+            "query": {"type": "contains", "caseSensitive": "false", "value": "val"},
+            "extractionFn": {"type": "regex", "expr": "([a-b])"},
+        }
+        assert actual == expected
+
     def test_like_filter(self):
         actual = filters.Filter.build_filter(
             filters.Filter(type="like", dimension="dim", pattern="%val%"))
         expected = {'type': 'like', 'dimension': 'dim', 'pattern': '%val%'}
         assert actual == expected
-        
+
+    def test_like_filter_with_extraction_function(self):
+        extraction_fn = dimensions.RegexExtraction("([a-b])")
+        actual = filters.Filter.build_filter(
+            filters.Filter(
+                type="like",
+                dimension="dim",
+                pattern="%val%",
+                extraction_function=extraction_fn,
+            )
+        )
+        expected = {
+            "type": "like",
+            "dimension": "dim",
+            "pattern": "%val%",
+            "extractionFn": {"type": "regex", "expr": "([a-b])"},
+        }
+        assert actual == expected