infrahub-server 1.1.8__py3-none-any.whl → 1.2.0__py3-none-any.whl

infrahub/menu/models.py

@@ -1,15 +1,15 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, computed_field
 from typing_extensions import Self
 
 from infrahub.core.account import GlobalPermission
 from infrahub.core.node import Node
 from infrahub.core.protocols import CoreMenuItem
-from infrahub.core.schema import GenericSchema, MainSchemaTypes, NodeSchema, ProfileSchema
+from infrahub.core.schema import GenericSchema, MainSchemaTypes, NodeSchema, ProfileSchema, TemplateSchema
 
 from .constants import MenuSection
 
@@ -17,8 +17,8 @@ if TYPE_CHECKING:
     from infrahub.database import InfrahubDatabase
 
 
-def get_full_name(obj: CoreMenuItem | NodeSchema | GenericSchema | ProfileSchema) -> str:
-    if isinstance(obj, (NodeSchema, GenericSchema, ProfileSchema)):
+def get_full_name(obj: CoreMenuItem | NodeSchema | GenericSchema | ProfileSchema | TemplateSchema) -> str:
+    if isinstance(obj, NodeSchema | GenericSchema | ProfileSchema | TemplateSchema):
         return _get_full_name_schema(obj)
     return _get_full_name_node(obj)
 
@@ -35,21 +35,26 @@ def _get_full_name_schema(node: MainSchemaTypes) -> str:
 class MenuDict:
     data: dict[str, MenuItemDict] = field(default_factory=dict)
 
+    def get_item_location(self, name: str) -> list[str]:
+        location, _ = self._find_child_item(name=name, children=self.data)
+        return location
+
     def find_item(self, name: str) -> MenuItemDict | None:
-        return self._find_child_item(name=name, children=self.data)
+        _, item = self._find_child_item(name=name, children=self.data)
+        return item
 
     @classmethod
-    def _find_child_item(cls, name: str, children: dict[str, MenuItemDict]) -> MenuItemDict | None:
+    def _find_child_item(cls, name: str, children: dict[str, MenuItemDict]) -> tuple[list[str], MenuItemDict | None]:
         if name in children.keys():
-            return children[name]
+            return [], children[name]
 
         for child in children.values():
             if not child.children:
                 continue
-            found = cls._find_child_item(name=name, children=child.children)
+            position, found = cls._find_child_item(name=name, children=child.children)
             if found:
-                return found
-        return None
+                return [str(child.identifier)] + position, found
+        return [], None
 
     def to_rest(self) -> Menu:
         data: dict[str, list[MenuItemList]] = {}
@@ -62,10 +67,15 @@ class MenuDict:
 
         return Menu(sections=data)
 
-    # @staticmethod
-    # def _sort_menu_items(items: dict[str, MenuItem]) -> dict[str, MenuItem]:
-    #     sorted_dict = dict(sorted(items.items(), key=lambda x: (x[1].order_weight, x[0]), reverse=False))
-    #     return sorted_dict
+    @classmethod
+    def from_definition_list(cls, definitions: list[MenuItemDefinition]) -> Self:
+        menu = cls()
+        for definition in definitions:
+            menu.data[definition.full_name] = MenuItemDict.from_definition(definition=definition)
+        return menu
+
+    def get_all_identifiers(self) -> set[str]:
+        return {identifier for item in self.data.values() for identifier in item.get_all_identifiers()}
 
 
 @dataclass
@@ -74,7 +84,11 @@ class Menu:
 
 
 class MenuItem(BaseModel):
-    identifier: str = Field(..., description="Unique identifier for this menu item")
+    id: str | None = None
+    namespace: str = Field(..., description="Namespace of the menu item")
+    name: str = Field(..., description="Name of the menu item")
+    description: str = Field(default="", description="Description of the menu item")
+    protected: bool = Field(default=False, description="Whether the menu item is protected")
     label: str = Field(..., description="Title of the menu item")
     path: str = Field(default="", description="URL endpoint if applicable")
     icon: str = Field(default="", description="The icon to show for the current view")
@@ -83,10 +97,27 @@ class MenuItem(BaseModel):
     section: MenuSection = MenuSection.OBJECT
     permissions: list[str] = Field(default_factory=list)
 
+    @computed_field
+    def identifier(self) -> str:
+        return f"{self.namespace}{self.name}"
+
+    def get_path(self) -> str | None:
+        if self.path:
+            return self.path
+
+        if self.kind:
+            return f"/objects/{self.kind}"
+
+        return None
+
     @classmethod
     def from_node(cls, obj: CoreMenuItem) -> Self:
         return cls(
-            identifier=get_full_name(obj),
+            id=obj.get_id(),
+            name=obj.name.value,
+            namespace=obj.namespace.value,
+            protected=obj.protected.value,
+            description=obj.description.value or "",
             label=obj.label.value or "",
             icon=obj.icon.value or "",
             order_weight=obj.order_weight.value,
@@ -96,10 +127,30 @@ class MenuItem(BaseModel):
             permissions=obj.required_permissions.value or [],
         )
 
+    async def to_node(self, db: InfrahubDatabase, parent: CoreMenuItem | None = None) -> CoreMenuItem:
+        obj = await Node.init(db=db, schema=CoreMenuItem)
+        await obj.new(
+            db=db,
+            namespace=self.namespace,
+            name=self.name,
+            label=self.label,
+            kind=self.kind,
+            path=self.get_path(),
+            description=self.description or None,
+            icon=self.icon or None,
+            protected=self.protected,
+            section=self.section.value,
+            order_weight=self.order_weight,
+            parent=parent.id if parent else None,
+            required_permissions=self.permissions,
+        )
+        return obj
+
     @classmethod
-    def from_schema(cls, model: NodeSchema | GenericSchema | ProfileSchema) -> Self:
+    def from_schema(cls, model: NodeSchema | GenericSchema | ProfileSchema | TemplateSchema) -> Self:
         return cls(
-            identifier=get_full_name(model),
+            name=model.name,
+            namespace=model.namespace,
             label=model.label or model.kind,
             path=f"/objects/{model.kind}",
             icon=model.icon or "",
@@ -111,8 +162,14 @@ class MenuItemDict(MenuItem):
     hidden: bool = False
     children: dict[str, MenuItemDict] = Field(default_factory=dict, description="Child objects")
 
+    def get_all_identifiers(self) -> set[str]:
+        identifiers: set[str] = {str(self.identifier)}
+        for child in self.children.values():
+            identifiers.update(child.get_all_identifiers())
+        return identifiers
+
     def to_list(self) -> MenuItemList:
-        data = self.model_dump(exclude={"children"})
+        data = self.model_dump(exclude={"children", "id"})
         unsorted_children = [child.to_list() for child in self.children.values() if child.hidden is False]
         data["children"] = sorted(unsorted_children, key=lambda d: d.order_weight)
         return MenuItemList(**data)
@@ -125,6 +182,35 @@ class MenuItemDict(MenuItem):
             permissions.append(GlobalPermission.from_string(input=permission))
         return permissions
 
+    def diff_attributes(self, other: Self) -> dict[str, Any]:
+        other_attributes = other.model_dump(exclude={"children"})
+        self_attributes = self.model_dump(exclude={"children"})
+        return {
+            key: value
+            for key, value in other_attributes.items()
+            if value != self_attributes[key] and key not in ["id", "children"]
+        }
+
+    @classmethod
+    def from_definition(cls, definition: MenuItemDefinition) -> Self:
+        menu_item = cls(
+            name=definition.name,
+            namespace=definition.namespace,
+            label=definition.label,
+            path=definition.get_path() or "",
+            icon=definition.icon,
+            kind=definition.kind,
+            protected=definition.protected,
+            section=definition.section,
+            permissions=definition.permissions,
+            order_weight=definition.order_weight,
+        )
+
+        for child in definition.children:
+            menu_item.children[child.full_name] = MenuItemDict.from_definition(definition=child)
+
+        return menu_item
+
 
 class MenuItemList(MenuItem):
     children: list[MenuItemList] = Field(default_factory=list, description="Child objects")
@@ -162,6 +248,21 @@ class MenuItemDefinition(BaseModel):
         )
         return obj
 
+    @classmethod
+    async def from_node(cls, node: CoreMenuItem) -> Self:
+        return cls(
+            namespace=node.namespace.value,
+            name=node.name.value,
+            label=node.label.value or "",
+            description=node.description.value or "",
+            icon=node.icon.value or "",
+            protected=node.protected.value,
+            path=node.path.value or "",
+            kind=node.kind.value or "",
+            section=node.section.value,
+            order_weight=node.order_weight.value,
+        )
+
     def get_path(self) -> str | None:
         if self.path:
             return self.path

infrahub/menu/repository.py

@@ -0,0 +1,111 @@
+from infrahub.core.manager import NodeManager
+from infrahub.core.protocols import CoreMenuItem
+from infrahub.database import InfrahubDatabase
+
+from .models import MenuDict, MenuItemDefinition, MenuItemDict
+
+
+class MenuRepository:
+    def __init__(self, db: InfrahubDatabase):
+        self.db = db
+
+    async def get_menu(self, nodes: dict[str, CoreMenuItem] | None = None) -> MenuDict:
+        menu_nodes = nodes or await self.get_menu_db()
+        return await self._convert_menu_from_db(nodes=menu_nodes)
+
+    async def _convert_menu_from_db(self, nodes: dict[str, CoreMenuItem]) -> MenuDict:
+        menu = MenuDict()
+        menu_by_ids = {menu_node.get_id(): MenuItemDict.from_node(menu_node) for menu_node in nodes.values()}
+
+        async def add_children(menu_item: MenuItemDict, menu_node: CoreMenuItem) -> MenuItemDict:
+            children = await menu_node.children.get_peers(db=self.db, peer_type=CoreMenuItem)
+            for child_id, child_node in children.items():
+                child_menu_item = menu_by_ids[child_id]
+                child = await add_children(child_menu_item, child_node)
+                menu_item.children[str(child.identifier)] = child
+            return menu_item
+
+        for menu_node in nodes.values():
+            menu_item = menu_by_ids[menu_node.get_id()]
+            parent = await menu_node.parent.get_peer(db=self.db, peer_type=CoreMenuItem)
+            if parent:
+                continue
+
+            children = await menu_node.children.get_peers(db=self.db, peer_type=CoreMenuItem)
+            for child_id, child_node in children.items():
+                child_menu_item = menu_by_ids[child_id]
+                child = await add_children(child_menu_item, child_node)
+                menu_item.children[str(child.identifier)] = child
+
+            menu.data[str(menu_item.identifier)] = menu_item
+
+        return menu
+
+    async def get_menu_db(self) -> dict[str, CoreMenuItem]:
+        menu_nodes = await NodeManager.query(
+            schema=CoreMenuItem,
+            filters={"namespace__value": "Builtin"},
+            prefetch_relationships=True,
+            db=self.db,
+        )
+        return {node.get_id(): node for node in menu_nodes}
+
+    async def create_menu(self, menu: list[MenuItemDefinition]) -> None:
+        for item in menu:
+            obj = await item.to_node(db=self.db)
+            await obj.save(db=self.db)
+            if item.children:
+                await self.create_menu_children(parent=obj, children=item.children)
+
+    async def create_menu_children(self, parent: CoreMenuItem, children: list[MenuItemDefinition]) -> None:
+        for child in children:
+            obj = await child.to_node(db=self.db, parent=parent)
+            await obj.save(db=self.db)
+            if child.children:
+                await self.create_menu_children(parent=obj, children=child.children)
+
+    async def update_menu(
+        self, existing_menu: MenuDict, new_menu: MenuDict, menu_nodes: dict[str, CoreMenuItem]
+    ) -> None:
+        async def process_menu_item(menu_item: MenuItemDict, parent: CoreMenuItem | None) -> None:
+            existing_item = existing_menu.find_item(name=str(menu_item.identifier))
+            if existing_item and existing_item.id:
+                node = menu_nodes[existing_item.id]
+                await self.update_menu_item(
+                    node=node, existing_menu_item=existing_item, new_menu_item=menu_item, parent=parent
+                )
+            else:
+                node = await self.create_menu_item(new_menu_item=menu_item, parent=parent)
+
+            for child_item in menu_item.children.values():
+                await process_menu_item(menu_item=child_item, parent=node)
+
+        for top_level_item in new_menu.data.values():
+            await process_menu_item(menu_item=top_level_item, parent=None)
+
+        # Delete items that are not in the new menu
+        menu_to_delete = existing_menu.get_all_identifiers() - new_menu.get_all_identifiers()
+        for item_to_delete in menu_to_delete:
+            existing_item = existing_menu.find_item(name=item_to_delete)
+            if existing_item and existing_item.id:
+                node = menu_nodes[existing_item.id]
+                await node.delete(db=self.db)
+
+    async def update_menu_item(
+        self,
+        node: CoreMenuItem,
+        existing_menu_item: MenuItemDict,
+        new_menu_item: MenuItemDict,
+        parent: CoreMenuItem | None,
+    ) -> None:
+        attrs_to_update = existing_menu_item.diff_attributes(new_menu_item)
+        for attr_name, value in attrs_to_update.items():
+            attr = getattr(node, attr_name)
+            attr.value = value
+        await node.parent.update(data=parent, db=self.db)  # type: ignore[arg-type]
+        await node.save(db=self.db)
+
+    async def create_menu_item(self, new_menu_item: MenuItemDict, parent: CoreMenuItem | None) -> CoreMenuItem:
+        obj = await new_menu_item.to_node(db=self.db, parent=parent)
+        await obj.save(db=self.db)
+        return obj

infrahub/menu/utils.py

@@ -1,12 +1,9 @@
-from infrahub.core.protocols import CoreMenuItem
 from infrahub.database import InfrahubDatabase
 
-from .models import MenuItemDefinition
+from .menu import default_menu
+from .repository import MenuRepository
 
 
-async def create_menu_children(db: InfrahubDatabase, parent: CoreMenuItem, children: list[MenuItemDefinition]) -> None:
-    for child in children:
-        obj = await child.to_node(db=db, parent=parent)
-        await obj.save(db=db)
-        if child.children:
-            await create_menu_children(db=db, parent=obj, children=child.children)
+async def create_default_menu(db: InfrahubDatabase) -> None:
+    repository = MenuRepository(db=db)
+    await repository.create_menu(menu=default_menu)

infrahub/message_bus/__init__.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Any, Optional, Union
+from typing import Any
 
 from pydantic import BaseModel, Field
 
@@ -11,21 +11,19 @@ from infrahub.log import set_log_data
 
 class Meta(BaseModel):
     request_id: str = ""
-    correlation_id: Optional[str] = Field(default=None)
-    reply_to: Optional[str] = Field(default=None)
-    initiator_id: Optional[str] = Field(
+    correlation_id: str | None = Field(default=None)
+    reply_to: str | None = Field(default=None)
+    initiator_id: str | None = Field(
         default=None, description="The worker identity of the initial sender of this message"
     )
-    retry_count: Optional[int] = Field(
-        default=None, description="Indicates how many times this message has been retried."
-    )
-    headers: Optional[dict[str, Any]] = Field(default=None)
-    validator_execution_id: Optional[str] = Field(
+    retry_count: int | None = Field(default=None, description="Indicates how many times this message has been retried.")
+    headers: dict[str, Any] | None = Field(default=None)
+    validator_execution_id: str | None = Field(
         default=None, description="Validator execution ID related to this message"
     )
-    check_execution_id: Optional[str] = Field(default=None, description="Check execution ID related to this message")
+    check_execution_id: str | None = Field(default=None, description="Check execution ID related to this message")
     priority: int = Field(default=3, description="Message Priority")
-    expiration: Optional[int] = Field(default=None, description="TTL before this message expires in seconds")
+    expiration: int | None = Field(default=None, description="TTL before this message expires in seconds")
 
     @classmethod
     def default(cls) -> Meta:
@@ -89,9 +87,9 @@ class InfrahubResponse(InfrahubMessage):
 
     passed: bool = True
     routing_key: str
-    data: Union[dict, InfrahubResponseData] = Field(default_factory=dict)
+    data: dict | InfrahubResponseData = Field(default_factory=dict)  # type: ignore[arg-type]
     errors: list[str] = Field(default_factory=list)
-    initial_message: Optional[dict] = Field(
+    initial_message: dict | None = Field(
         default=None,
         description="Initial message in dict format, the primary goal of this field is to provide additional context when there is an error",
     )

infrahub/message_bus/messages/__init__.py

@@ -1,16 +1,7 @@
 from infrahub.message_bus import InfrahubMessage, InfrahubResponse
 
-from .check_artifact_create import CheckArtifactCreate
 from .check_generator_run import CheckGeneratorRun
-from .check_repository_checkdefinition import CheckRepositoryCheckDefinition
-from .check_repository_mergeconflicts import CheckRepositoryMergeConflicts
-from .check_repository_usercheck import CheckRepositoryUserCheck
-from .event_branch_create import EventBranchCreate
-from .event_branch_delete import EventBranchDelete
 from .event_branch_merge import EventBranchMerge
-from .event_branch_rebased import EventBranchRebased
-from .event_node_mutated import EventNodeMutated
-from .event_schema_update import EventSchemaUpdate
 from .event_worker_newprimaryapi import EventWorkerNewPrimaryAPI
 from .finalize_validator_execution import FinalizeValidatorExecution
 from .git_file_get import GitFileGet, GitFileGetResponse
@@ -19,25 +10,13 @@ from .proposed_change.request_proposedchange_refreshartifacts import RequestProp
 from .refresh_git_fetch import RefreshGitFetch
 from .refresh_registry_branches import RefreshRegistryBranches
 from .refresh_registry_rebasedbranch import RefreshRegistryRebasedBranch
-from .request_artifactdefinition_check import RequestArtifactDefinitionCheck
 from .request_generatordefinition_check import RequestGeneratorDefinitionCheck
 from .request_proposedchange_pipeline import RequestProposedChangePipeline
-from .request_repository_checks import RequestRepositoryChecks
-from .request_repository_userchecks import RequestRepositoryUserChecks
 from .send_echo_request import SendEchoRequest, SendEchoRequestResponse
 
 MESSAGE_MAP: dict[str, type[InfrahubMessage]] = {
-    "check.artifact.create": CheckArtifactCreate,
     "check.generator.run": CheckGeneratorRun,
-    "check.repository.check_definition": CheckRepositoryCheckDefinition,
-    "check.repository.merge_conflicts": CheckRepositoryMergeConflicts,
-    "check.repository.user_check": CheckRepositoryUserCheck,
-    "event.branch.create": EventBranchCreate,
-    "event.branch.delete": EventBranchDelete,
     "event.branch.merge": EventBranchMerge,
-    "event.branch.rebased": EventBranchRebased,
-    "event.node.mutated": EventNodeMutated,
-    "event.schema.update": EventSchemaUpdate,
     "event.worker.new_primary_api": EventWorkerNewPrimaryAPI,
     "finalize.validator.execution": FinalizeValidatorExecution,
     "git.file.get": GitFileGet,
@@ -45,12 +24,9 @@ MESSAGE_MAP: dict[str, type[InfrahubMessage]] = {
     "refresh.git.fetch": RefreshGitFetch,
     "refresh.registry.branches": RefreshRegistryBranches,
     "refresh.registry.rebased_branch": RefreshRegistryRebasedBranch,
-    "request.artifact_definition.check": RequestArtifactDefinitionCheck,
     "request.generator_definition.check": RequestGeneratorDefinitionCheck,
     "request.proposed_change.pipeline": RequestProposedChangePipeline,
     "request.proposed_change.refresh_artifacts": RequestProposedChangeRefreshArtifacts,
-    "request.repository.checks": RequestRepositoryChecks,
-    "request.repository.user_checks": RequestRepositoryUserChecks,
     "send.echo.request": SendEchoRequest,
 }
 
@@ -63,7 +39,7 @@ PRIORITY_MAP = {
     "check.artifact.create": 2,
     "check.repository.check_definition": 2,
     "check.repository.merge_conflicts": 2,
-    "event.branch.create": 5,
+    "send.echo.request": 5,  # Currently only for testing purposes, will be removed once all message bus have been migrated to prefect
     "event.branch.delete": 5,
     "event.branch.merge": 5,
     "event.schema.update": 5,

infrahub/message_bus/messages/check_generator_run.py

@@ -1,7 +1,6 @@
-from typing import Optional
-
 from pydantic import Field
 
+from infrahub.context import InfrahubContext
 from infrahub.generators.models import ProposedChangeGeneratorDefinition
 from infrahub.message_bus import InfrahubMessage
 
@@ -10,7 +9,7 @@ class CheckGeneratorRun(InfrahubMessage):
     """A check that runs a generator."""
 
     generator_definition: ProposedChangeGeneratorDefinition = Field(..., description="The Generator definition")
-    generator_instance: Optional[str] = Field(
+    generator_instance: str | None = Field(
         default=None, description="The id of the generator instance if it previously existed"
     )
     commit: str = Field(..., description="The commit to target")
@@ -24,3 +23,4 @@ class CheckGeneratorRun(InfrahubMessage):
     variables: dict = Field(..., description="Input variables when running the generator")
     validator_id: str = Field(..., description="The ID of the validator")
     proposed_change: str | None = Field(None, description="The unique ID of the Proposed Change")
+    context: InfrahubContext = Field(..., description="The Infrahub context")

infrahub/message_bus/messages/event_branch_merge.py

@@ -1,5 +1,6 @@
 from pydantic import Field
 
+from infrahub.context import InfrahubContext
 from infrahub.message_bus import InfrahubMessage
 
 
@@ -8,3 +9,5 @@ class EventBranchMerge(InfrahubMessage):
 
     source_branch: str = Field(..., description="The source branch")
     target_branch: str = Field(..., description="The target branch")
+
+    context: InfrahubContext = Field(..., description="The context of the event")

infrahub/message_bus/messages/finalize_validator_execution.py

@@ -1,5 +1,6 @@
 from pydantic import Field
 
+from infrahub.context import InfrahubContext
 from infrahub.message_bus import InfrahubMessage
 
 
@@ -10,3 +11,5 @@ class FinalizeValidatorExecution(InfrahubMessage):
     validator_execution_id: str = Field(..., description="The id of current execution of the associated validator")
     start_time: str = Field(..., description="Start time when the message was first created")
     validator_type: str = Field(..., description="The type of validator to complete")
+    context: InfrahubContext = Field(..., description="The Infrahub context")
+    proposed_change: str = Field(..., description="The ID of the proposed change")

infrahub/message_bus/messages/proposed_change/request_proposedchange_refreshartifacts.py

@@ -1,5 +1,11 @@
+from pydantic import Field
+
+from infrahub.context import InfrahubContext
+
 from .base_with_diff import BaseProposedChangeWithDiffMessage
 
 
 class RequestProposedChangeRefreshArtifacts(BaseProposedChangeWithDiffMessage):
     """Sent trigger the refresh of artifacts that are impacted by the proposed change."""
+
+    context: InfrahubContext = Field(..., description="The context of the task")

infrahub/message_bus/messages/request_generatordefinition_check.py

@@ -1,5 +1,6 @@
 from pydantic import ConfigDict, Field
 
+from infrahub.context import InfrahubContext
 from infrahub.generators.models import ProposedChangeGeneratorDefinition
 from infrahub.message_bus import InfrahubMessage
 from infrahub.message_bus.types import ProposedChangeBranchDiff
@@ -16,3 +17,4 @@ class RequestGeneratorDefinitionCheck(InfrahubMessage):
     source_branch: str = Field(..., description="The source branch")
     source_branch_sync_with_git: bool = Field(..., description="Indicates if the source branch should sync with git")
     destination_branch: str = Field(..., description="The target branch")
+    context: InfrahubContext = Field(..., description="The Infrahub context")

infrahub/message_bus/messages/request_proposedchange_pipeline.py

@@ -1,5 +1,6 @@
 from pydantic import Field
 
+from infrahub.context import InfrahubContext
 from infrahub.core.constants import CheckType
 from infrahub.message_bus import InfrahubMessage
 
@@ -14,3 +15,4 @@ class RequestProposedChangePipeline(InfrahubMessage):
     check_type: CheckType = Field(
         default=CheckType.ALL, description="Can be used to restrict the pipeline to a specific type of job"
     )
+    context: InfrahubContext = Field(..., description="The context of the task")

infrahub/message_bus/messages/send_echo_request.py

@@ -6,7 +6,7 @@ ROUTING_KEY = "send.echo.request"
 
 
 class SendEchoRequest(InfrahubMessage):
-    """Sent a echo request, i.e., ping message."""
+    """Sent a echo request, a ping message for example."""
 
     message: str = Field(..., description="The message to send")
 

infrahub/message_bus/operations/__init__.py

@@ -1,5 +1,3 @@
-from typing import Optional
-
 import ujson
 from prefect import Flow
 
@@ -18,14 +16,8 @@ from infrahub.services import InfrahubServices
 from infrahub.tasks.check import set_check_status
 
 COMMAND_MAP = {
-    "check.artifact.create": check.artifact.create,
     "check.generator.run": check.generator.run,
-    "check.repository.check_definition": check.repository.check_definition,
-    "check.repository.merge_conflicts": check.repository.merge_conflicts,
-    "check.repository.user_check": check.repository.user_check,
     "event.branch.merge": event.branch.merge,
-    "event.node.mutated": event.node.mutated,
-    "event.schema.update": event.schema.update,
     "event.worker.new_primary_api": event.worker.new_primary_api,
     "finalize.validator.execution": finalize.validator.execution,
     "git.file.get": git.file.get,
@@ -34,18 +26,15 @@ COMMAND_MAP = {
     "refresh.registry.branches": refresh.registry.branches,
     "refresh.registry.rebased_branch": refresh.registry.rebased_branch,
     "request.generator_definition.check": requests.generator_definition.check,
-    "request.artifact_definition.check": requests.artifact_definition.check,
     "request.proposed_change.pipeline": requests.proposed_change.pipeline,
     "request.proposed_change.refresh_artifacts": requests.proposed_change.refresh_artifacts,
-    "request.repository.checks": requests.repository.checks,
-    "request.repository.user_checks": requests.repository.user_checks,
     "send.echo.request": send.echo.request,
 }
 
 
 async def execute_message(
     routing_key: str, message_body: bytes, service: InfrahubServices, skip_flow: bool = False
-) -> Optional[MessageTTL]:
+) -> MessageTTL | None:
     message_data = ujson.loads(message_body)
     message = messages.MESSAGE_MAP[routing_key](**message_data)
     message.set_log_data(routing_key=routing_key)
@@ -54,15 +43,15 @@ async def execute_message(
         if skip_flow and isinstance(func, Flow):
             func = func.fn
         await func(message=message, service=service)
-    except Exception as exc:  # pylint: disable=broad-except
+    except Exception as exc:
         if message.reply_requested:
             response = RPCErrorResponse(errors=[str(exc)], initial_message=message.model_dump())
-            await service.reply(message=response, initiator=message)
+            await service.message_bus.reply_if_initiator_meta(message=response, initiator=message)
             return None
         if message.reached_max_retries:
             service.log.exception("Message failed after maximum number of retries", error=exc)
             await set_check_status(message, conclusion="failure", service=service)
             return None
         message.increase_retry_count()
-        await service.send(message, delay=MessageTTL.FIVE, is_retry=True)
+        await service.message_bus.send(message, delay=MessageTTL.FIVE, is_retry=True)
         return MessageTTL.FIVE

infrahub/message_bus/operations/check/__init__.py

@@ -1,3 +1,3 @@
-from . import artifact, generator, repository
+from . import generator
 
-__all__ = ["artifact", "generator", "repository"]
+__all__ = ["generator"]

infrahub/message_bus/operations/check/generator.py

@@ -14,8 +14,6 @@ from infrahub.services import InfrahubServices
 from infrahub.tasks.check import set_check_status
 from infrahub.workflows.utils import add_tags
 
-# pylint: disable=duplicate-code
-
 
 @flow(
     name="git-repository-check-generator-run",
@@ -74,6 +72,7 @@ async def run(message: messages.CheckGeneratorRun, service: InfrahubServices) ->
             convert_query_response=generator_definition.convert_query_response,
             infrahub_node=InfrahubNode,
         )
+        generator._init_client.request_context = message.context.to_request_context()
         await generator.run(identifier=generator_definition.name)
         generator_instance.status.value = GeneratorInstanceStatus.READY.value
     except ModuleImportError as exc:
@@ -81,7 +80,7 @@ async def run(message: messages.CheckGeneratorRun, service: InfrahubServices) ->
         generator_instance.status.value = GeneratorInstanceStatus.ERROR.value
         check_message = f"Failed to import generator: {exc.message}"
         log.exception(check_message, exc_info=exc)
-    except Exception as exc:  # pylint: disable=broad-exception-caught
+    except Exception as exc:
         conclusion = ValidatorConclusion.FAILURE
         generator_instance.status.value = GeneratorInstanceStatus.ERROR.value
         check_message = f"Failed to execute generator: {str(exc)}"