infrahub-server 1.1.9__py3-none-any.whl → 1.2.0__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

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, AsyncIterator, Optional
+from typing import TYPE_CHECKING, AsyncIterator
 
 from fastapi import Depends, Query, Request
 from fastapi.security import APIKeyHeader, HTTPAuthorizationCredentials, HTTPBearer
@@ -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
@@ -24,7 +25,7 @@ jwt_scheme = HTTPBearer(auto_error=False)
 api_key_scheme = APIKeyHeader(name="X-INFRAHUB-KEY", auto_error=False)
 
 
-async def cookie_auth_scheme(request: Request) -> Optional[str]:
+async def cookie_auth_scheme(request: Request) -> str | None:
     return request.cookies.get("access_token")  # Replace with the actual name of your JWT cookie
 
 
@@ -61,7 +62,7 @@ async def get_access_token(
 async def get_refresh_token(
     request: Request,
     db: InfrahubDatabase = Depends(get_db),
-    jwt_header: Optional[HTTPAuthorizationCredentials] = Depends(jwt_scheme),
+    jwt_header: HTTPAuthorizationCredentials | None = Depends(jwt_scheme),
 ) -> RefreshTokenData:
     token = None
 
@@ -82,8 +83,8 @@ async def get_refresh_token(
 
 async def get_branch_params(
     db: InfrahubDatabase = Depends(get_db),
-    branch_name: Optional[str] = Query(None, alias="branch", description="Name of the branch to use for the query"),
-    at: Optional[str] = Query(None, description="Time to use for the query, in absolute or relative format"),
+    branch_name: str | None = Query(None, alias="branch", description="Name of the branch to use for the query"),
+    at: str | None = Query(None, description="Time to use for the query, in absolute or relative format"),
 ) -> BranchParams:
     branch = await registry.get_branch(db=db, branch=branch_name)
 
@@ -92,7 +93,7 @@ async def get_branch_params(
 
 async def get_branch_dep(
     db: InfrahubDatabase = Depends(get_db),
-    branch_name: Optional[str] = Query(None, alias="branch", description="Name of the branch to use for the query"),
+    branch_name: str | None = Query(None, alias="branch", description="Name of the branch to use for the query"),
 ) -> Branch:
     return await registry.get_branch(db=db, branch=branch_name)
 
@@ -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/diff/validation_models.py

@@ -1,4 +1,4 @@
-from typing import Any, Optional
+from typing import Any
 
 from pydantic import BaseModel, ConfigDict, field_validator, model_validator
 
@@ -9,13 +9,13 @@ from infrahub.core.timestamp import Timestamp
 class DiffQueryValidated(BaseModel):
     model_config = ConfigDict(arbitrary_types_allowed=True)
     branch: Branch
-    time_from: Optional[str] = None
-    time_to: Optional[str] = None
+    time_from: str | None = None
+    time_to: str | None = None
     branch_only: bool
 
     @field_validator("time_from", "time_to", mode="before")
     @classmethod
-    def validate_time(cls, value: Optional[str]) -> Optional[str]:
+    def validate_time(cls, value: str | None) -> str | None:
         if not value:
             return None
         Timestamp(value)
@@ -24,12 +24,12 @@ class DiffQueryValidated(BaseModel):
     @model_validator(mode="before")
     @classmethod
     def validate_time_from_if_required(cls, values: dict[str, Any]) -> dict[str, Any]:
-        branch: Optional[Branch] = values.get("branch")
-        time_from: Optional[Timestamp] = values.get("time_from")
+        branch: Branch | None = values.get("branch")
+        time_from: Timestamp | None = values.get("time_from")
         if getattr(branch, "is_default", False) and not time_from:
             branch_name = getattr(branch, "name", "")
             raise ValueError(f"time_from is mandatory when diffing on the default branch `{branch_name}`.")
-        time_to: Optional[Timestamp] = values.get("time_to")
+        time_to: Timestamp | None = values.get("time_to")
         if time_to and time_from and time_to < time_from:
             raise ValueError("time_from and time_to are not a valid time range")
         return values

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
@@ -67,7 +68,7 @@ class APISchemaMixin:
     @classmethod
     def set_kind(cls, values: Any) -> Any:
         if isinstance(values, dict):
-            values["kind"] = f'{values["namespace"]}{values["name"]}'
+            values["kind"] = f"{values['namespace']}{values['name']}"
         return values
 
 
@@ -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,9 @@
-from typing import Optional
+from pydantic import BaseModel, Field
 
-from pydantic import Field
+from infrahub.context import InfrahubContext
 
-from infrahub.message_bus import InfrahubMessage
 
-
-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 +17,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")
+    artifact_id: str | None = 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}

@@ -1,41 +1,39 @@
-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]] = {
+    artifact_result: dict[str, str | bool | None] = {
         "changed": None,
         "checksum": None,
         "artifact_id": None,
@@ -44,22 +42,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 +67,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 +79,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 +94,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

@@ -1,20 +1,21 @@
 from asyncio import run as aiorun
 
 import typer
+from infrahub_sdk.async_typer import AsyncTyper
 
 from infrahub import config
-from infrahub.cli.context import CliContext
-from infrahub.cli.db import app as db_app
-from infrahub.cli.events import app as events_app
-from infrahub.cli.git_agent import app as git_app
-from infrahub.cli.server import app as server_app
-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
+from .context import CliContext
+from .db import app as db_app
+from .events import app as events_app
+from .git_agent import app as git_app
+from .server import app as server_app
+from .tasks import app as tasks_app
+from .upgrade import upgrade_cmd
 
-app = typer.Typer(name="Infrahub CLI", pretty_exceptions_enable=False)
+app = AsyncTyper(name="Infrahub CLI", pretty_exceptions_enable=False)
 
 
 @app.callback()
@@ -24,10 +25,11 @@ def common(ctx: typer.Context) -> None:
 
 
 app.add_typer(server_app, name="server")
-app.add_typer(git_app, name="git-agent")
+app.add_typer(git_app, name="git-agent", hidden=True)
 app.add_typer(db_app, name="db")
-app.add_typer(events_app, name="events", help="Interact with the events system.")
+app.add_typer(events_app, name="events", help="Interact with the events system.", hidden=True)
 app.add_typer(tasks_app, name="tasks", hidden=True)
+app.command(name="upgrade")(upgrade_cmd)
 
 
 async def _init_shell(config_file: str) -> None:

infrahub/cli/constants.py

@@ -0,0 +1,3 @@
+ERROR_BADGE = "[bold red]ERROR[/bold red]"
+SUCCESS_BADGE = "[green]SUCCESS[/green]"
+FAILED_BADGE = "[bold red]FAILED[/bold red]"