✨(backends) add support for queries with multiple authorities

CHANGELOG.md

@@ -8,6 +8,10 @@ and this project adheres to
 
 ## [Unreleased]
 
+### Added
+
+- Backends: (internal) add support for multiple authority queries
+
 ### Removed
 
 - Drop support for Python 3.8

src/ralph/backends/lrs/base.py

@@ -117,7 +117,9 @@ class RalphStatementsQuery(LRSStatementsQuery):
     agent: Optional[AgentParameters] = AgentParameters.model_construct()
     search_after: Optional[str] = None
     pit_id: Optional[str] = None
-    authority: Optional[AgentParameters] = AgentParameters.model_construct()
+    authority: Optional[Union[AgentParameters, list[AgentParameters]]] = (
+        AgentParameters.model_construct()
+    )
     ignore_order: Optional[bool] = None
 
 

src/ralph/backends/lrs/clickhouse.py

@@ -1,7 +1,8 @@
 """ClickHouse LRS backend for Ralph."""
 
 import logging
-from typing import Generator, Iterator, List, Optional
+from functools import reduce
+from typing import Generator, Iterator, List, Optional, Union
 
 from pydantic_settings import SettingsConfigDict
 
@@ -162,48 +163,80 @@ def chunk_id_list(chunk_size: int = self.settings.IDS_CHUNK_SIZE) -> Generator:
     def _add_agent_filters(
         ch_params: dict,
         where: list,
-        agent_params: AgentParameters,
+        agent_params: Union[AgentParameters, list[AgentParameters]],
         target_field: str,
     ) -> None:
         """Add filters relative to agents to `where`."""
-        if not agent_params:
-            return
 
-        if not isinstance(agent_params, dict):
-            agent_params = agent_params.model_dump()
+        def _get_agent_filters(
+            _params: AgentParameters, idx: Optional[int] = None
+        ) -> Union[tuple[list[str], dict], None]:
+            if not _params:
+                return None
 
-        if agent_params.get("mbox"):
-            ch_params[f"{target_field}__mbox"] = agent_params.get("mbox")
-            where.append(
-                f"JSONExtractString(event, '{target_field}', 'mbox') = "
-                f"{{{target_field}__mbox:String}}"
-            )
-        elif agent_params.get("mbox_sha1sum"):
-            ch_params[f"{target_field}__mbox_sha1sum"] = agent_params.get(
-                "mbox_sha1sum"
-            )
-            where.append(
-                f"JSONExtractString(event, '{target_field}', 'mbox_sha1sum') = "
-                f"{{{target_field}__mbox_sha1sum:String}}"
-            )
-        elif agent_params.get("openid"):
-            ch_params[f"{target_field}__openid"] = agent_params.get("openid")
-            where.append(
-                f"JSONExtractString(event, '{target_field}', 'openid') = "
-                f"{{{target_field}__openid:String}}"
-            )
-        elif agent_params.get("account__name"):
-            ch_params[f"{target_field}__account__name"] = agent_params.get(
-                "account__name"
-            )
-            where.append(
-                f"JSONExtractString(event, '{target_field}', 'account', 'name') = "
-                f"{{{target_field}__account__name:String}}"
-            )
-            ch_params[f"{target_field}__account__home_page"] = agent_params.get(
-                "account__home_page"
-            )
-            where.append(
-                f"JSONExtractString(event, '{target_field}', 'account', 'homePage') = "
-                f"{{{target_field}__account__home_page:String}}"
-            )
+            if not isinstance(_params, dict):
+                _params = _params.model_dump()
+
+            target_param = f"{target_field}_{idx}" if idx is not None else target_field
+
+            if _params.get("mbox"):
+                return (
+                    [
+                        f"JSONExtractString(event, '{target_field}', 'mbox') = "
+                        f"{{{target_param}__mbox:String}}"
+                    ],
+                    {f"{target_param}__mbox": _params.get("mbox")},
+                )
+            elif _params.get("mbox_sha1sum"):
+                return (
+                    [
+                        f"JSONExtractString(event, '{target_field}', 'mbox_sha1sum') = "
+                        f"{{{target_param}__mbox_sha1sum:String}}"
+                    ],
+                    {f"{target_param}__mbox_sha1sum": _params.get("mbox_sha1sum")},
+                )
+            elif _params.get("openid"):
+                return (
+                    [
+                        f"JSONExtractString(event, '{target_field}', 'openid') = "
+                        f"{{{target_param}__openid:String}}"
+                    ],
+                    {f"{target_param}__openid": _params.get("openid")},
+                )
+            elif _params.get("account__name"):
+                return (
+                    [
+                        f"JSONExtractString(event, '{target_field}', 'account',"
+                        f" 'name') = "
+                        f"{{{target_param}__account__name:String}}",
+                        f"JSONExtractString(event, '{target_field}', 'account',"
+                        f" 'homePage') = "
+                        f"{{{target_param}__account__home_page:String}}",
+                    ],
+                    {
+                        f"{target_param}__account__name": _params.get("account__name"),
+                        f"{target_param}__account__home_page": _params.get(
+                            "account__home_page"
+                        ),
+                    },
+                )
+            return None
+
+        if not agent_params:
+            return
+        elif not isinstance(agent_params, list):
+            filters = _get_agent_filters(agent_params)
+            if filters:
+                _where, _ch_params = filters
+                ch_params.update(_ch_params)
+                where.extend(_where)
+        else:
+            filters = [
+                _get_agent_filters(params, idx=idx)
+                for idx, params in enumerate(agent_params)
+                if params
+            ]
+            _ch_params = reduce(lambda acc, el: acc | el[1], filters, {})
+            _where = [" OR ".join([" AND ".join(el[0]) for el in filters])]
+            ch_params.update(_ch_params)
+            where.extend(_where)

src/ralph/backends/lrs/es.py

@@ -1,7 +1,7 @@
 """Elasticsearch LRS backend for Ralph."""
 
 import logging
-from typing import Iterator, List, Optional
+from typing import Iterator, List, Optional, Union
 
 from pydantic_settings import SettingsConfigDict
 
@@ -112,28 +112,51 @@ def get_query(params: RalphStatementsQuery) -> ESQuery:
 
     @staticmethod
     def _add_agent_filters(
-        es_query_filters: list, agent_params: AgentParameters, target_field: str
+        es_query_filters: list,
+        agent_params: Union[AgentParameters, list[AgentParameters]],
+        target_field: str,
     ) -> None:
         """Add filters relative to agents to `es_query_filters`."""
+
+        def _get_agent_filters(_params: AgentParameters) -> Union[dict, None]:
+            if not _params:
+                return None
+
+            if not isinstance(_params, dict):
+                _params = _params.model_dump()
+
+            if _params.get("mbox"):
+                field = f"{target_field}.mbox.keyword"
+                return {"term": {field: _params.get("mbox")}}
+            elif _params.get("mbox_sha1sum"):
+                field = f"{target_field}.mbox_sha1sum.keyword"
+                return {"term": {field: _params.get("mbox_sha1sum")}}
+            elif _params.get("openid"):
+                field = f"{target_field}.openid.keyword"
+                return {"term": {field: _params.get("openid")}}
+            elif _params.get("account__name"):
+                field_name = f"{target_field}.account.name.keyword"
+                field_homepage = f"{target_field}.account.homePage.keyword"
+                return {
+                    "bool": {
+                        "filter": [
+                            {"term": {field_name: _params.get("account__name")}},
+                            {
+                                "term": {
+                                    field_homepage: _params.get("account__home_page")
+                                }
+                            },
+                        ]
+                    }
+                }
+            return None
+
         if not agent_params:
             return
-
-        if not isinstance(agent_params, dict):
-            agent_params = agent_params.model_dump()
-
-        if agent_params.get("mbox"):
-            field = f"{target_field}.mbox.keyword"
-            es_query_filters += [{"term": {field: agent_params.get("mbox")}}]
-        elif agent_params.get("mbox_sha1sum"):
-            field = f"{target_field}.mbox_sha1sum.keyword"
-            es_query_filters += [{"term": {field: agent_params.get("mbox_sha1sum")}}]
-        elif agent_params.get("openid"):
-            field = f"{target_field}.openid.keyword"
-            es_query_filters += [{"term": {field: agent_params.get("openid")}}]
-        elif agent_params.get("account__name"):
-            field = f"{target_field}.account.name.keyword"
-            es_query_filters += [{"term": {field: agent_params.get("account__name")}}]
-            field = f"{target_field}.account.homePage.keyword"
-            es_query_filters += [
-                {"term": {field: agent_params.get("account__home_page")}}
-            ]
+        elif not isinstance(agent_params, list):
+            filters = _get_agent_filters(agent_params)
+            if filters:
+                es_query_filters += [filters]
+        else:
+            filters = [_get_agent_filters(params) for params in agent_params if params]
+            es_query_filters += [{"bool": {"should": filters}}]

src/ralph/backends/lrs/fs.py

@@ -136,29 +136,43 @@ def _add_filter_by_agent(
     @staticmethod
     def _add_filter_by_authority(
         filters: list,
-        authority: Optional[AgentParameters],
+        authority: Optional[Union[AgentParameters, list[AgentParameters]]],
     ) -> None:
         """Add authority filters to `filters` if `authority` is set."""
+
+        def _add_filter_by_authority_single(
+            _filter: list, _params: AgentParameters
+        ) -> None:
+            if not _params:
+                return
+            if not isinstance(_params, dict):
+                _params = _params.model_dump()
+
+            FSLRSBackend._add_filter_by_mbox(
+                _filter, _params.get("mbox", None), field="authority"
+            )
+            FSLRSBackend._add_filter_by_sha1sum(
+                _filter, _params.get("mbox_sha1sum", None), field="authority"
+            )
+            FSLRSBackend._add_filter_by_openid(
+                _filter, _params.get("openid", None), field="authority"
+            )
+            FSLRSBackend._add_filter_by_account(
+                _filter,
+                _params.get("account__name", None),
+                _params.get("account__home_page", None),
+                field="authority",
+            )
+
         if not authority:
             return
-
-        if not isinstance(authority, dict):
-            authority = authority.model_dump()
-        FSLRSBackend._add_filter_by_mbox(
-            filters, authority.get("mbox", None), field="authority"
-        )
-        FSLRSBackend._add_filter_by_sha1sum(
-            filters, authority.get("mbox_sha1sum", None), field="authority"
-        )
-        FSLRSBackend._add_filter_by_openid(
-            filters, authority.get("openid", None), field="authority"
-        )
-        FSLRSBackend._add_filter_by_account(
-            filters,
-            authority.get("account__name", None),
-            authority.get("account__home_page", None),
-            field="authority",
-        )
+        elif not isinstance(authority, list):
+            _add_filter_by_authority_single(filters, authority)
+        else:
+            or_filters = []
+            for params in authority:
+                _add_filter_by_authority_single(or_filters, params)
+            FSLRSBackend._add_or_filter(filters, or_filters=or_filters)
 
     @staticmethod
     def _add_filter_by_id(filters: list, statement_id: Optional[str]) -> None:
@@ -334,6 +348,14 @@ def match_related_object_id(statement: dict) -> bool:
         if object_id:
             filters.append(match_related_object_id if related else match_object_id)
 
+    @staticmethod
+    def _add_or_filter(filters: list, or_filters: Optional[list]) -> None:
+        """Add a filter that passes if the input passes any of `or_filters`."""
+        if or_filters:
+            filters.append(
+                lambda statement: any(filter(statement) for filter in or_filters)
+            )
+
     def _add_filter_by_timestamp_since(
         self, filters: list, timestamp: Optional[datetime]
     ) -> None:

src/ralph/backends/lrs/mongo.py

@@ -1,7 +1,7 @@
 """MongoDB LRS backend for Ralph."""
 
 import logging
-from typing import Iterator, List, Optional
+from typing import Iterator, List, Optional, Union
 
 from bson.objectid import ObjectId
 from pydantic_settings import SettingsConfigDict
@@ -122,7 +122,9 @@ def get_query(params: RalphStatementsQuery) -> MongoQuery:
 
     @staticmethod
     def _add_agent_filters(
-        mongo_query_filters: dict, agent_params: AgentParameters, target_field: str
+        mongo_query_filters: dict,
+        agent_params: Union[AgentParameters, list[AgentParameters]],
+        target_field: str,
     ) -> None:
         """Add filters relative to agents to mongo_query_filters.
 
@@ -131,26 +133,42 @@ def _add_agent_filters(
             agent_params (AgentParameters): Agent query parameters to search for.
             target_field (str): The target agent field name to perform the search.
         """
-        if not agent_params:
-            return
-
-        if not isinstance(agent_params, dict):
-            agent_params = agent_params.model_dump()
-
-        if agent_params.get("mbox"):
-            key = f"_source.{target_field}.mbox"
-            mongo_query_filters.update({key: agent_params.get("mbox")})
 
-        if agent_params.get("mbox_sha1sum"):
-            key = f"_source.{target_field}.mbox_sha1sum"
-            mongo_query_filters.update({key: agent_params.get("mbox_sha1sum")})
+        def _get_agent_filters(
+            _params: AgentParameters,
+        ) -> Union[dict, None]:
+            if not _params:
+                return None
+            if not isinstance(_params, dict):
+                _params = _params.model_dump()
+
+            if _params.get("mbox"):
+                key = f"_source.{target_field}.mbox"
+                return {key: _params.get("mbox")}
+
+            if _params.get("mbox_sha1sum"):
+                key = f"_source.{target_field}.mbox_sha1sum"
+                return {key: _params.get("mbox_sha1sum")}
+
+            if _params.get("openid"):
+                key = f"_source.{target_field}.openid"
+                return {key: _params.get("openid")}
+
+            if _params.get("account__name"):
+                key_name = f"_source.{target_field}.account.name"
+                key_homepage = f"_source.{target_field}.account.homePage"
+                return {
+                    key_name: _params.get("account__name"),
+                    key_homepage: _params.get("account__home_page"),
+                }
+            return None
 
-        if agent_params.get("openid"):
-            key = f"_source.{target_field}.openid"
-            mongo_query_filters.update({key: agent_params.get("openid")})
-
-        if agent_params.get("account__name"):
-            key = f"_source.{target_field}.account.name"
-            mongo_query_filters.update({key: agent_params.get("account__name")})
-            key = f"_source.{target_field}.account.homePage"
-            mongo_query_filters.update({key: agent_params.get("account__home_page")})
+        if not agent_params:
+            return
+        elif not isinstance(agent_params, list):
+            filters = _get_agent_filters(agent_params)
+            if filters:
+                mongo_query_filters.update(filters)
+        else:
+            filters = [_get_agent_filters(params) for params in agent_params if params]
+            mongo_query_filters.update({"$or": filters})