feat(search): configurable facets size

udata/core/dataservices/search.py

@@ -83,6 +83,7 @@ class DataserviceSearch(ModelSearchAdapter):
     model = Dataservice
     service_class = DataserviceService
     consumer_class = DataserviceConsumer
+    configurable_size_facets = ["organization_id_with_name"]
 
     sorts = {"created": "created_at", "views": "views", "followers": "followers"}
 

udata/core/dataset/search.py

@@ -34,6 +34,7 @@ class DatasetSearch(ModelSearchAdapter):
     model = Dataset
     service_class = DatasetService
     consumer_class = DatasetConsumer
+    configurable_size_facets = ["organization_id_with_name"]
 
     sorts = {
         "created": "created_at_internal",

udata/core/reuse/search.py

@@ -25,6 +25,7 @@ class ReuseSearch(ModelSearchAdapter):
     model = Reuse
     service_class = ReuseService
     consumer_class = ReuseConsumer
+    configurable_size_facets = ["organization_id_with_name"]
 
     sorts = {
         "created": "created_at",

udata/core/topic/search.py

@@ -22,6 +22,7 @@ class TopicSearch(ModelSearchAdapter):
     model = Topic
     service_class = TopicService
     consumer_class = TopicConsumer
+    configurable_size_facets = ["organization_id_with_name"]
 
     sorts = {
         "name": "name",

udata/search/adapter.py

@@ -13,6 +13,7 @@ class ModelSearchAdapter:
     model = None
     sorts = None
     filters = {}
+    configurable_size_facets = []
     service_class = None
     consumer_class = None
 
@@ -57,6 +58,13 @@ def as_request_parser(cls, paginate=True, store_missing: bool = True):
             parser.add_argument(
                 "page_size", type=int, location="args", default=20, help="The page size"
             )
+        for facet_name in cls.configurable_size_facets:
+            parser.add_argument(
+                f"facet_size__{facet_name}",
+                type=int,
+                location="args",
+                help=f"Number of {facet_name} facet values to return",
+            )
         return parser
 
     @classmethod

udata/search/query.py

@@ -11,9 +11,22 @@
 DEFAULT_PAGE_SIZE = 20
 # Elasticsearch default max_result_window is 10000
 ES_MAX_RESULT_WINDOW = 10000
+DEFAULT_MAX_FACET_SIZE = 500
 log = logging.getLogger(__name__)
 
 
+def parse_facet_size(key, raw, max_facet_size):
+    try:
+        value = int(raw)
+    except (ValueError, TypeError):
+        abort(400, f"Invalid value for {key}: {raw!r} is not an integer.")
+    if value < 1:
+        abort(400, f"Invalid value for {key}: must be at least 1.")
+    if value > max_facet_size:
+        abort(400, f"Invalid value for {key}: must be at most {max_facet_size}.")
+    return value
+
+
 class SearchQuery:
     adapter = None
     model = None
@@ -32,6 +45,11 @@ def __init__(self, params):
             )
         self._query = params.pop("q", "")
         self.sort = params.pop("sort", None)
+        max_facet_size = current_app.config.get("MAX_FACET_SIZE", DEFAULT_MAX_FACET_SIZE)
+        self._facet_sizes = {}
+        for key in [k for k in list(params.keys()) if k.startswith("facet_size__")]:
+            facet_name = key[len("facet_size__") :]
+            self._facet_sizes[facet_name] = parse_facet_size(key, params.pop(key), max_facet_size)
         self._filters = {}
         self.extract_filters(params)
 
@@ -81,6 +99,7 @@ def to_search_params(self):
             "page": self.page,
             "page_size": self.page_size,
             "sort": self.sort,
+            "facet_sizes": self._facet_sizes,
         }
         params.update(self._filters)
         return params

udata/tests/search/test_adapter.py

@@ -560,6 +560,22 @@ def test_serialize_includes_producer_type_user(self):
         assert serialized["producer_type"] == [USER]
 
 
+class ConfigurableSizeFacetsTest(APITestCase):
+    def test_facet_size_params_in_request_parser(self):
+        for adapter in [DatasetSearch, ReuseSearch, DataserviceSearch]:
+            parser = adapter.as_request_parser()
+            arg_names = [arg.name for arg in parser.args]
+            assert "facet_size__organization_id_with_name" in arg_names, (
+                f"{adapter.__name__} parser is missing facet_size__organization_id_with_name — "
+                f"it would be silently dropped from API requests"
+            )
+
+    def test_facet_size_param_is_int(self):
+        parser = DatasetSearch.as_request_parser()
+        arg = next(a for a in parser.args if a.name == "facet_size__organization_id_with_name")
+        assert arg.type is int
+
+
 class DataserviceSearchAdapterTest(APITestCase):
     def test_serialize_includes_access_type(self):
         """Test that DataserviceSearch.serialize includes access_type in the serialized document"""

udata/tests/search/test_query.py

@@ -1,4 +1,9 @@
-from udata.search.query import DEFAULT_PAGE_SIZE, SearchQuery
+from udata.search.query import (
+    DEFAULT_MAX_FACET_SIZE,
+    DEFAULT_PAGE_SIZE,
+    SearchQuery,
+    parse_facet_size,
+)
 from udata.tests.api import APITestCase
 
 
@@ -55,3 +60,84 @@ def test_search_query_to_search_params(self):
         assert params["sort"] == "-created"
         assert params["organization"] == "534fff81a3a7292c64a77e5c"
         assert params["tag"] == ["tag-1", "tag-2"]
+
+    def test_facet_sizes_default_empty(self):
+        search_query = SearchQuery(params={})
+        assert search_query._facet_sizes == {}
+
+    def test_facet_size_param_parsed(self):
+        query = {"facet_size__organization_id_with_name": "200"}
+        search_query = SearchQuery(params=query)
+        assert search_query._facet_sizes == {"organization_id_with_name": 200}
+
+    def test_multiple_facet_size_params_parsed(self):
+        query = {"facet_size__organization_id_with_name": "200", "facet_size__tag": "100"}
+        search_query = SearchQuery(params=query)
+        assert search_query._facet_sizes == {"organization_id_with_name": 200, "tag": 100}
+
+    def test_facet_size_params_not_in_filters(self):
+        query = {"facet_size__organization_id_with_name": "200", "tag": "transport"}
+        search_query = SearchQuery(params=query)
+        assert "facet_size__organization_id_with_name" not in search_query._filters
+        assert search_query._filters == {"tag": "transport"}
+
+    def test_facet_sizes_included_in_search_params(self):
+        query = {"facet_size__organization_id_with_name": "200", "q": "test"}
+        search_query = SearchQuery(params=query)
+        params = search_query.to_search_params()
+        assert params["facet_sizes"] == {"organization_id_with_name": 200}
+
+    def test_facet_size_zero_raises_400(self):
+        with self.assertRaises(Exception) as ctx:
+            SearchQuery(params={"facet_size__tag": "0"})
+        assert ctx.exception.code == 400
+
+    def test_facet_size_negative_raises_400(self):
+        with self.assertRaises(Exception) as ctx:
+            SearchQuery(params={"facet_size__tag": "-1"})
+        assert ctx.exception.code == 400
+
+    def test_facet_size_non_integer_raises_400(self):
+        with self.assertRaises(Exception) as ctx:
+            SearchQuery(params={"facet_size__tag": "big"})
+        assert ctx.exception.code == 400
+
+    def test_facet_size_exceeds_max_raises_400(self):
+        with self.assertRaises(Exception) as ctx:
+            SearchQuery(params={"facet_size__tag": str(DEFAULT_MAX_FACET_SIZE + 1)})
+        assert ctx.exception.code == 400
+
+    def test_facet_size_at_max_is_valid(self):
+        search_query = SearchQuery(params={"facet_size__tag": str(DEFAULT_MAX_FACET_SIZE)})
+        assert search_query._facet_sizes["tag"] == DEFAULT_MAX_FACET_SIZE
+
+
+class ParseFacetSizeTest(APITestCase):
+    def test_valid_integer(self):
+        assert parse_facet_size("facet_size__tag", "50", 500) == 50
+
+    def test_non_integer_raises_400(self):
+        with self.assertRaises(Exception) as ctx:
+            parse_facet_size("facet_size__tag", "abc", 500)
+        assert ctx.exception.code == 400
+
+    def test_zero_raises_400(self):
+        with self.assertRaises(Exception) as ctx:
+            parse_facet_size("facet_size__tag", "0", 500)
+        assert ctx.exception.code == 400
+
+    def test_negative_raises_400(self):
+        with self.assertRaises(Exception) as ctx:
+            parse_facet_size("facet_size__tag", "-5", 500)
+        assert ctx.exception.code == 400
+
+    def test_exceeds_max_raises_400(self):
+        with self.assertRaises(Exception) as ctx:
+            parse_facet_size("facet_size__tag", "501", 500)
+        assert ctx.exception.code == 400
+
+    def test_at_max_is_valid(self):
+        assert parse_facet_size("facet_size__tag", "500", 500) == 500
+
+    def test_at_min_is_valid(self):
+        assert parse_facet_size("facet_size__tag", "1", 500) == 1

udata/tests/search/test_services.py

@@ -0,0 +1,38 @@
+from unittest.mock import MagicMock
+
+from udata_search_service.services import DatasetService
+
+
+def make_service():
+    mock_client = MagicMock()
+    mock_client.query_datasets.return_value = (0, [], {})
+    return DatasetService(mock_client), mock_client
+
+
+def base_filters():
+    return {"q": "", "page": 1, "page_size": 20, "sort": None}
+
+
+def test_facet_sizes_passed_to_client_query():
+    service, mock_client = make_service()
+    filters = {**base_filters(), "facet_sizes": {"organization_id_with_name": 200}}
+    service.search(filters)
+    _, kwargs = mock_client.query_datasets.call_args
+    assert kwargs["facet_sizes"] == {"organization_id_with_name": 200}
+
+
+def test_empty_facet_sizes_when_not_provided():
+    service, mock_client = make_service()
+    service.search(base_filters())
+    _, kwargs = mock_client.query_datasets.call_args
+    assert kwargs["facet_sizes"] == {}
+
+
+def test_facet_sizes_not_passed_as_filter():
+    service, mock_client = make_service()
+    filters = {**base_filters(), "facet_sizes": {"tag": 100}}
+    service.search(filters)
+    args, _ = mock_client.query_datasets.call_args
+    # 4th positional arg is the filters dict
+    filters_arg = args[3]
+    assert "facet_sizes" not in filters_arg