Merge pull request #332 from jbothma/schema-facet

Add schema facet and option to specify which facets are included in the response

tests/unit/test_search.py

@@ -132,6 +132,39 @@ def test_search_facet_topics():
     assert "sanction" in names, names
 
 
+def test_search_facet_schema():
+    res = client.get("/search/default?schema=Address")
+    assert res.status_code == 200, res
+    addresses = res.json()["total"]["value"]
+    assert addresses > 0, addresses
+
+    res = client.get("/search/default?facets=schema")
+    assert res.status_code == 200, res
+    schemata = res.json()["facets"]["schema"]
+    names = [c["name"] for c in schemata["values"]]
+    assert "Address" in names, names
+
+
+def test_search_facet_parameter():
+    res = client.get("/search/default")
+    assert res.status_code == 200, res
+    facets = res.json()["facets"]
+    assert len(list(facets.keys())) == 3
+    assert "schema" not in facets
+    assert "topics" in facets
+    assert "datasets" in facets
+    assert "countries" in facets
+
+    res = client.get("/search/default?facets=schema&facets=topics")
+    assert res.status_code == 200, res
+    facets = res.json()["facets"]
+    assert len(list(facets.keys())) == 2
+    assert "schema" in facets
+    assert "topics" in facets
+    assert "datasets" not in facets
+    assert "countries" not in facets
+
+
 def test_search_no_targets():
     res = client.get("/search/default?schema=LegalEntity&target=false")
     assert res.status_code == 200, res

yente/routers/search.py

@@ -1,7 +1,8 @@
-from typing import List, Optional, Union
+from typing import List, Optional, Union, Annotated
 from fastapi import APIRouter, Path, Query, Response, HTTPException
 from fastapi.responses import RedirectResponse
 from followthemoney import model
+import enum
 
 from yente import settings
 from yente.logs import get_logger
@@ -22,6 +23,16 @@
 router = APIRouter()
 
 
+DEFAULT_FACETS = ["countries", "topics", "datasets"]
+
+
+class Facets(str, enum.Enum):
+    COUNTRIES = "countries"
+    TOPICS = "topics"
+    DATASETS = "datasets"
+    SCHEMA = "schema"
+
+
 @router.get(
     "/search/{dataset}",
     summary="Simple entity search",
@@ -63,6 +74,15 @@ async def search(
     target: Optional[bool] = Query(None, title="Include only targeted entities"),
     fuzzy: bool = Query(False, title="Allow fuzzy query syntax"),
     simple: bool = Query(False, title="Use simple syntax for user-facing query boxes"),
+    facets: Annotated[
+        List[Facets],
+        Query(
+            title=(
+                "Facet counts to include in response. "
+                "`schema` facet includes descendents of schema provided as query parameter."
+            )
+        ),
+    ] = DEFAULT_FACETS,
 ) -> SearchResponse:
     """Search endpoint for matching entities based on a simple piece of text, e.g.
     a name. This can be used to implement a simple, user-facing search. For proper
@@ -95,7 +115,7 @@ async def search(
         exclude_dataset=exclude_dataset,
         changed_since=changed_since,
     )
-    aggregations = facet_aggregations([f for f in filters.keys()])
+    aggregations = facet_aggregations(facets)
     resp = await search_entities(
         query,
         limit=limit,