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

infrahub/api/artifact.py

@@ -5,7 +5,14 @@ from typing import TYPE_CHECKING
 from fastapi import APIRouter, Body, Depends, Request, Response
 from pydantic import BaseModel, Field
 
-from infrahub.api.dependencies import BranchParams, get_branch_params, get_current_user, get_db, get_permission_manager
+from infrahub.api.dependencies import (
+    BranchParams,
+    get_branch_params,
+    get_context,
+    get_current_user,
+    get_db,
+    get_permission_manager,
+)
 from infrahub.core import registry
 from infrahub.core.account import ObjectPermission
 from infrahub.core.constants import GLOBAL_BRANCH_NAME, InfrahubKind, PermissionAction
@@ -19,7 +26,9 @@ from infrahub.workflows.catalogue import REQUEST_ARTIFACT_DEFINITION_GENERATE
 
 if TYPE_CHECKING:
     from infrahub.auth import AccountSession
+    from infrahub.context import InfrahubContext
     from infrahub.permissions import PermissionManager
+    from infrahub.services import InfrahubServices
 
 log = get_logger()
 router = APIRouter(prefix="/artifact")
@@ -57,12 +66,13 @@ async def generate_artifact(
     request: Request,
     artifact_definition_id: str,
     payload: ArtifactGeneratePayload = Body(
-        ArtifactGeneratePayload(),
+        ArtifactGeneratePayload(),  # noqa: B008
         description="Payload of the request, can be used to limit the scope of the query to a specific list of hosts",
     ),
     db: InfrahubDatabase = Depends(get_db),
     branch_params: BranchParams = Depends(get_branch_params),
     permission_manager: PermissionManager = Depends(get_permission_manager),
+    context: InfrahubContext = Depends(get_context),
 ) -> None:
     permission_decision = (
         PermissionDecisionFlag.ALLOW_DEFAULT
@@ -83,7 +93,7 @@ async def generate_artifact(
         branch=branch_params.branch,
     )
 
-    service = request.app.state.service
+    service: InfrahubServices = request.app.state.service
     model = RequestArtifactDefinitionGenerate(
         artifact_definition_id=artifact_definition.id,
         artifact_definition_name=artifact_definition.name.value,
@@ -91,4 +101,6 @@ async def generate_artifact(
         limit=payload.nodes,
     )
 
-    await service.workflow.submit_workflow(workflow=REQUEST_ARTIFACT_DEFINITION_GENERATE, parameters={"model": model})
+    await service.workflow.submit_workflow(
+        workflow=REQUEST_ARTIFACT_DEFINITION_GENERATE, context=context, parameters={"model": model}
+    )

infrahub/api/dependencies.py

@@ -8,6 +8,7 @@ from pydantic import BaseModel, ConfigDict
 
 from infrahub import config
 from infrahub.auth import AccountSession, authentication_token, validate_jwt_access_token, validate_jwt_refresh_token
+from infrahub.context import InfrahubContext
 from infrahub.core.branch import Branch  # noqa: TC001
 from infrahub.core.registry import registry
 from infrahub.core.timestamp import Timestamp
@@ -133,3 +134,10 @@ async def get_permission_manager(
     await permission_manager.load_permissions(db=db, branch=branch_params.branch)
 
     return permission_manager
+
+
+async def get_context(
+    branch: Branch = Depends(get_branch_dep),
+    account_session: AccountSession = Depends(get_current_user),
+) -> InfrahubContext:
+    return InfrahubContext.init(branch=branch, account=account_session)

infrahub/api/oauth2.py

@@ -22,7 +22,6 @@ if TYPE_CHECKING:
     from infrahub.database import InfrahubDatabase
     from infrahub.services import InfrahubServices
 
-# pylint: disable=R0801
 
 log = get_logger()
 router = APIRouter(prefix="/oauth2")

infrahub/api/oidc.py

@@ -24,7 +24,6 @@ if TYPE_CHECKING:
     from infrahub.database import InfrahubDatabase
     from infrahub.services import InfrahubServices
 
-# pylint: disable=R0801
 
 log = get_logger()
 router = APIRouter(prefix="/oidc")

infrahub/api/query.py

@@ -7,6 +7,7 @@ from graphql import graphql
 from pydantic import BaseModel, Field
 
 from infrahub.api.dependencies import BranchParams, get_branch_params, get_current_user, get_db
+from infrahub.context import InfrahubContext
 from infrahub.core import registry
 from infrahub.core.constants import InfrahubKind
 from infrahub.core.protocols import CoreGraphQLQuery
@@ -57,12 +58,21 @@ async def execute_query(
         db=db, id=query_id, kind=CoreGraphQLQuery, branch=branch_params.branch, at=branch_params.at
     )
 
+    context = InfrahubContext.init(branch=branch_params.branch, account=account_session)
+
     gql_params = await prepare_graphql_params(
-        db=db, branch=branch_params.branch, at=branch_params.at, account_session=account_session
+        db=db,
+        branch=branch_params.branch,
+        at=branch_params.at,
+        account_session=account_session,
+        service=request.app.state.service,
     )
+    schema_branch = db.schema.get_schema_branch(name=branch_params.branch.name)
+
     analyzed_query = InfrahubGraphQLQueryAnalyzer(
         query=gql_query.query.value,
         schema=gql_params.schema,
+        schema_branch=schema_branch,
         branch=branch_params.branch,
     )
     await permission_checker.check(
@@ -97,9 +107,7 @@ async def execute_query(
     GRAPHQL_QUERY_HEIGHT_METRICS.labels(**labels).observe(await analyzed_query.calculate_height())
     GRAPHQL_QUERY_VARS_METRICS.labels(**labels).observe(len(analyzed_query.variables))
     GRAPHQL_TOP_LEVEL_QUERIES_METRICS.labels(**labels).observe(analyzed_query.nbr_queries)
-    GRAPHQL_QUERY_OBJECTS_METRICS.labels(**labels).observe(
-        len(await analyzed_query.get_models_in_use(types=gql_params.context.types))
-    )
+    GRAPHQL_QUERY_OBJECTS_METRICS.labels(**labels).observe(len(analyzed_query.query_report.impacted_models))
 
     response_payload: dict[str, Any] = {"data": data}
 
@@ -111,11 +119,13 @@ async def execute_query(
             branch=branch_params.branch.name,
             query_id=gql_query.id,
             query_name=gql_query.name.value,
-            related_node_ids=sorted(list(related_node_ids)),
+            related_node_ids=sorted(related_node_ids),
             subscribers=sorted(subscribers),
             params=params,
         )
-        await service.workflow.submit_workflow(workflow=GRAPHQL_QUERY_GROUP_UPDATE, parameters={"model": model})
+        await service.workflow.submit_workflow(
+            workflow=GRAPHQL_QUERY_GROUP_UPDATE, context=context, parameters={"model": model}
+        )
 
     return response_payload
 
@@ -124,7 +134,8 @@ async def execute_query(
 async def graphql_query_post(
     request: Request,
     payload: QueryPayload = Body(
-        QueryPayload(), description="Payload of the request, must be used to provide the variables"
+        QueryPayload(),  # noqa: B008
+        description="Payload of the request, must be used to provide the variables",
     ),
     query_id: str = Path(description="ID or Name of the GraphQL query to execute"),
     subscribers: list[str] = Query(

infrahub/api/schema.py

@@ -13,7 +13,7 @@ from pydantic import (
 from starlette.responses import JSONResponse
 
 from infrahub import lock
-from infrahub.api.dependencies import get_branch_dep, get_current_user, get_db, get_permission_manager
+from infrahub.api.dependencies import get_branch_dep, get_context, get_current_user, get_db, get_permission_manager
 from infrahub.api.exceptions import SchemaNotValidError
 from infrahub.core import registry
 from infrahub.core.account import GlobalPermission
@@ -25,7 +25,7 @@ from infrahub.core.models import (  # noqa: TC001
     SchemaDiff,
     SchemaUpdateValidationResult,
 )
-from infrahub.core.schema import GenericSchema, MainSchemaTypes, NodeSchema, ProfileSchema, SchemaRoot
+from infrahub.core.schema import GenericSchema, MainSchemaTypes, NodeSchema, ProfileSchema, SchemaRoot, TemplateSchema
 from infrahub.core.schema.constants import SchemaNamespace  # noqa: TC001
 from infrahub.core.validators.models.validate_migration import (
     SchemaValidateMigrationData,
@@ -44,6 +44,7 @@ if TYPE_CHECKING:
     from typing_extensions import Self
 
     from infrahub.auth import AccountSession
+    from infrahub.context import InfrahubContext
     from infrahub.core.schema.schema_branch import SchemaBranch
     from infrahub.permissions import PermissionManager
     from infrahub.services import InfrahubServices
@@ -86,11 +87,17 @@ class APIProfileSchema(ProfileSchema, APISchemaMixin):
     hash: str
 
 
+class APITemplateSchema(TemplateSchema, APISchemaMixin):
+    api_kind: str | None = Field(default=None, alias="kind", validate_default=True)
+    hash: str
+
+
 class SchemaReadAPI(BaseModel):
     main: str = Field(description="Main hash for the entire schema")
     nodes: list[APINodeSchema] = Field(default_factory=list)
     generics: list[APIGenericSchema] = Field(default_factory=list)
     profiles: list[APIProfileSchema] = Field(default_factory=list)
+    templates: list[APITemplateSchema] = Field(default_factory=list)
     namespaces: list[SchemaNamespace] = Field(default_factory=list)
 
 
@@ -190,6 +197,11 @@ async def get_schema(
             for value in all_schemas
             if isinstance(value, ProfileSchema) and value.namespace != "Internal"
         ],
+        templates=[
+            APITemplateSchema.from_schema(value)
+            for value in all_schemas
+            if isinstance(value, TemplateSchema) and value.namespace != "Internal"
+        ],
         namespaces=schema_branch.get_namespaces(),
     )
 
@@ -206,15 +218,16 @@ async def get_schema_summary(
 @router.get("/{schema_kind}")
 async def get_schema_by_kind(
     schema_kind: str, branch: Branch = Depends(get_branch_dep), _: AccountSession = Depends(get_current_user)
-) -> APIProfileSchema | APINodeSchema | APIGenericSchema:
+) -> APIProfileSchema | APINodeSchema | APIGenericSchema | APITemplateSchema:
     log.debug("schema_kind_request", branch=branch.name)
 
     schema = registry.schema.get(name=schema_kind, branch=branch, duplicate=False)
 
-    api_schema: dict[str, type[APIProfileSchema | APINodeSchema | APIGenericSchema]] = {
+    api_schema: dict[str, type[APIProfileSchema | APINodeSchema | APIGenericSchema | APITemplateSchema]] = {
         "profile": APIProfileSchema,
         "node": APINodeSchema,
         "generic": APIGenericSchema,
+        "template": APITemplateSchema,
     }
     key = ""
 
@@ -224,6 +237,8 @@ async def get_schema_by_kind(
         key = "node"
     if isinstance(schema, GenericSchema):
         key = "generic"
+    if isinstance(schema, TemplateSchema):
+        key = "template"
 
     return api_schema[key].from_schema(schema=schema)
 
@@ -269,6 +284,7 @@ async def load_schema(
     branch: Branch = Depends(get_branch_dep),
     account_session: AccountSession = Depends(get_current_user),
     permission_manager: PermissionManager = Depends(get_permission_manager),
+    context: InfrahubContext = Depends(get_context),
 ) -> SchemaUpdate:
     permission_manager.raise_for_permission(
         permission=GlobalPermission(
@@ -317,6 +333,7 @@ async def load_schema(
         )
         responses = await service.workflow.execute_workflow(
             workflow=SCHEMA_VALIDATE_MIGRATION,
+            context=context,
             expected_return=list[SchemaValidatorPathResponseData],
             parameters={"message": validate_migration_data},
         )
@@ -361,6 +378,7 @@ async def load_schema(
         )
         migration_error_msgs = await service.workflow.execute_workflow(
             workflow=SCHEMA_APPLY_MIGRATION,
+            context=context,
             expected_return=list[str],
             parameters={"message": apply_migration_data},
         )
@@ -373,9 +391,15 @@ async def load_schema(
     log_data = get_log_data()
     request_id = log_data.get("request_id", "")
     event = SchemaUpdatedEvent(
-        branch=branch.name,
+        branch_name=branch.name,
         schema_hash=branch.active_schema_hash.main,
-        meta=EventMeta(initiator_id=WORKER_IDENTITY, request_id=request_id, account_id=account_session.account_id),
+        meta=EventMeta(
+            initiator_id=WORKER_IDENTITY,
+            request_id=request_id,
+            account_id=account_session.account_id,
+            branch=branch,
+            context=context,
+        ),
     )
     await service.event.send(event=event)
 
@@ -387,6 +411,7 @@ async def check_schema(
     request: Request,
     schemas: SchemasLoadAPI,
     branch: Branch = Depends(get_branch_dep),
+    context: InfrahubContext = Depends(get_context),
     _: AccountSession = Depends(get_current_user),
 ) -> JSONResponse:
     service: InfrahubServices = request.app.state.service
@@ -413,6 +438,7 @@ async def check_schema(
     )
     responses = await service.workflow.execute_workflow(
         workflow=SCHEMA_VALIDATE_MIGRATION,
+        context=context,
         expected_return=list[SchemaValidatorPathResponseData],
         parameters={"message": validate_migration_data},
     )

infrahub/api/transformation.py

@@ -9,6 +9,7 @@ from starlette.responses import JSONResponse, PlainTextResponse
 from infrahub.api.dependencies import (
     BranchParams,
     get_branch_params,
+    get_context,
     get_current_user,
     get_db,
 )
@@ -28,8 +29,8 @@ from infrahub.workflows.catalogue import TRANSFORM_JINJA2_RENDER, TRANSFORM_PYTH
 
 if TYPE_CHECKING:
     from infrahub.auth import AccountSession
+    from infrahub.context import InfrahubContext
     from infrahub.services import InfrahubServices
-
 router = APIRouter()
 
 
@@ -38,6 +39,7 @@ async def transform_python(
     request: Request,
     transform_id: str,
     db: InfrahubDatabase = Depends(get_db),
+    context: InfrahubContext = Depends(get_context),
     branch_params: BranchParams = Depends(get_branch_params),
     _: AccountSession = Depends(get_current_user),
 ) -> JSONResponse:
@@ -62,7 +64,9 @@ async def transform_python(
             message="Repository doesn't have a commit",
         )
 
-    gql_params = await prepare_graphql_params(db=request.app.state.db, branch=branch_params.branch, at=branch_params.at)
+    gql_params = await prepare_graphql_params(
+        db=request.app.state.db, branch=branch_params.branch, at=branch_params.at, service=request.app.state.service
+    )
 
     result = await graphql(
         schema=gql_params.schema,
@@ -88,7 +92,7 @@ async def transform_python(
     )
 
     response = await service.workflow.execute_workflow(
-        workflow=TRANSFORM_PYTHON_RENDER, parameters={"message": message}
+        workflow=TRANSFORM_PYTHON_RENDER, context=context, parameters={"message": message}
     )
     return JSONResponse(content=response)
 
@@ -99,6 +103,7 @@ async def transform_jinja2(
     transform_id: str = Path(description="ID or Name of the Jinja2 Transform to render"),
     db: InfrahubDatabase = Depends(get_db),
     branch_params: BranchParams = Depends(get_branch_params),
+    context: InfrahubContext = Depends(get_context),
     _: AccountSession = Depends(get_current_user),
 ) -> PlainTextResponse:
     params = {key: value for key, value in request.query_params.items() if key not in ["branch", "at"]}
@@ -122,7 +127,9 @@ async def transform_jinja2(
             message="Repository doesn't have a commit",
         )
 
-    gql_params = await prepare_graphql_params(db=request.app.state.db, branch=branch_params.branch, at=branch_params.at)
+    gql_params = await prepare_graphql_params(
+        db=request.app.state.db, branch=branch_params.branch, at=branch_params.at, service=request.app.state.service
+    )
 
     result = await graphql(
         schema=gql_params.schema,
@@ -148,6 +155,6 @@ async def transform_jinja2(
     service: InfrahubServices = request.app.state.service
 
     response = await service.workflow.execute_workflow(
-        workflow=TRANSFORM_JINJA2_RENDER, expected_return=str, parameters={"message": message}
+        workflow=TRANSFORM_JINJA2_RENDER, context=context, expected_return=str, parameters={"message": message}
     )
     return PlainTextResponse(content=response)

infrahub/{message_bus/messages/check_artifact_create.py → artifacts/models.py}

@@ -1,11 +1,11 @@
 from typing import Optional
 
-from pydantic import Field
+from pydantic import BaseModel, Field
 
-from infrahub.message_bus import InfrahubMessage
+from infrahub.context import InfrahubContext
 
 
-class CheckArtifactCreate(InfrahubMessage):
+class CheckArtifactCreate(BaseModel):
     """Runs a check to verify the creation of an artifact."""
 
     artifact_name: str = Field(..., description="Name of the artifact")
@@ -19,9 +19,11 @@ class CheckArtifactCreate(InfrahubMessage):
     repository_kind: str = Field(..., description="The kind of the Repository")
     branch_name: str = Field(..., description="The branch where the check is run")
     target_id: str = Field(..., description="The ID of the target object for this artifact")
+    target_kind: str = Field(..., description="The kind of the target object for this artifact")
     target_name: str = Field(..., description="Name of the artifact target")
     artifact_id: Optional[str] = Field(default=None, description="The id of the artifact if it previously existed")
     query: str = Field(..., description="The name of the query to use when collecting data")
     timeout: int = Field(..., description="Timeout for requests used to generate this artifact")
     variables: dict = Field(..., description="Input variables when generating the artifact")
     validator_id: str = Field(..., description="The ID of the validator")
+    context: InfrahubContext = Field(..., description="The context of the task")

infrahub/{message_bus/operations/check/artifact.py → artifacts/tasks.py}

@@ -2,38 +2,38 @@ from typing import Union
 
 from prefect import flow
 
+from infrahub.artifacts.models import CheckArtifactCreate
 from infrahub.core.constants import InfrahubKind, ValidatorConclusion
 from infrahub.core.timestamp import Timestamp
-from infrahub.git.repository import InfrahubReadOnlyRepository, InfrahubRepository
-from infrahub.log import get_logger
-from infrahub.message_bus import messages
+from infrahub.git import InfrahubReadOnlyRepository, InfrahubRepository
 from infrahub.services import InfrahubServices
 from infrahub.tasks.artifact import define_artifact
-from infrahub.tasks.check import set_check_status
 from infrahub.workflows.utils import add_tags
 
-log = get_logger()
-
 
 @flow(name="git-repository-check-artifact-create", flow_run_name="Check artifact creation")
-async def create(message: messages.CheckArtifactCreate, service: InfrahubServices) -> None:
-    await add_tags(branches=[message.branch_name], nodes=[message.target_id])
-    validator = await service.client.get(
-        kind=InfrahubKind.ARTIFACTVALIDATOR, id=message.validator_id, include=["checks"]
-    )
+async def create(model: CheckArtifactCreate, service: InfrahubServices) -> ValidatorConclusion:
+    await add_tags(branches=[model.branch_name], nodes=[model.target_id])
+    validator = await service.client.get(kind=InfrahubKind.ARTIFACTVALIDATOR, id=model.validator_id, include=["checks"])
 
+    repo: InfrahubReadOnlyRepository | InfrahubRepository
     if InfrahubKind.READONLYREPOSITORY:
         repo = await InfrahubReadOnlyRepository.init(
-            id=message.repository_id, name=message.repository_name, client=service.client
+            id=model.repository_id,
+            name=model.repository_name,
+            client=service.client,
+            service=service,
         )
     else:
         repo = await InfrahubRepository.init(
-            id=message.repository_id, name=message.repository_name, client=service.client
+            id=model.repository_id,
+            name=model.repository_name,
+            client=service.client,
+            service=service,
         )
 
-    artifact = await define_artifact(message=message, service=service)
+    artifact, artifact_created = await define_artifact(model=model, service=service)
 
-    conclusion = ValidatorConclusion.SUCCESS.value
     severity = "info"
     artifact_result: dict[str, Union[str, bool, None]] = {
         "changed": None,
@@ -44,22 +44,23 @@ async def create(message: messages.CheckArtifactCreate, service: InfrahubService
     check_message = "Failed to render artifact"
 
     try:
-        result = await repo.render_artifact(artifact=artifact, message=message)
+        result = await repo.render_artifact(artifact=artifact, artifact_created=artifact_created, message=model)
         artifact_result["changed"] = result.changed
         artifact_result["checksum"] = result.checksum
         artifact_result["artifact_id"] = result.artifact_id
         artifact_result["storage_id"] = result.storage_id
         check_message = "Artifact rendered successfully"
+        conclusion = ValidatorConclusion.SUCCESS
 
-    except Exception as exc:  # pylint: disable=broad-except
-        conclusion = ValidatorConclusion.FAILURE.value
+    except Exception as exc:
         artifact.status.value = "Error"
+        await artifact.save()
         severity = "critical"
+        conclusion = ValidatorConclusion.FAILURE
         check_message += f": {str(exc)}"
-        await artifact.save()
 
     check = None
-    check_name = f"{message.artifact_name}: {message.target_name}"
+    check_name = f"{model.artifact_name}: {model.target_name}"
     existing_check = await service.client.filters(
         kind=InfrahubKind.ARTIFACTCHECK, validator__ids=validator.id, name__value=check_name
     )
@@ -68,7 +69,7 @@ async def create(message: messages.CheckArtifactCreate, service: InfrahubService
 
     if check:
         check.created_at.value = Timestamp().to_string()
-        check.conclusion.value = conclusion
+        check.conclusion.value = conclusion.value
         check.severity.value = severity
         check.changed.value = artifact_result["changed"]
         check.checksum.value = artifact_result["checksum"]
@@ -80,12 +81,12 @@ async def create(message: messages.CheckArtifactCreate, service: InfrahubService
             kind=InfrahubKind.ARTIFACTCHECK,
             data={
                 "name": check_name,
-                "origin": message.repository_id,
+                "origin": model.repository_id,
                 "kind": "ArtifactDefinition",
-                "validator": message.validator_id,
+                "validator": model.validator_id,
                 "created_at": Timestamp().to_string(),
                 "message": check_message,
-                "conclusion": conclusion,
+                "conclusion": conclusion.value,
                 "severity": severity,
                 "changed": artifact_result["changed"],
                 "checksum": artifact_result["checksum"],
@@ -95,4 +96,4 @@ async def create(message: messages.CheckArtifactCreate, service: InfrahubService
         )
         await check.save()
 
-    await set_check_status(message=message, conclusion=conclusion, service=service)
+    return conclusion

infrahub/cli/__init__.py

@@ -12,8 +12,6 @@ from infrahub.cli.tasks import app as tasks_app
 from infrahub.core.initialization import initialization
 from infrahub.database import InfrahubDatabase, get_db
 
-# pylint: disable=import-outside-toplevel
-
 app = typer.Typer(name="Infrahub CLI", pretty_exceptions_enable=False)
 
 

infrahub/cli/db.py

@@ -44,7 +44,7 @@ from infrahub.core.validators.models.validate_migration import SchemaValidateMig
 from infrahub.core.validators.tasks import schema_validate_migrations
 from infrahub.database import DatabaseType
 from infrahub.log import get_logger
-from infrahub.services import InfrahubServices, services
+from infrahub.services import InfrahubServices
 from infrahub.services.adapters.message_bus.local import BusSimulator
 from infrahub.services.adapters.workflow.local import WorkflowLocalExecution
 
@@ -193,7 +193,7 @@ async def migrate(
 
 
 @app.command()
-async def update_core_schema(  # pylint: disable=too-many-statements
+async def update_core_schema(
     ctx: typer.Context,
     debug: bool = typer.Option(False, help="Enable advanced logging and troubleshooting"),
     config_file: str = typer.Argument("infrahub.toml", envvar="INFRAHUB_CONFIG"),
@@ -216,10 +216,9 @@ async def update_core_schema(  # pylint: disable=too-many-statements
             # ----------------------------------------------------------
             # Initialize Schema and Registry
             # ----------------------------------------------------------
-            service = InfrahubServices(
-                database=db, message_bus=BusSimulator(database=db), workflow=WorkflowLocalExecution()
+            service = await InfrahubServices.new(
+                database=db, message_bus=BusSimulator(), workflow=WorkflowLocalExecution()
             )
-            services.prepare(service)
             await initialize_registry(db=db)
 
             default_branch = registry.get_branch_from_registry(branch=registry.default_branch)
@@ -269,7 +268,7 @@ async def update_core_schema(  # pylint: disable=too-many-statements
                 schema_branch=candidate_schema,
                 constraints=result.constraints,
             )
-            responses = await schema_validate_migrations(message=validate_migration_data)
+            responses = await schema_validate_migrations(message=validate_migration_data, service=service)
             error_messages = [violation.message for response in responses for violation in response.violations]
             if error_messages:
                 rprint(f"{error_badge} | Unable to update the schema, due to failed validations")
@@ -311,7 +310,7 @@ async def update_core_schema(  # pylint: disable=too-many-statements
                 previous_schema=origin_schema,
                 migrations=result.migrations,
             )
-            migration_error_msgs = await schema_apply_migrations(message=apply_migration_data)
+            migration_error_msgs = await schema_apply_migrations(message=apply_migration_data, service=service)
 
             if migration_error_msgs:
                 rprint(f"{error_badge} | Some error(s) happened while running the schema migrations")

infrahub/cli/events.py

@@ -7,6 +7,7 @@ from infrahub_sdk.async_typer import AsyncTyper
 from rich import print as rprint
 
 from infrahub import config
+from infrahub.components import ComponentType
 from infrahub.services import InfrahubServices
 from infrahub.services.adapters.message_bus.rabbitmq import RabbitMQMessageBus
 
@@ -24,9 +25,13 @@ async def listen(
 ) -> None:
     """Listen to event in the Events bus and print them."""
     config.load_and_exit(config_file)
-    broker = RabbitMQMessageBus()
-    service = InfrahubServices()
-    await broker.initialize(service=service)
+
+    # Below code seems to be a duplicate of `RabbitMQMessageBus._initialize_api_server`.
+    #  Also, here we are using ComponentType.NONE so RabbitMQMessageBus will NOT instantiate connection.
+    #  We might want to factorize here and use ComponentType.GitAgent so broker instantiates connection correctly.
+    component_type = ComponentType.NONE
+    broker = await RabbitMQMessageBus.new(component_type=component_type)
+    _ = await InfrahubServices.new(message_bus=broker, component_type=component_type)
 
     queue = await broker.channel.declare_queue(exclusive=True)
     exchange_name = f"{config.SETTINGS.broker.namespace}.events"

infrahub/cli/git_agent.py

@@ -36,7 +36,7 @@ log = get_logger()
 shutdown_event = asyncio.Event()
 
 
-def signal_handler(*args: Any, **kwargs: Any) -> None:  # pylint: disable=unused-argument
+def signal_handler(*args: Any, **kwargs: Any) -> None:  # noqa: ARG001
     shutdown_event.set()
 
 
@@ -99,7 +99,7 @@ async def start(
         await client.branch.all()
     except SdkError as exc:
         log.error(f"Error in communication with Infrahub: {exc.message}")
-        raise typer.Exit(1)
+        raise typer.Exit(1) from None
 
     # Initialize trace
     if config.SETTINGS.trace.enable:
@@ -119,22 +119,24 @@ async def start(
         else WorkflowLocalExecution()
     )
 
+    component_type = ComponentType.GIT_AGENT
     message_bus = config.OVERRIDE.message_bus or (
-        NATSMessageBus() if config.SETTINGS.broker.driver == config.BrokerDriver.NATS else RabbitMQMessageBus()
+        await NATSMessageBus.new(component_type=component_type)
+        if config.SETTINGS.broker.driver == config.BrokerDriver.NATS
+        else await RabbitMQMessageBus.new(component_type=component_type)
     )
     cache = config.OVERRIDE.cache or (
-        NATSCache() if config.SETTINGS.cache.driver == config.CacheDriver.NATS else RedisCache()
+        await NATSCache.new() if config.SETTINGS.cache.driver == config.CacheDriver.NATS else RedisCache()
     )
 
-    service = InfrahubServices(
+    service = await InfrahubServices.new(
         cache=cache,
         client=client,
         database=database,
         workflow=workflow,
         message_bus=message_bus,
-        component_type=ComponentType.GIT_AGENT,
+        component_type=component_type,
     )
-    await service.initialize()
 
     # Initialize the lock
     initialize_lock(service=service)