feat: return objets counts in orgs suggest for search facets

udata/core/organization/api.py

@@ -1,7 +1,9 @@
 from datetime import UTC, datetime
+from types import SimpleNamespace
 
-from flask import make_response, redirect, request, url_for
+from flask import current_app, make_response, redirect, request, url_for
 from mongoengine.queryset.visitor import Q
+from werkzeug.datastructures import MultiDict
 
 from udata.api import API, api, errors
 from udata.api.parsers import ModelApiParser
@@ -32,6 +34,7 @@
 from udata.mongo import db
 from udata.mongo.errors import FieldValidationError
 from udata.rdf import RDF_EXTENSIONS, graph_response, negociate_content
+from udata.search import adapter_for, get_elastic_client
 
 from .api_fields import (
     invite_fields,
@@ -674,27 +677,111 @@ class FollowOrgAPI(FollowAPI):
     model = Organization
 
 
+# Models whose search index can back an organization suggest count (count_for=...).
+COUNT_FOR_MODELS = {"dataset": Dataset, "reuse": Reuse, "dataservice": Dataservice}
+# Prefix namespacing the context-search filters from the suggest's own params: the
+# suggest `q` is an organization name, while `count_filter.q` is the dataset/reuse/...
+# full-text query the count is scoped to (e.g. count_filter.tag=transport).
+COUNT_FILTER_PREFIX = "count_filter."
+# Upper bound on the number of organizations we send to ES in a single count request.
+# The candidate set comes from a name `icontains`, which is barely selective for short
+# queries (on data.gouv.fr "a" matches ~5000 of ~6300 orgs, "le" ~1500), so we always
+# keep only the top organizations by followers. The count ranking is therefore exact
+# whenever the name match already fits under this cap (typically queries of 3+ chars,
+# e.g. "data" matches ~90 orgs) and a best-effort top-by-followers otherwise.
+COUNT_CANDIDATE_POOL = 100
+
+
+def org_suggestion(org):
+    return {
+        "id": org.id,
+        "name": org.name,
+        "acronym": org.acronym,
+        "slug": org.slug,
+        "image_url": org.logo,
+        "page": org.self_web_url(),
+    }
+
+
+org_suggest_parser = suggest_parser.copy()
+org_suggest_parser.add_argument(
+    "count_for",
+    type=str,
+    choices=list(COUNT_FOR_MODELS),
+    location="args",
+    help=(
+        "Annotate each suggestion with the number of matching objects of this kind "
+        "(dataset, reuse or dataservice), scoped to the search passed via "
+        f"`{COUNT_FILTER_PREFIX}*` params (e.g. {COUNT_FILTER_PREFIX}tag=transport). "
+        "Empty organizations are pushed to the end, the follower ranking is kept."
+    ),
+)
+
+
+def organization_match_counts(count_for, organizations):
+    """Return {organization_id: matching_count} for the given candidate organizations.
+
+    Counts come from the search index of ``count_for`` (datasets/reuses/dataservices) —
+    the only place a search-scoped per-organization count exists — via an exact ES
+    `include` of the candidates. Returns an empty mapping when search is disabled.
+    """
+    if not organizations or not current_app.config["ELASTICSEARCH_URL"]:
+        return {}
+
+    adapter = adapter_for(COUNT_FOR_MODELS[count_for])
+
+    # Parse the count_filter.* params through the target model's own request parser so
+    # filters are validated/typed exactly like a real search on that model.
+    stripped = MultiDict()
+    for key in request.args:
+        if key.startswith(COUNT_FILTER_PREFIX):
+            for value in request.args.getlist(key):
+                stripped.add(key[len(COUNT_FILTER_PREFIX) :], value)
+    parser = adapter.as_request_parser(paginate=False, store_missing=False)
+    params = dict(parser.parse_args(req=SimpleNamespace(args=stripped)))
+
+    params.setdefault("q", "")
+    params["page"] = 1
+    # Ignored: count-only mode forces ES `size: 0`, so no hits are ever fetched.
+    params["page_size"] = 1
+    # The search layer owns how organizations are indexed; we only pass ids and get back
+    # a {organization_id: count} mapping.
+    params["count_organizations"] = [str(org.id) for org in organizations]
+
+    service = adapter.service_class(get_elastic_client())
+    _, _, _, facets = service.search(params)
+    return facets.get("organization_counts", {})
+
+
 @ns.route("/suggest/", endpoint="suggest_organizations")
 class OrganizationSuggestAPI(API):
     @api.doc("suggest_organizations")
-    @api.expect(suggest_parser)
+    @api.expect(org_suggest_parser)
     @api.marshal_list_with(org_suggestion_fields)
     def get(self):
         """Organizations suggest endpoint using mongoDB contains"""
-        args = suggest_parser.parse_args()
+        args = org_suggest_parser.parse_args()
         orgs = Organization.objects(
             Q(name__icontains=args["q"]) | Q(acronym__icontains=args["q"]), deleted=None
-        )
+        ).order_by(SUGGEST_SORTING)
+
+        if not args["count_for"]:
+            return [org_suggestion(org) for org in orgs.limit(args["size"])]
+
+        # Count the top organizations by followers matching the name. We cap the pool
+        # because a name `icontains` is barely selective on short queries (see
+        # COUNT_CANDIDATE_POOL); the counted set stays exact for queries narrow enough
+        # to fit under the cap.
+        candidates = list(orgs.limit(COUNT_CANDIDATE_POOL))
+        counts = organization_match_counts(args["count_for"], candidates)
+
+        # Keep the follower order (stable sort), only pushing empty organizations to the
+        # end. We never sort by count: a huge low-quality organization must not outrank a
+        # well-followed one.
+        candidates.sort(key=lambda org: counts.get(str(org.id), 0) == 0)
         return [
-            {
-                "id": org.id,
-                "name": org.name,
-                "acronym": org.acronym,
-                "slug": org.slug,
-                "image_url": org.logo,
-                "page": org.self_web_url(),
-            }
-            for org in orgs.order_by(SUGGEST_SORTING).limit(args["size"])
+            {**org_suggestion(org), "matching_count": counts.get(str(org.id), 0)}
+            for org in candidates[: args["size"]]
         ]
 
 

udata/core/organization/api_fields.py

@@ -99,5 +99,13 @@
             size=BIGGEST_LOGO_SIZE, description="The organization logo URL", readonly=True
         ),
         "page": fields.String(description="The organization web page URL", readonly=True),
+        "matching_count": fields.Integer(
+            description=(
+                "Number of objects (set by the `count_for` param) owned by this "
+                "organization and matching the `count_filter.*` search. Null unless "
+                "`count_for` was requested."
+            ),
+            readonly=True,
+        ),
     },
 )

udata/tests/api/test_organizations_api.py

@@ -1058,6 +1058,28 @@ def test_unfollow_org(self):
         assert Follow.objects.following(user).count() == 0
         assert Follow.objects.followers(user).count() == 0
 
+    def test_suggest_organizations_without_count_for_has_null_count(self):
+        """Without count_for the endpoint stays the plain Mongo suggest."""
+        OrganizationFactory(name="plain-org")
+        response = self.get(url_for("api.suggest_organizations", q="plain", size=5))
+        assert200(response)
+        assert response.json[0]["name"] == "plain-org"
+        assert response.json[0]["matching_count"] is None
+
+    def test_suggest_organizations_count_for_without_search_service(self):
+        """With count_for but no search service, suggestions are returned without crashing."""
+        OrganizationFactory(name="degraded-org")
+        response = self.get(
+            url_for("api.suggest_organizations", q="degraded", size=5, count_for="dataset")
+        )
+        assert200(response)
+        assert response.json[0]["name"] == "degraded-org"
+
+    def test_suggest_organizations_invalid_count_for(self):
+        """count_for only accepts known model kinds."""
+        response = self.get(url_for("api.suggest_organizations", q="x", size=5, count_for="banana"))
+        assert response.status_code == 400
+
     def test_suggest_organizations_api(self):
         """It should suggest organizations"""
         for i in range(3):

udata/tests/search/test_search_integration.py

@@ -1,6 +1,7 @@
 import time
 
 import pytest
+from flask import url_for
 
 from udata.core.access_type.constants import AccessType
 from udata.core.dataservices.factories import DataserviceFactory
@@ -658,3 +659,60 @@ def test_dataservice_filter_by_tags(self):
         titles = [d["title"] for d in response.json["data"]]
         assert "DS tagged" in titles
         assert "DS other" not in titles
+
+    def test_suggest_organizations_with_dataset_count(self):
+        """count_for=dataset annotates each suggestion and pushes empty orgs to the end."""
+        org_two = OrganizationFactory(name="Count Alpha")
+        org_one = OrganizationFactory(name="Count Beta")
+        OrganizationFactory(name="Count Empty")
+        DatasetFactory(organization=org_two)
+        DatasetFactory(organization=org_two)
+        DatasetFactory(organization=org_one)
+
+        time.sleep(1)
+
+        response = self.get(
+            url_for("api.suggest_organizations", q="Count", size=10, count_for="dataset")
+        )
+        self.assert200(response)
+        counts = {org["name"]: org["matching_count"] for org in response.json}
+        assert counts == {"Count Alpha": 2, "Count Beta": 1, "Count Empty": 0}
+        # The empty organization is demoted to the end, the others keep their ranking.
+        assert response.json[-1]["name"] == "Count Empty"
+
+    def test_suggest_organizations_count_scoped_by_filter(self):
+        """The count respects the count_filter.* search context."""
+        org = OrganizationFactory(name="Scoped Org")
+        DatasetFactory(organization=org, tags=["transport"])
+        DatasetFactory(organization=org, tags=["sante"])
+
+        time.sleep(1)
+
+        response = self.get(
+            url_for(
+                "api.suggest_organizations",
+                q="Scoped",
+                size=10,
+                count_for="dataset",
+                **{"count_filter.tag": "transport"},
+            )
+        )
+        self.assert200(response)
+        assert response.json[0]["name"] == "Scoped Org"
+        assert response.json[0]["matching_count"] == 1
+
+    def test_suggest_organizations_dataservice_count(self):
+        """count_for=dataservice replaces the old "orgs with at least one API" flag."""
+        org_api = OrganizationFactory(name="Api Provider")
+        OrganizationFactory(name="Api Nothing")
+        DataserviceFactory(organization=org_api)
+
+        time.sleep(1)
+
+        response = self.get(
+            url_for("api.suggest_organizations", q="Api", size=10, count_for="dataservice")
+        )
+        self.assert200(response)
+        counts = {org["name"]: org["matching_count"] for org in response.json}
+        assert counts == {"Api Provider": 1, "Api Nothing": 0}
+        assert response.json[-1]["name"] == "Api Nothing"

udata/tests/search/test_services.py

@@ -0,0 +1,39 @@
+from udata_search_service.services import DatasetService
+
+
+class FakeElasticClient:
+    """Records the last query_datasets call so we can assert on the plumbing."""
+
+    def __init__(self):
+        self.last_call = None
+
+    def index_dataset(self, *args, **kwargs):
+        pass
+
+    def find_one_dataset(self, *args, **kwargs):
+        pass
+
+    def delete_one_dataset(self, *args, **kwargs):
+        pass
+
+    def query_datasets(self, search_text, offset, page_size, filters, sort=None, **kwargs):
+        self.last_call = {"filters": filters, "kwargs": kwargs}
+        return 0, [], {}
+
+
+def base_filters(**extra):
+    return {"q": "", "page": 1, "page_size": 20, "sort": None, **extra}
+
+
+def test_count_organizations_forwarded_when_set():
+    client = FakeElasticClient()
+    DatasetService(client).search(base_filters(count_organizations=["org-id-1", "org-id-2"]))
+    assert client.last_call["kwargs"]["count_organizations"] == ["org-id-1", "org-id-2"]
+    # It must not leak into the filters sent to ES as a regular filter.
+    assert "count_organizations" not in client.last_call["filters"]
+
+
+def test_count_organizations_not_forwarded_when_absent():
+    client = FakeElasticClient()
+    DatasetService(client).search(base_filters())
+    assert "count_organizations" not in client.last_call["kwargs"]