infrahub-server 1.1.7__py3-none-any.whl → 1.2.0b1__py3-none-any.whl

infrahub/services/scheduler.py

@@ -7,12 +7,15 @@ from typing import TYPE_CHECKING
 
 from infrahub import config
 from infrahub.components import ComponentType
+from infrahub.log import get_logger
 from infrahub.tasks.keepalive import refresh_heartbeat
 from infrahub.tasks.recurring import trigger_branch_refresh
 
 if TYPE_CHECKING:
     from infrahub.services import InfrahubServices, ServiceFunction
 
+log = get_logger()
+
 
 @dataclass
 class Schedule:
@@ -23,18 +26,17 @@ class Schedule:
 
 
 class InfrahubScheduler:
-    def __init__(self) -> None:
-        self.service: InfrahubServices
+    # TODO we could remove service dependency by adding kwargs to Schedule instead of passing services
+    service: InfrahubServices | None
+
+    def __init__(self, component_type: ComponentType) -> None:
         self.running: bool = False
         self.schedules: list[Schedule] = []
 
-    async def initialize(self, service: InfrahubServices) -> None:
-        self.service = service
-
         self.running = config.SETTINGS.miscellaneous.start_background_runner
         # Add some randomness to the interval to avoid having all workers pulling the latest update at the same time
         random_number = random.randint(0, 5)
-        if self.service.component_type == ComponentType.API_SERVER:
+        if component_type == ComponentType.API_SERVER:
             schedules = [
                 Schedule(name="refresh_api_components", interval=10, function=refresh_heartbeat, start_delay=0),
                 Schedule(
@@ -43,7 +45,7 @@ class InfrahubScheduler:
             ]
             self.schedules.extend(schedules)
 
-        if self.service.component_type == ComponentType.GIT_AGENT:
+        if component_type == ComponentType.GIT_AGENT:
             schedules = [
                 Schedule(name="refresh_components", interval=10, function=refresh_heartbeat),
                 Schedule(
@@ -54,32 +56,31 @@ class InfrahubScheduler:
 
     async def start_schedule(self) -> None:
         for schedule in self.schedules:
-            asyncio.create_task(
-                run_schedule(schedule=schedule, service=self.service), name=f"scheduled_task_{schedule.name}"
-            )
+            asyncio.create_task(self.run_schedule(schedule=schedule), name=f"scheduled_task_{schedule.name}")
 
     async def shutdown(self) -> None:
         self.running = False
 
+    async def run_schedule(self, schedule: Schedule) -> None:
+        """Execute the task provided in the schedule as per the defined interval
 
-async def run_schedule(schedule: Schedule, service: InfrahubServices) -> None:
-    """Execute the task provided in the schedule as per the defined interval
-
-    Once the service is marked to be shutdown the scheduler will stop executing tasks.
-    """
-    for _ in range(schedule.start_delay):
-        if not service.scheduler.running:
-            return
-        await asyncio.sleep(delay=1)
-
-    service.log.info("Started recurring task", task=schedule.name)
-
-    while service.scheduler.running:
-        try:
-            await schedule.function(service)
-        except Exception as exc:  # pylint: disable=broad-exception-caught
-            service.log.error(str(exc))
-        for _ in range(schedule.interval):
-            if not service.scheduler.running:
+        Once the service is marked to be shutdown the scheduler will stop executing tasks.
+        """
+        for _ in range(schedule.start_delay):
+            if not self.running:
                 return
             await asyncio.sleep(delay=1)
+
+        if self.service is None:
+            raise ValueError("InfrahubScheduler.service is None")
+
+        self.service.log.info("Started recurring task", task=schedule.name)
+        while self.running:
+            try:
+                await schedule.function(self.service)
+            except Exception as exc:
+                self.service.log.error(str(exc))
+            for _ in range(schedule.interval):
+                if not self.running:
+                    return
+                await asyncio.sleep(delay=1)

infrahub/storage.py

@@ -53,7 +53,5 @@ class InfrahubObjectStorage:
         try:
             with self._storage.open(identifier) as f:
                 return f.read().decode()
-        except (FileNotFoundError, botocore.exceptions.ClientError):
-            raise NodeNotFoundError(  # pylint: disable=raise-missing-from
-                node_type="StorageObject", identifier=identifier
-            )
+        except (FileNotFoundError, botocore.exceptions.ClientError) as err:
+            raise NodeNotFoundError(node_type="StorageObject", identifier=identifier) from err

infrahub/task_manager/constants.py

@@ -2,7 +2,7 @@ from infrahub.core.constants import TaskConclusion
 
 LOG_LEVEL_MAPPING = {10: "debug", 20: "info", 30: "warning", 40: "error", 50: "critical"}
 
-CONCLUSION_STATE_MAPPING = {
+CONCLUSION_STATE_MAPPING: dict[str, TaskConclusion] = {
     "Scheduled": TaskConclusion.UNKNOWN,
     "Pending": TaskConclusion.UNKNOWN,
     "Running": TaskConclusion.UNKNOWN,

infrahub/task_manager/event.py

@@ -0,0 +1,261 @@
+from __future__ import annotations
+
+import asyncio
+from typing import TYPE_CHECKING, Any
+
+from prefect.client.orchestration import PrefectClient, get_client
+from prefect.events.schemas.events import Event as PrefectEventModel
+from pydantic import BaseModel, Field, TypeAdapter
+
+from infrahub.core.constants import GLOBAL_BRANCH_NAME
+from infrahub.exceptions import ServiceUnavailableError
+from infrahub.log import get_logger
+from infrahub.utils import get_nested_dict
+
+log = get_logger()
+
+if TYPE_CHECKING:
+    from .models import InfrahubEventFilter
+
+
+class PrefectEventData(PrefectEventModel):
+    def get_branch(self) -> str | None:
+        for resource in self.related:
+            if resource.get("prefect.resource.role") != "infrahub.branch":
+                continue
+            if "infrahub.resource.label" not in resource:
+                continue
+            if resource.get("infrahub.resource.label") == GLOBAL_BRANCH_NAME:
+                return None
+            return resource.get("infrahub.resource.label")
+        return None
+
+    def get_level(self) -> int:
+        for resource in self.related:
+            level = resource.get("infrahub.event.level")
+            if level is None:
+                continue
+            try:
+                return int(level)
+            except ValueError:
+                return 0
+
+        return 0
+
+    def get_parent(self) -> str | None:
+        for resource in self.related:
+            if resource.get("prefect.resource.role") != "infrahub.child_event":
+                continue
+            return resource.get("infrahub.event_parent.id")
+        return None
+
+    def get_primary_node(self) -> dict[str, str] | None:
+        node_id = self.resource.get("infrahub.node.id")
+        node_kind = self.resource.get("infrahub.node.kind")
+        if node_id and node_kind:
+            return {"id": node_id, "kind": node_kind}
+
+        return None
+
+    def get_related_nodes(self) -> list[dict[str, str]]:
+        related_nodes = []
+        for resource in self.related:
+            if resource.get("prefect.resource.role") != "infrahub.related.node":
+                continue
+
+            node_id = resource.get("prefect.resource.id")
+            node_kind = resource.get("infrahub.node.kind")
+            if node_id == self.resource.get("infrahub.node.id"):
+                # Don't include the primary node as a related node.
+                continue
+            if node_id and node_kind:
+                related_nodes.append({"id": node_id, "kind": node_kind})
+
+        return related_nodes
+
+    def get_account_id(self) -> str | None:
+        for resource in self.related:
+            if resource.get("prefect.resource.role") != "infrahub.account":
+                continue
+            return resource.get("infrahub.resource.id")
+        return None
+
+    def has_children(self) -> bool:
+        for resource in self.related:
+            if resource.get("prefect.resource.role") != "infrahub.event":
+                continue
+            if resource.get("infrahub.event.has_children") == "true":
+                return True
+            return False
+        return False
+
+    def _return_node_mutation(self) -> dict[str, Any]:
+        attributes = []
+
+        for resource in self.related:
+            if resource.get("prefect.resource.role") == "infrahub.node.field_update" and resource.get(
+                "infrahub.attribute.name"
+            ):
+                attributes.append(
+                    {
+                        "name": resource.get("infrahub.attribute.name", ""),
+                        "kind": resource.get("infrahub.attribute.kind", ""),
+                        "value": None
+                        if resource.get("infrahub.attribute.value") == "NULL"
+                        else resource.get("infrahub.attribute.value"),
+                        "value_previous": None
+                        if resource.get("infrahub.attribute.value_previous") == "NULL"
+                        else resource.get("infrahub.attribute.value_previous"),
+                        "action": resource.get("infrahub.attribute.action", "unchanged"),
+                    }
+                )
+
+        return {"attributes": attributes}
+
+    def _get_branch_name_from_resource(self) -> str:
+        return self.resource.get("infrahub.branch.name") or ""
+
+    def _return_artifact_event(self) -> dict[str, Any]:
+        checksum = ""
+        checksum_previous: str | None = None
+        storage_id = ""
+        storage_id_previous: str | None = None
+        artifact_definition_id = ""
+        for resource in self.related:
+            if resource.role == "infrahub.artifact":
+                checksum = resource.get("infrahub.artifact.checksum") or ""
+                checksum_previous = resource.get("infrahub.artifact.checksum_previous")
+                storage_id = resource.get("infrahub.artifact.storage_id") or ""
+                storage_id_previous = resource.get("infrahub.artifact.storage_id_previous")
+                artifact_definition_id = resource.get("infrahub.artifact.artifact_definition_id") or ""
+
+        return {
+            "checksum": checksum,
+            "checksum_previous": checksum_previous,
+            "storage_id": storage_id,
+            "storage_id_previous": storage_id_previous,
+            "artifact_definition_id": artifact_definition_id,
+        }
+
+    def _return_branch_created(self) -> dict[str, Any]:
+        return {"created_branch": self._get_branch_name_from_resource()}
+
+    def _return_branch_deleted(self) -> dict[str, Any]:
+        return {"deleted_branch": self._get_branch_name_from_resource()}
+
+    def _return_branch_merged(self) -> dict[str, Any]:
+        return {"source_branch": self._get_branch_name_from_resource()}
+
+    def _return_branch_rebased(self) -> dict[str, Any]:
+        return {"rebased_branch": self._get_branch_name_from_resource()}
+
+    def _return_group_event(self) -> dict[str, Any]:
+        members = []
+        ancestors = []
+
+        for resource in self.related:
+            if resource.role == "infrahub.group.member" and resource.get("infrahub.node.kind"):
+                members.append({"id": resource.id, "kind": resource.get("infrahub.node.kind")})
+            elif resource.role == "infrahub.group.ancestor" and resource.get("infrahub.node.kind"):
+                ancestors.append({"id": resource.id, "kind": resource.get("infrahub.node.kind")})
+
+        return {"members": members, "ancestors": ancestors}
+
+    def _return_event_specifics(self) -> dict[str, Any]:
+        """Return event specific data based on the type of event being processed"""
+
+        event_specifics = {}
+
+        match self.event:
+            case "infrahub.artifact.created" | "infrahub.artifact.updated":
+                event_specifics = self._return_artifact_event()
+            case "infrahub.node.created" | "infrahub.node.updated" | "infrahub.node.deleted":
+                event_specifics = self._return_node_mutation()
+            case "infrahub.branch.created":
+                event_specifics = self._return_branch_created()
+            case "infrahub.branch.deleted":
+                event_specifics = self._return_branch_deleted()
+            case "infrahub.branch.merged":
+                event_specifics = self._return_branch_merged()
+            case "infrahub.branch.rebased":
+                event_specifics = self._return_branch_rebased()
+            case "infrahub.group.member_added" | "infrahub.group.member_removed":
+                event_specifics = self._return_group_event()
+
+        return event_specifics
+
+    def to_graphql(self) -> dict[str, Any]:
+        response = {
+            "id": str(self.id),
+            "event": self.event,
+            "branch": self.get_branch(),
+            "account_id": self.get_account_id(),
+            "occurred_at": self.occurred,
+            "has_children": self.has_children(),
+            "payload": self.payload,
+            "level": self.get_level(),
+            "primary_node": self.get_primary_node(),
+            "parent_id": self.get_parent(),
+            "related_nodes": self.get_related_nodes(),
+        }
+        response.update(self._return_event_specifics())
+        return response
+
+
+class PrefectEventResponse(BaseModel):
+    count: int = Field(..., description="Number of matching events")
+    events: list[PrefectEventData] = Field(..., description="Returned events")
+
+
+class PrefectEvent:
+    @classmethod
+    async def query_events(
+        cls,
+        client: PrefectClient,
+        limit: int,
+        filters: InfrahubEventFilter,
+        offset: int | None = None,
+    ) -> PrefectEventResponse:
+        body = {"limit": limit, "filter": filters.model_dump(mode="json", exclude_none=True), "offset": offset}
+
+        # Retry due to https://github.com/PrefectHQ/prefect/issues/16299
+        for _ in range(1, 5):
+            response = await client._client.post("/infrahub/events/filter", json=body)
+            if response.status_code == 200:
+                break
+            await asyncio.sleep(0.1)
+
+        if response.status_code != 200:
+            raise ServiceUnavailableError(
+                message=f"Unable to query prefect due to invalid response from the server (status_code={response.status_code})"
+            )
+        data: dict[str, Any] = response.json()
+
+        return PrefectEventResponse(
+            count=data.get("total", 0),
+            events=TypeAdapter(list[PrefectEventData]).validate_python(data.get("events")),
+        )
+
+    @classmethod
+    async def query(
+        cls,
+        fields: dict[str, Any],
+        event_filter: InfrahubEventFilter,
+        limit: int | None = None,
+        offset: int | None = None,
+    ) -> dict[str, Any]:
+        nodes: list[dict] = []
+        limit = limit or 50
+
+        node_fields = get_nested_dict(nested_dict=fields, keys=["edges", "node"])
+
+        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=event_filter, limit=limit, offset=offset)
+            nodes = [{"node": event.to_graphql()} for event in response.events]
+
+        return {"count": response.count, "edges": nodes}

infrahub/task_manager/models.py

@@ -1,11 +1,29 @@
+from __future__ import annotations
+
+import uuid
 from collections import defaultdict
+from typing import TYPE_CHECKING, Any
 from uuid import UUID
 
-from prefect.client.schemas.objects import Log as PrefectLog
+from prefect.client.schemas.objects import Log as PrefectLog  # noqa: TC002
+from prefect.events.filters import (
+    EventFilter,
+    EventIDFilter,
+    EventNameFilter,
+    EventOccurredFilter,
+    EventRelatedFilter,
+    EventResourceFilter,
+)
+from prefect.events.schemas.events import ResourceSpecification
 from pydantic import BaseModel, Field
 
+from infrahub.core.timestamp import Timestamp
+
 from .constants import LOG_LEVEL_MAPPING
 
+if TYPE_CHECKING:
+    from datetime import datetime
+
 
 class RelatedNodeInfo(BaseModel):
     id: str
@@ -13,7 +31,7 @@ class RelatedNodeInfo(BaseModel):
 
 
 class RelatedNodesInfo(BaseModel):
-    flows: dict[UUID, dict[str, RelatedNodeInfo]] = Field(default_factory=lambda: defaultdict(dict))
+    flows: dict[UUID, dict[str, RelatedNodeInfo]] = Field(default_factory=lambda: defaultdict(dict))  # type: ignore[arg-type]
     nodes: dict[str, RelatedNodeInfo] = Field(default_factory=dict)
 
     def add_nodes(self, flow_id: UUID, node_ids: list[str]) -> None:
@@ -46,7 +64,7 @@ class RelatedNodesInfo(BaseModel):
 
 
 class FlowLogs(BaseModel):
-    logs: defaultdict[UUID, list[PrefectLog]] = Field(default_factory=lambda: defaultdict(list))
+    logs: defaultdict[UUID, list[PrefectLog]] = Field(default_factory=lambda: defaultdict(list))  # type: ignore[arg-type]
 
     def to_graphql(self, flow_id: UUID) -> list[dict]:
         return [
@@ -63,3 +81,129 @@ class FlowLogs(BaseModel):
 
 class FlowProgress(BaseModel):
     data: dict[UUID, float] = Field(default_factory=dict)
+
+
+class InfrahubEventFilter(EventFilter):
+    matching_related: list[EventRelatedFilter] = Field(default_factory=list)
+
+    def add_account_filter(self, account__ids: list[str] | None) -> None:
+        if account__ids:
+            self.matching_related.append(
+                EventRelatedFilter(
+                    labels=ResourceSpecification(
+                        {"prefect.resource.role": "infrahub.account", "infrahub.resource.id": account__ids}
+                    )
+                )
+            )
+
+    def add_branch_filter(self, branches: list[str] | None = None) -> None:
+        if branches:
+            self.matching_related.append(
+                EventRelatedFilter(
+                    labels=ResourceSpecification(
+                        {"prefect.resource.role": "infrahub.branch", "infrahub.resource.label": branches}
+                    )
+                )
+            )
+
+    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, event_type_filter: dict[str, Any] | None = None
+    ) -> None:
+        event_type = event_type or []
+        event_type_filter = event_type_filter or {}
+
+        if branch_merged := event_type_filter.get("branch_merged"):
+            branches: list[str] = branch_merged.get("branches") or []
+            if "infrahub.branch.created" not in event_type:
+                event_type.append("infrahub.branch.merged")
+            if branches:
+                self.resource = EventResourceFilter(labels=ResourceSpecification({"infrahub.branch.name": branches}))
+
+        if branch_rebased := event_type_filter.get("branch_rebased"):
+            branches = branch_rebased.get("branches") or []
+            if "infrahub.branch.created" not in event_type:
+                event_type.append("infrahub.branch.rebased")
+            if branches:
+                self.resource = EventResourceFilter(labels=ResourceSpecification({"infrahub.branch.name": branches}))
+
+        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__ids: 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,
+        event_type_filter: dict[str, Any] | None = None,
+        branches: 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()).to_datetime()
+
+        if until:
+            occurred_filter["until"] = Timestamp(until.isoformat()).to_datetime()
+
+        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, event_type_filter=event_type_filter)
+        filters.add_branch_filter(branches=branches)
+        filters.add_account_filter(account__ids=account__ids)
+        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

infrahub/task_manager/task.py

@@ -250,7 +250,7 @@ class PrefectTask:
                             "node": {
                                 "title": flow.name,
                                 "conclusion": CONCLUSION_STATE_MAPPING.get(
-                                    flow.state_name, TaskConclusion.UNKNOWN
+                                    str(flow.state_name), TaskConclusion.UNKNOWN
                                 ).value,
                                 "state": flow.state_type,
                                 "progress": progress_flow.data.get(flow.id, None),

infrahub/tasks/artifact.py

@@ -5,41 +5,42 @@ from prefect import task
 from prefect.cache_policies import NONE
 
 from infrahub import lock
+from infrahub.artifacts.models import CheckArtifactCreate
 from infrahub.core.constants import InfrahubKind
 from infrahub.git.models import RequestArtifactGenerate
-from infrahub.message_bus import messages
 from infrahub.services import InfrahubServices
 
 
-@task(name="define-artifact", task_run_name="Define Artifact", cache_policy=NONE)
+@task(name="define-artifact", task_run_name="Define Artifact", cache_policy=NONE)  # type: ignore[arg-type]
 async def define_artifact(
-    message: Union[messages.CheckArtifactCreate, RequestArtifactGenerate], service: InfrahubServices
-) -> InfrahubNode:
-    if message.artifact_id:
-        artifact = await service.client.get(
-            kind=InfrahubKind.ARTIFACT, id=message.artifact_id, branch=message.branch_name
-        )
+    model: Union[CheckArtifactCreate, RequestArtifactGenerate], service: InfrahubServices
+) -> tuple[InfrahubNode, bool]:
+    """Return an artifact together with a flag to indicate if the artifact is created now or already existed."""
+    created = False
+    if model.artifact_id:
+        artifact = await service.client.get(kind=InfrahubKind.ARTIFACT, id=model.artifact_id, branch=model.branch_name)
     else:
-        async with lock.registry.get(f"{message.target_id}-{message.artifact_definition}", namespace="artifact"):
+        async with lock.registry.get(f"{model.target_id}-{model.artifact_definition}", namespace="artifact"):
             artifacts = await service.client.filters(
                 kind=InfrahubKind.ARTIFACT,
-                branch=message.branch_name,
-                definition__ids=[message.artifact_definition],
-                object__ids=[message.target_id],
+                branch=model.branch_name,
+                definition__ids=[model.artifact_definition],
+                object__ids=[model.target_id],
             )
             if artifacts:
                 artifact = artifacts[0]
             else:
                 artifact = await service.client.create(
                     kind=InfrahubKind.ARTIFACT,
-                    branch=message.branch_name,
+                    branch=model.branch_name,
                     data={
-                        "name": message.artifact_name,
+                        "name": model.artifact_name,
                         "status": "Pending",
-                        "object": message.target_id,
-                        "definition": message.artifact_definition,
-                        "content_type": message.content_type,
+                        "object": model.target_id,
+                        "definition": model.artifact_definition,
+                        "content_type": model.content_type,
                     },
                 )
                 await artifact.save()
-    return artifact
+                created = True
+    return artifact, created

infrahub/tasks/registry.py

@@ -18,7 +18,7 @@ async def refresh_branches(db: InfrahubDatabase) -> None:
     If a branch is already present with a different value for the hash
     We pull the new schema from the database and we update the registry.
     """
-    from infrahub.graphql.manager import GraphQLSchemaManager  # pylint: disable=import-outside-toplevel,cyclic-import
+    from infrahub.graphql.manager import GraphQLSchemaManager
 
     async with lock.registry.local_schema_lock():
         branches = await registry.branch_object.get_list(db=db)