Refactor InfrahubEventFilter

Author: ogenstad

backend/infrahub/graphql/queries/event.py

@@ -8,6 +8,7 @@
 from infrahub.exceptions import ValidationError
 from infrahub.graphql.types.event import EventNodes
 from infrahub.task_manager.event import PrefectEvent
+from infrahub.task_manager.models import InfrahubEventFilter
 
 if TYPE_CHECKING:
     from datetime import datetime
@@ -41,58 +42,43 @@ async def resolve(
         if limit > 50:
             # Prefect restricts this to 50
             raise ValidationError(input_value="The parameter 'limit' can't be above 50")
-        return await Events.query(
-            info=info,
-            branch=branches,
-            account=account__ids,
-            limit=limit,
-            level=level,
-            event_type=event_type,
+
+        event_filter = InfrahubEventFilter.from_filters(
+            ids=ids,
+            branches=branches,
+            account__ids=account__ids,
             has_children=has_children,
-            offset=offset,
+            event_type=event_type,
             related_node__ids=related_node__ids,
             primary_node__ids=primary_node__ids,
             parent__ids=parent__ids,
-            ids=ids,
             since=since,
             until=until,
+            level=level,
+        )
+
+        return await Events.query(
+            info=info,
+            event_filter=event_filter,
+            limit=limit,
+            offset=offset,
         )
 
     @classmethod
     async def query(
         cls,
         info: GraphQLResolveInfo,
+        event_filter: InfrahubEventFilter,
         limit: int,
         offset: int | None = None,
-        level: int | None = None,
-        has_children: bool | None = None,
-        ids: list[str] | None = None,
-        event_type: list[str] | None = None,
-        related_node__ids: list[str] | None = None,
-        primary_node__ids: list[str] | None = None,
-        parent__ids: list[str] | None = None,
-        branch: list[str] | None = None,
-        account: list[str] | None = None,
-        since: datetime | None = None,
-        until: datetime | None = None,
     ) -> dict[str, Any]:
         fields = await extract_fields_first_node(info)
 
         prefect_tasks = await PrefectEvent.query(
             fields=fields,
-            ids=ids,
-            related_node__ids=related_node__ids,
-            primary_node__ids=primary_node__ids,
-            parent__ids=parent__ids,
-            branch=branch,
-            event_type=event_type,
-            has_children=has_children,
-            account=account,
-            level=level,
+            event_filter=event_filter,
             limit=limit,
             offset=offset,
-            since=since,
-            until=until,
         )
         return {
             "count": prefect_tasks.get("count", 0),

backend/infrahub/task_manager/event.py

@@ -1,135 +1,18 @@
 from __future__ import annotations
 
-import uuid
 from typing import TYPE_CHECKING, Any
 
 from prefect.client.orchestration import PrefectClient, get_client
-from prefect.events.filters import (
-    EventFilter,
-    EventIDFilter,
-    EventNameFilter,
-    EventOccurredFilter,
-    EventRelatedFilter,
-    EventResourceFilter,
-)
 from prefect.events.schemas.events import Event as PrefectEventModel
-from prefect.events.schemas.events import ResourceSpecification
 from pydantic import BaseModel, Field, TypeAdapter
 
-from infrahub.core.timestamp import Timestamp
 from infrahub.log import get_logger
 from infrahub.utils import get_nested_dict
 
 log = get_logger()
 
 if TYPE_CHECKING:
-    from datetime import datetime
-
-
-class InfrahubEventFilter(EventFilter):
-    matching_related: list[EventRelatedFilter] = Field(default_factory=list)
-
-    def add_account_filter(self, account: list[str] | None) -> None:
-        if account:
-            self.matching_related.append(
-                EventRelatedFilter(
-                    labels=ResourceSpecification(
-                        {"prefect.resource.role": "infrahub.account", "infrahub.resource.id": account}
-                    )
-                )
-            )
-
-    def add_branch_filter(self, branch: list[str] | None = None) -> None:
-        if branch:
-            self.matching_related.append(
-                EventRelatedFilter(
-                    labels=ResourceSpecification(
-                        {"prefect.resource.role": "infrahub.branch", "infrahub.resource.label": branch}
-                    )
-                )
-            )
-
-    def add_event_filter(self, level: int | None = None, has_children: bool | None = None) -> None:
-        event_filter: dict[str, list[str] | str] = {}
-        if level is not None:
-            event_filter["infrahub.event.level"] = str(level)
-
-        if has_children is not None:
-            event_filter["infrahub.event.has_children"] = str(has_children).lower()
-
-        if event_filter:
-            event_filter["prefect.resource.role"] = "infrahub.event"
-            self.matching_related.append(EventRelatedFilter(labels=ResourceSpecification(event_filter)))
-
-    def add_event_id_filter(self, ids: list[str] | None = None) -> None:
-        if ids:
-            self.id = EventIDFilter(id=[uuid.UUID(id) for id in ids])
-
-    def add_event_type_filter(self, event_type: list[str] | None = None) -> None:
-        if event_type:
-            self.event = EventNameFilter(name=event_type)
-
-    def add_primary_node_filter(self, primary_node__ids: list[str] | None) -> None:
-        if primary_node__ids:
-            self.resource = EventResourceFilter(labels=ResourceSpecification({"infrahub.node.id": primary_node__ids}))
-
-    def add_parent_filter(self, parent__ids: list[str] | None) -> None:
-        if parent__ids:
-            self.matching_related.append(
-                EventRelatedFilter(
-                    labels=ResourceSpecification(
-                        {"prefect.resource.role": "infrahub.child_event", "infrahub.event_parent.id": parent__ids}
-                    )
-                )
-            )
-
-    def add_related_node_filter(self, related_node__ids: list[str] | None) -> None:
-        if related_node__ids:
-            self.matching_related.append(
-                EventRelatedFilter(
-                    labels=ResourceSpecification(
-                        {"prefect.resource.role": "infrahub.related.node", "prefect.resource.id": related_node__ids}
-                    )
-                )
-            )
-
-    @classmethod
-    def from_filters(
-        cls,
-        ids: list[str] | None = None,
-        account: list[str] | None = None,
-        related_node__ids: list[str] | None = None,
-        parent__ids: list[str] | None = None,
-        primary_node__ids: list[str] | None = None,
-        event_type: list[str] | None = None,
-        branch: list[str] | None = None,
-        level: int | None = None,
-        has_children: bool | None = None,
-        since: datetime | None = None,
-        until: datetime | None = None,
-    ) -> InfrahubEventFilter:
-        occurred_filter = {}
-        if since:
-            occurred_filter["since"] = Timestamp(since.isoformat()).obj
-
-        if until:
-            occurred_filter["until"] = Timestamp(until.isoformat()).obj
-
-        if occurred_filter:
-            filters = cls(occurred=EventOccurredFilter(**occurred_filter))
-        else:
-            filters = cls()
-
-        filters.add_event_filter(level=level, has_children=has_children)
-        filters.add_event_id_filter(ids=ids)
-        filters.add_event_type_filter(event_type=event_type)
-        filters.add_branch_filter(branch=branch)
-        filters.add_account_filter(account=account)
-        filters.add_parent_filter(parent__ids=parent__ids)
-        filters.add_primary_node_filter(primary_node__ids=primary_node__ids)
-        filters.add_related_node_filter(related_node__ids=related_node__ids)
-
-        return filters
+    from .models import InfrahubEventFilter
 
 
 class PrefectEventData(PrefectEventModel):
@@ -260,7 +143,7 @@ async def query_events(
         cls,
         client: PrefectClient,
         limit: int,
-        filters: EventFilter,
+        filters: InfrahubEventFilter,
         offset: int | None = None,
     ) -> PrefectEventResponse:
         body = {"limit": limit, "filter": filters.model_dump(mode="json", exclude_none=True), "offset": offset}
@@ -278,45 +161,22 @@ async def query_events(
     async def query(
         cls,
         fields: dict[str, Any],
+        event_filter: InfrahubEventFilter,
         limit: int | None = None,
         offset: int | None = None,
-        level: int | None = None,
-        ids: list[str] | None = None,
-        branch: list[str] | None = None,
-        has_children: bool | None = None,
-        account: list[str] | None = None,
-        event_type: list[str] | None = None,
-        related_node__ids: list[str] | None = None,
-        primary_node__ids: list[str] | None = None,
-        parent__ids: list[str] | None = None,
-        since: datetime | None = None,
-        until: datetime | None = None,
     ) -> dict[str, Any]:
         nodes: list[dict] = []
         limit = limit or 50
 
         node_fields = get_nested_dict(nested_dict=fields, keys=["edges", "node"])
-        filters = InfrahubEventFilter.from_filters(
-            ids=ids,
-            branch=branch,
-            account=account,
-            has_children=has_children,
-            event_type=event_type,
-            related_node__ids=related_node__ids,
-            primary_node__ids=primary_node__ids,
-            parent__ids=parent__ids,
-            since=since,
-            until=until,
-            level=level,
-        )
 
         if not node_fields:
             # This means that it's purely a count query and as such we can override the limit to avoid
             # returning data that will only be discarded
             limit = 1
 
         async with get_client(sync_client=False) as client:
-            response = await cls.query_events(client=client, filters=filters, limit=limit, offset=offset)
+            response = await cls.query_events(client=client, filters=event_filter, limit=limit, offset=offset)
             nodes = [{"node": event.to_graphql()} for event in response.events]
 
         return {"count": response.count, "edges": nodes}