infrahub-server 1.1.6__py3-none-any.whl → 1.2.0rc0__py3-none-any.whl

infrahub/services/adapters/message_bus/nats.py

@@ -33,7 +33,7 @@ async def _add_request_id(message: InfrahubMessage) -> None:
 
 
 class NATSMessageBus(InfrahubMessageBus):
-    def __init__(self, settings: Optional[BrokerSettings] = None) -> None:
+    def __init__(self, component_type: ComponentType, settings: Optional[BrokerSettings] = None) -> None:
         self.settings = settings or config.SETTINGS.broker
 
         self.service: InfrahubServices
@@ -46,9 +46,15 @@ class NATSMessageBus(InfrahubMessageBus):
         self.loop = asyncio.get_running_loop()
         self.futures: MutableMapping[str, asyncio.Future] = {}
 
-    async def initialize(self, service: InfrahubServices) -> None:
-        self.service = service
+        self.component_type: ComponentType = component_type
 
+    @classmethod
+    async def new(cls, component_type: ComponentType, settings: Optional[BrokerSettings] = None) -> NATSMessageBus:
+        message_bus = cls(component_type=component_type, settings=settings)
+        await message_bus._initialize()
+        return message_bus
+
+    async def _initialize(self) -> None:
         tls_context = None
         if self.settings.tls_enabled:
             tls_context = ssl.create_default_context(purpose=ssl.Purpose.SERVER_AUTH)
@@ -67,9 +73,9 @@ class NATSMessageBus(InfrahubMessageBus):
 
         self.jetstream = self.connection.jetstream()
 
-        if self.service.component_type == ComponentType.API_SERVER:
+        if self.component_type == ComponentType.API_SERVER:
             await self._initialize_api_server()
-        elif self.service.component_type == ComponentType.GIT_AGENT:
+        elif self.component_type == ComponentType.GIT_AGENT:
             await self._initialize_git_worker()
 
     async def shutdown(self) -> None:
@@ -282,7 +288,7 @@ class NATSMessageBus(InfrahubMessageBus):
             request_id=request_id, correlation_id=correlation_id, reply_to=self.callback_queue.config.name
         )
 
-        await self.service.send(message=message)
+        await self.service.message_bus.send(message=message)
 
         response: nats.aio.msg.Msg = await future
         data = ujson.loads(response.data)

infrahub/services/adapters/message_bus/rabbitmq.py

@@ -65,12 +65,14 @@ async def _add_request_id(message: InfrahubMessage) -> None:
 
 
 class RabbitMQMessageBus(InfrahubMessageBus):
-    def __init__(self, settings: Optional[BrokerSettings] = None) -> None:
+    def __init__(
+        self, component_type: ComponentType = ComponentType.NONE, settings: Optional[BrokerSettings] = None
+    ) -> None:
         self.settings = settings or config.SETTINGS.broker
         self.channel: AbstractChannel
         self.exchange: AbstractExchange
         self.delayed_exchange: AbstractExchange
-        self.service: InfrahubServices
+        self.service: InfrahubServices  # Needed because we inject `service` while sending messages.
         self.connection: AbstractRobustConnection
         self.callback_queue: AbstractQueue
         self.events_queue: AbstractQueue
@@ -80,10 +82,17 @@ class RabbitMQMessageBus(InfrahubMessageBus):
         self.loop = asyncio.get_running_loop()
         self.futures: MutableMapping[str, asyncio.Future] = {}
 
-    async def initialize(self, service: InfrahubServices) -> None:
+        self.component_type: ComponentType = component_type
+
+    @classmethod
+    async def new(cls, component_type: ComponentType, settings: Optional[BrokerSettings] = None) -> RabbitMQMessageBus:
+        message_bus = cls(component_type=component_type, settings=settings)
+        await message_bus._initialize()
+        return message_bus
+
+    async def _initialize(self) -> None:
         patch_spanbuilder_set_channel()
 
-        self.service = service
         self.connection = await aio_pika.connect_robust(
             host=self.settings.address,
             login=self.settings.username,
@@ -103,9 +112,9 @@ class RabbitMQMessageBus(InfrahubMessageBus):
     async def _initialize_connection(self) -> None:
         self.channel = await self.connection.channel()
 
-        if self.service.component_type == ComponentType.API_SERVER:
+        if self.component_type == ComponentType.API_SERVER:
             await self._initialize_api_server()
-        elif self.service.component_type == ComponentType.GIT_AGENT:
+        elif self.component_type == ComponentType.GIT_AGENT:
             await self._initialize_git_worker()
 
     async def shutdown(self) -> None:
@@ -135,7 +144,7 @@ class RabbitMQMessageBus(InfrahubMessageBus):
 
     async def on_reconnect(
         self,
-        weak: bool = False,  # pylint: disable=unused-argument
+        weak: bool = False,  # noqa: ARG002
     ) -> None:
         self.service.log.info("Reconnected to RabbitMQ, reinitializing connection")
         await self._initialize_connection()
@@ -209,7 +218,11 @@ class RabbitMQMessageBus(InfrahubMessageBus):
         await queue.consume(callback=self.on_message, no_ack=False)
 
     async def publish(
-        self, message: InfrahubMessage, routing_key: str, delay: Optional[MessageTTL] = None, is_retry: bool = False
+        self,
+        message: InfrahubMessage,
+        routing_key: str,
+        delay: Optional[MessageTTL] = None,
+        is_retry: bool = False,  # noqa: ARG002
     ) -> None:
         for enricher in self.message_enrichers:
             await enricher(message)
@@ -233,7 +246,7 @@ class RabbitMQMessageBus(InfrahubMessageBus):
         request_id = log_data.get("request_id", "")
         message.meta = Meta(request_id=request_id, correlation_id=correlation_id, reply_to=self.callback_queue.name)
 
-        await self.service.send(message=message)
+        await self.service.message_bus.send(message=message)
 
         response: AbstractIncomingMessage = await future
         data = ujson.loads(response.body)

infrahub/services/adapters/workflow/__init__.py

@@ -1,26 +1,24 @@
 from __future__ import annotations
 
+from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, Any, Callable, ParamSpec, TypeVar, overload
 
 if TYPE_CHECKING:
-    from infrahub.services import InfrahubServices
+    from infrahub.context import InfrahubContext
     from infrahub.workflows.models import WorkflowDefinition, WorkflowInfo
-
 Return = TypeVar("Return")
 Params = ParamSpec("Params")
 
 FuncType = Callable[Params, Return]
 
 
-class InfrahubWorkflow:
-    async def initialize(self, service: InfrahubServices) -> None:
-        """Initialize the Workflow engine"""
-
+class InfrahubWorkflow(ABC):
     @overload
     async def execute_workflow(
         self,
         workflow: WorkflowDefinition,
         expected_return: type[Return],
+        context: InfrahubContext | None = ...,
         parameters: dict[str, Any] | None = ...,
         tags: list[str] | None = ...,
     ) -> Return: ...
@@ -30,23 +28,28 @@ class InfrahubWorkflow:
         self,
         workflow: WorkflowDefinition,
         expected_return: None = ...,
+        context: InfrahubContext | None = ...,
         parameters: dict[str, Any] | None = ...,
         tags: list[str] | None = ...,
     ) -> Any: ...
 
+    @abstractmethod
     async def execute_workflow(
         self,
         workflow: WorkflowDefinition,
         expected_return: type[Return] | None = None,
+        context: InfrahubContext | None = None,
         parameters: dict[str, Any] | None = None,
         tags: list[str] | None = None,
     ) -> Any:
-        raise NotImplementedError("InfrahubWorkflow.execute_workflow is an abstract method")
+        raise NotImplementedError()
 
+    @abstractmethod
     async def submit_workflow(
         self,
         workflow: WorkflowDefinition,
+        context: InfrahubContext | None = None,
         parameters: dict[str, Any] | None = None,
         tags: list[str] | None = None,
     ) -> WorkflowInfo:
-        raise NotImplementedError("InfrahubWorkflow.submit_workflow is an abstract method")
+        raise NotImplementedError()

infrahub/services/adapters/workflow/local.py

@@ -1,27 +1,48 @@
+from __future__ import annotations
+
 import uuid
-from typing import Any
+from typing import Any, Optional
+
+from typing_extensions import TYPE_CHECKING
 
+from infrahub.workers.utils import inject_context_parameter, inject_service_parameter
 from infrahub.workflows.models import WorkflowDefinition, WorkflowInfo
 
 from . import InfrahubWorkflow, Return
 
+if TYPE_CHECKING:
+    from infrahub.context import InfrahubContext
+    from infrahub.services import InfrahubServices
+
 
 class WorkflowLocalExecution(InfrahubWorkflow):
+    service: Optional[InfrahubServices] = None  # needed for local injections
+
     async def execute_workflow(
         self,
         workflow: WorkflowDefinition,
-        expected_return: type[Return] | None = None,
+        expected_return: type[Return] | None = None,  # noqa: ARG002
+        context: InfrahubContext | None = None,
         parameters: dict[str, Any] | None = None,
-        tags: list[str] | None = None,
+        tags: list[str] | None = None,  # noqa: ARG002
     ) -> Any:
-        fn = workflow.get_function()
-        return await fn(**parameters or {})
+        if self.service is None:
+            raise ValueError("WorkflowLocalExecution.service is not initialized")
+
+        flow_func = workflow.load_function()
+        parameters = dict(parameters) if parameters is not None else {}  # avoid mutating input parameters
+        inject_service_parameter(func=flow_func, parameters=parameters, service=self.service)
+        inject_context_parameter(func=flow_func, parameters=parameters, context=context)
+
+        parameters = flow_func.validate_parameters(parameters=parameters)
+        return await flow_func(**parameters)
 
     async def submit_workflow(
         self,
         workflow: WorkflowDefinition,
+        context: InfrahubContext | None = None,
         parameters: dict[str, Any] | None = None,
-        tags: list[str] | None = None,
+        tags: list[str] | None = None,  # noqa: ARG002
     ) -> WorkflowInfo:
-        await self.execute_workflow(workflow=workflow, parameters=parameters)
+        await self.execute_workflow(workflow=workflow, context=context, parameters=parameters)
         return WorkflowInfo(id=uuid.uuid4())

infrahub/services/adapters/workflow/worker.py

@@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Any, overload
 from prefect.client.schemas.objects import StateType
 from prefect.deployments import run_deployment
 
+from infrahub.workers.utils import inject_context_parameter
 from infrahub.workflows.initialization import setup_task_manager
 from infrahub.workflows.models import WorkflowInfo
 
@@ -13,15 +14,14 @@ from . import InfrahubWorkflow, Return
 if TYPE_CHECKING:
     from prefect.client.schemas.objects import FlowRun
 
-    from infrahub.services import InfrahubServices
+    from infrahub.context import InfrahubContext
     from infrahub.workflows.models import WorkflowDefinition
 
 
 class WorkflowWorkerExecution(InfrahubWorkflow):
-    async def initialize(self, service: InfrahubServices) -> None:
-        """Initialize the Workflow engine"""
-
-        if await service.component.is_primary_api():
+    @staticmethod
+    async def initialize(component_is_primary_server: bool) -> None:
+        if component_is_primary_server:
             await setup_task_manager()
 
     @overload
@@ -29,6 +29,7 @@ class WorkflowWorkerExecution(InfrahubWorkflow):
         self,
         workflow: WorkflowDefinition,
         expected_return: type[Return],
+        context: InfrahubContext | None = None,
         parameters: dict[str, Any] | None = ...,
         tags: list[str] | None = ...,
     ) -> Return: ...
@@ -38,17 +39,24 @@ class WorkflowWorkerExecution(InfrahubWorkflow):
         self,
         workflow: WorkflowDefinition,
         expected_return: None = ...,
+        context: InfrahubContext | None = ...,
         parameters: dict[str, Any] | None = ...,
         tags: list[str] | None = ...,
     ) -> Any: ...
 
+    # TODO Make expected_return mandatory and remove above overloads.
     async def execute_workflow(
         self,
         workflow: WorkflowDefinition,
-        expected_return: type[Return] | None = None,
+        expected_return: type[Return] | None = None,  # noqa: ARG002
+        context: InfrahubContext | None = None,
         parameters: dict[str, Any] | None = None,
         tags: list[str] | None = None,
     ) -> Any:
+        flow_func = workflow.load_function()
+        parameters = dict(parameters) if parameters is not None else {}
+        inject_context_parameter(func=flow_func, parameters=parameters, context=context)
+
         response: FlowRun = await run_deployment(
             name=workflow.full_name, poll_interval=1, parameters=parameters or {}, tags=tags
         )  # type: ignore[return-value, misc]
@@ -61,7 +69,15 @@ class WorkflowWorkerExecution(InfrahubWorkflow):
         return await response.state.result(raise_on_failure=True, fetch=True)  # type: ignore[call-overload]
 
     async def submit_workflow(
-        self, workflow: WorkflowDefinition, parameters: dict[str, Any] | None = None, tags: list[str] | None = None
+        self,
+        workflow: WorkflowDefinition,
+        context: InfrahubContext | None = None,
+        parameters: dict[str, Any] | None = None,
+        tags: list[str] | None = None,
     ) -> WorkflowInfo:
+        flow_func = workflow.load_function()
+        parameters = dict(parameters) if parameters is not None else {}
+        inject_context_parameter(func=flow_func, parameters=parameters, context=context)
+
         flow_run = await run_deployment(name=workflow.full_name, timeout=0, parameters=parameters or {}, tags=tags)  # type: ignore[return-value, misc]
         return WorkflowInfo.from_flow(flow_run=flow_run)

infrahub/services/component.py

@@ -3,54 +3,57 @@ from __future__ import annotations
 import re
 from typing import TYPE_CHECKING, Any, Optional
 
+from attr import dataclass
+
 from infrahub.components import ComponentType
 from infrahub.core.constants import GLOBAL_BRANCH_NAME
 from infrahub.core.registry import registry
 from infrahub.core.timestamp import Timestamp
-from infrahub.exceptions import InitializationError
+from infrahub.log import get_logger
 from infrahub.message_bus import messages
 from infrahub.message_bus.types import KVTTL
 from infrahub.worker import WORKER_IDENTITY
 
 if TYPE_CHECKING:
-    from infrahub.services import InfrahubServices
+    from infrahub.database import InfrahubDatabase
+    from infrahub.services import InfrahubCache, InfrahubMessageBus
 
 PRIMARY_API_SERVER = "workers:primary:api_server"
 WORKER_MATCH = re.compile(r":worker:([^:]+)")
 
+log = get_logger()
 
-class InfrahubComponent:
-    def __init__(self) -> None:
-        self._service: Optional[InfrahubServices] = None
-
-    @property
-    def service(self) -> InfrahubServices:
-        if not self._service:
-            raise InitializationError("Component has not been initialized")
 
-        return self._service
+@dataclass
+class InfrahubComponent:
+    cache: InfrahubCache
+    db: InfrahubDatabase
+    message_bus: InfrahubMessageBus
+    component_type: ComponentType
+
+    @classmethod
+    async def new(
+        cls, cache: InfrahubCache, db: InfrahubDatabase, message_bus: InfrahubMessageBus, component_type: ComponentType
+    ) -> InfrahubComponent:
+        component = cls(cache=cache, db=db, message_bus=message_bus, component_type=component_type)
+        await component.refresh_heartbeat()
+        return component
 
     @property
     def component_names(self) -> list[str]:
         names = []
-        if self.service.component_type == ComponentType.API_SERVER:
+        if self.component_type == ComponentType.API_SERVER:
             names.append("api_server")
-        elif self.service.component_type == ComponentType.GIT_AGENT:
+        elif self.component_type == ComponentType.GIT_AGENT:
             names.append("git_agent")
         return names
 
-    async def initialize(self, service: InfrahubServices) -> None:
-        """Initialize the Message bus"""
-        self._service = service
-
-        await self.refresh_heartbeat()
-
-    async def is_primary_api(self) -> bool:
-        primary_identity = await self.service.cache.get(PRIMARY_API_SERVER)
+    async def is_primary_gunicorn_worker(self) -> bool:
+        primary_identity = await self.cache.get(PRIMARY_API_SERVER)
         return primary_identity == WORKER_IDENTITY
 
     async def list_workers(self, branch: str, schema_hash: bool) -> list[WorkerInfo]:
-        keys = await self.service.cache.list_keys(filter_pattern="workers:*")
+        keys = await self.cache.list_keys(filter_pattern="workers:*")
 
         workers: dict[str, WorkerInfo] = {}
         for key in keys:
@@ -64,9 +67,9 @@ class InfrahubComponent:
         schema_hash_keys = []
         if schema_hash:
             schema_hash_keys = [key for key in keys if f":schema_hash:branch:{branch}" in key]
-            response = await self.service.cache.get_values(keys=schema_hash_keys)
+            response = await self.cache.get_values(keys=schema_hash_keys)
 
-        for key, value in zip(schema_hash_keys, response):
+        for key, value in zip(schema_hash_keys, response, strict=False):
             if match := WORKER_MATCH.search(key):
                 identity = match.group(1)
                 workers[identity].add_value(key=key, value=value)
@@ -82,14 +85,14 @@ class InfrahubComponent:
 
             # Use branch name if we cannot find branch id in cache
             branch_id: Optional[str] = None
-            if branch_obj := await registry.get_branch(branch=branch, db=self.service.database):
+            if branch_obj := await registry.get_branch(branch=branch, db=self.db):
                 branch_id = str(branch_obj.uuid)
 
             if not branch_id:
                 branch_id = branch
 
             for component in self.component_names:
-                await self.service.cache.set(
+                await self.cache.set(
                     key=f"workers:schema_hash:branch:{branch_id}:{component}:worker:{WORKER_IDENTITY}",
                     value=hash_value,
                     expires=KVTTL.TWO_HOURS,
@@ -97,29 +100,29 @@ class InfrahubComponent:
 
     async def refresh_heartbeat(self) -> None:
         for component in self.component_names:
-            await self.service.cache.set(
+            await self.cache.set(
                 key=f"workers:active:{component}:worker:{WORKER_IDENTITY}",
                 value=Timestamp().to_string(),
                 expires=KVTTL.FIFTEEN,
             )
-        if self.service.component_type == ComponentType.API_SERVER:
+        if self.component_type == ComponentType.API_SERVER:
             await self._set_primary_api_server()
-        await self.service.cache.set(
+        await self.cache.set(
             key=f"workers:worker:{WORKER_IDENTITY}", value=Timestamp().to_string(), expires=KVTTL.TWO_HOURS
         )
 
     async def _set_primary_api_server(self) -> None:
-        result = await self.service.cache.set(
+        result = await self.cache.set(
             key=PRIMARY_API_SERVER, value=WORKER_IDENTITY, expires=KVTTL.FIFTEEN, not_exists=True
         )
         if result:
-            await self.service.send(message=messages.EventWorkerNewPrimaryAPI(worker_id=WORKER_IDENTITY))
+            await self.message_bus.send(message=messages.EventWorkerNewPrimaryAPI(worker_id=WORKER_IDENTITY))
         else:
-            self.service.log.debug("Primary node already set")
-            primary_id = await self.service.cache.get(key=PRIMARY_API_SERVER)
+            log.debug("Primary node already set")
+            primary_id = await self.cache.get(key=PRIMARY_API_SERVER)
             if primary_id == WORKER_IDENTITY:
-                self.service.log.debug("Primary node set but same as ours, refreshing lifetime")
-                await self.service.cache.set(key=PRIMARY_API_SERVER, value=WORKER_IDENTITY, expires=KVTTL.FIFTEEN)
+                log.debug("Primary node set but same as ours, refreshing lifetime")
+                await self.cache.set(key=PRIMARY_API_SERVER, value=WORKER_IDENTITY, expires=KVTTL.FIFTEEN)
 
 
 class WorkerInfo:

infrahub/services/scheduler.py

@@ -5,14 +5,19 @@ import random
 from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
+from typing_extensions import Optional
+
 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 +28,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: Optional[InfrahubServices]
+
+    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 +47,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 +58,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,