infrahub-server 1.1.2__py3-none-any.whl → 1.1.3__py3-none-any.whl

infrahub/api/__init__.py

@@ -1,11 +1,13 @@
-from typing import NoReturn
+from __future__ import annotations
 
-from fastapi import APIRouter
+from typing import TYPE_CHECKING, NoReturn
+
+from fastapi import APIRouter, Depends
 from fastapi.openapi.docs import (
     get_redoc_html,
     get_swagger_ui_html,
 )
-from starlette.responses import HTMLResponse
+from starlette.responses import HTMLResponse  # noqa: TC002
 
 from infrahub.api import (
     artifact,
@@ -21,8 +23,12 @@ from infrahub.api import (
     storage,
     transformation,
 )
+from infrahub.api.dependencies import get_current_user
 from infrahub.exceptions import ResourceNotFoundError
 
+if TYPE_CHECKING:
+    from infrahub.auth import AccountSession
+
 router = APIRouter(prefix="/api")
 
 router.include_router(artifact.router)
@@ -40,7 +46,9 @@ router.include_router(transformation.router)
 
 
 @router.get("/docs", include_in_schema=False)
-async def custom_swagger_ui_html() -> HTMLResponse:
+async def custom_swagger_ui_html(
+    _: AccountSession = Depends(get_current_user),
+) -> HTMLResponse:
     return get_swagger_ui_html(
         openapi_url="/api/openapi.json",
         title="Infrahub - Swagger UI",
@@ -50,7 +58,7 @@ async def custom_swagger_ui_html() -> HTMLResponse:
 
 
 @router.get("/redoc", include_in_schema=False)
-async def redoc_html() -> HTMLResponse:
+async def redoc_html(_: AccountSession = Depends(get_current_user)) -> HTMLResponse:
     return get_redoc_html(
         openapi_url="/api/openapi.json",
         title="Infrahub - ReDoc",

infrahub/api/artifact.py

@@ -18,6 +18,7 @@ from infrahub.permissions.constants import PermissionDecisionFlag
 from infrahub.workflows.catalogue import REQUEST_ARTIFACT_DEFINITION_GENERATE
 
 if TYPE_CHECKING:
+    from infrahub.auth import AccountSession
     from infrahub.permissions import PermissionManager
 
 log = get_logger()
@@ -37,7 +38,7 @@ async def get_artifact(
     artifact_id: str,
     db: InfrahubDatabase = Depends(get_db),
     branch_params: BranchParams = Depends(get_branch_params),
-    _: str = Depends(get_current_user),
+    _: AccountSession = Depends(get_current_user),
 ) -> Response:
     artifact = await registry.manager.get_one(db=db, id=artifact_id, branch=branch_params.branch, at=branch_params.at)
     if not artifact:

infrahub/api/auth.py

@@ -1,3 +1,7 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
 from fastapi import APIRouter, Depends, Response
 
 from infrahub import config, models
@@ -8,7 +12,9 @@ from infrahub.auth import (
     create_fresh_access_token,
     invalidate_refresh_token,
 )
-from infrahub.database import InfrahubDatabase
+
+if TYPE_CHECKING:
+    from infrahub.database import InfrahubDatabase
 
 router = APIRouter(prefix="/auth")
 

infrahub/api/diff/diff.py

@@ -1,13 +1,12 @@
 from __future__ import annotations
 
 from collections import defaultdict
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING
 
 from fastapi import APIRouter, Depends, Request
 
 from infrahub.api.dependencies import get_branch_dep, get_current_user, get_db
 from infrahub.core import registry
-from infrahub.core.branch import Branch  # noqa: TC001
 from infrahub.core.diff.artifacts.calculator import ArtifactDiffCalculator
 from infrahub.core.diff.branch_differ import BranchDiffer
 from infrahub.core.diff.model.diff import (
@@ -15,9 +14,11 @@ from infrahub.core.diff.model.diff import (
     BranchDiffFile,
     BranchDiffRepository,
 )
-from infrahub.database import InfrahubDatabase  # noqa: TC001
 
 if TYPE_CHECKING:
+    from infrahub.auth import AccountSession
+    from infrahub.core.branch import Branch
+    from infrahub.database import InfrahubDatabase
     from infrahub.services import InfrahubServices
 
 
@@ -29,17 +30,22 @@ async def get_diff_files(
     request: Request,
     db: InfrahubDatabase = Depends(get_db),
     branch: Branch = Depends(get_branch_dep),
-    time_from: Optional[str] = None,
-    time_to: Optional[str] = None,
+    time_from: str | None = None,
+    time_to: str | None = None,
     branch_only: bool = True,
-    _: str = Depends(get_current_user),
+    _: AccountSession = Depends(get_current_user),
 ) -> dict[str, dict[str, BranchDiffRepository]]:
     response: dict[str, dict[str, BranchDiffRepository]] = defaultdict(dict)
     service: InfrahubServices = request.app.state.service
 
     # Query the Diff for all files and repository from the database
     diff = await BranchDiffer.init(
-        db=db, branch=branch, diff_from=time_from, diff_to=time_to, branch_only=branch_only, service=service
+        db=db,
+        branch=branch,
+        diff_from=time_from,
+        diff_to=time_to,
+        branch_only=branch_only,
+        service=service,
     )
     diff_files = await diff.get_files()
 

infrahub/api/file.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Optional, Union
+from typing import TYPE_CHECKING
 
 from fastapi import APIRouter, Depends, Request
 from starlette.responses import PlainTextResponse
@@ -27,13 +27,13 @@ async def get_file(
     file_path: str,
     branch_params: BranchParams = Depends(get_branch_params),
     db: InfrahubDatabase = Depends(get_db),
-    commit: Optional[str] = None,
+    commit: str | None = None,
     _: str = Depends(get_current_user),
 ) -> PlainTextResponse:
     """Retrieve a file from a git repository."""
     service: InfrahubServices = request.app.state.service
 
-    repo: Union[CoreRepository, CoreReadOnlyRepository] = await NodeManager.get_one_by_id_or_default_filter(
+    repo: CoreRepository | CoreReadOnlyRepository = await NodeManager.get_one_by_id_or_default_filter(
         db=db,
         id=repository_id,
         kind=InfrahubKind.GENERICREPOSITORY,

infrahub/api/internal.py

@@ -1,16 +1,27 @@
+from __future__ import annotations
+
 import re
-from typing import Optional
+from typing import TYPE_CHECKING
 
 import ujson
-from fastapi import APIRouter, Request
+from fastapi import APIRouter, Depends, Request
 from lunr.index import Index
 from pydantic import BaseModel
 
 from infrahub import config
-from infrahub.config import AnalyticsSettings, ExperimentalFeaturesSettings, LoggingSettings, MainSettings
+from infrahub.api.dependencies import get_current_user
+from infrahub.config import (  # noqa: TC001
+    AnalyticsSettings,
+    ExperimentalFeaturesSettings,
+    LoggingSettings,
+    MainSettings,
+)
 from infrahub.core import registry
 from infrahub.exceptions import NodeNotFoundError
 
+if TYPE_CHECKING:
+    from infrahub.auth import AccountSession
+
 router = APIRouter()
 
 
@@ -39,7 +50,7 @@ async def get_config() -> ConfigAPI:
 
 
 @router.get("/info")
-async def get_info(request: Request) -> InfoAPI:
+async def get_info(request: Request, _: AccountSession = Depends(get_current_user)) -> InfoAPI:
     return InfoAPI(deployment_id=str(registry.id), version=request.app.version)
 
 
@@ -47,7 +58,7 @@ class SearchDocs:
     def __init__(self) -> None:
         self._title_documents: list[dict] = []
         self._heading_documents: list[dict] = []
-        self._heading_index: Optional[Index] = None
+        self._heading_index: Index | None = None
 
     def _load_json(self) -> None:
         """
@@ -142,7 +153,9 @@ class SearchResultAPI(BaseModel):
 
 
 @router.get("/search/docs", include_in_schema=False)
-async def search_docs(query: str, limit: Optional[int] = None) -> list[SearchResultAPI]:
+async def search_docs(
+    query: str, limit: int | None = None, _: AccountSession = Depends(get_current_user)
+) -> list[SearchResultAPI]:
     smart_query = smart_queries(query)
     search_results = search_docs_loader.heading_index.search(smart_query)
     heading_results = [

infrahub/api/oauth2.py

@@ -3,9 +3,11 @@ from __future__ import annotations
 from typing import TYPE_CHECKING
 from urllib.parse import urljoin
 
+import ujson
 from authlib.integrations.httpx_client import AsyncOAuth2Client
 from fastapi import APIRouter, Depends, Request, Response
 from fastapi.responses import JSONResponse, RedirectResponse
+from opentelemetry import trace
 
 from infrahub import config, models
 from infrahub.api.dependencies import get_db
@@ -20,6 +22,8 @@ if TYPE_CHECKING:
     from infrahub.database import InfrahubDatabase
     from infrahub.services import InfrahubServices
 
+# pylint: disable=R0801
+
 log = get_logger()
 router = APIRouter(prefix="/oauth2")
 
@@ -33,11 +37,16 @@ def _get_redirect_url(request: Request, provider_name: str) -> str:
 @router.get("/{provider_name:str}/authorize")
 async def authorize(request: Request, provider_name: str, final_url: str | None = None) -> Response:
     provider = config.SETTINGS.security.get_oauth2_provider(provider=provider_name)
-    client = AsyncOAuth2Client(
-        client_id=provider.client_id,
-        client_secret=provider.client_secret,
-        scope=provider.scopes,
-    )
+
+    with trace.get_tracer(__name__).start_as_current_span("sso_oauth2_client_configuration") as span:
+        span.set_attribute("provider_name", provider_name)
+        span.set_attribute("scopes", provider.scopes)
+
+        client = AsyncOAuth2Client(
+            client_id=provider.client_id,
+            client_secret=provider.client_secret,
+            scope=provider.scopes,
+        )
 
     redirect_uri = _get_redirect_url(request=request, provider_name=provider_name)
     final_url = final_url or config.SETTINGS.main.public_url or str(request.base_url)
@@ -88,7 +97,10 @@ async def token(
 
     token_response = await service.http.post(provider.token_url, data=token_data)
     _validate_response(response=token_response)
-    payload = token_response.json()
+
+    with trace.get_tracer(__name__).start_as_current_span("sso_token_request") as span:
+        span.set_attribute("token_request_data", ujson.dumps(token_response.json()))
+        payload = token_response.json()
 
     headers = {"Authorization": f"{payload.get('token_type')} {payload.get('access_token')}"}
     if provider.userinfo_method == config.UserInfoMethod.GET:
@@ -102,7 +114,10 @@ async def token(
     if not sso_groups and config.SETTINGS.security.sso_user_default_group:
         sso_groups = [config.SETTINGS.security.sso_user_default_group]
 
-    user_token = await signin_sso_account(db=db, account_name=user_info["name"], sso_groups=sso_groups)
+    with trace.get_tracer(__name__).start_as_current_span("signin_sso_account") as span:
+        span.set_attribute("account_name", ujson.dumps(userinfo_response.json()))
+        span.set_attribute("sso_groups", sso_groups)
+        user_token = await signin_sso_account(db=db, account_name=user_info["name"], sso_groups=sso_groups)
 
     response.set_cookie(
         "access_token", user_token.access_token, httponly=True, max_age=config.SETTINGS.security.access_token_lifetime

infrahub/api/oidc.py

@@ -3,9 +3,11 @@ from __future__ import annotations
 from typing import TYPE_CHECKING
 from urllib.parse import urljoin
 
+import ujson
 from authlib.integrations.httpx_client import AsyncOAuth2Client
 from fastapi import APIRouter, Depends, Request, Response
 from fastapi.responses import JSONResponse, RedirectResponse
+from opentelemetry import trace
 from pydantic import BaseModel, HttpUrl
 
 from infrahub import config, models
@@ -21,6 +23,8 @@ if TYPE_CHECKING:
     from infrahub.database import InfrahubDatabase
     from infrahub.services import InfrahubServices
 
+# pylint: disable=R0801
+
 log = get_logger()
 router = APIRouter(prefix="/oidc")
 
@@ -68,11 +72,16 @@ async def authorize(request: Request, provider_name: str, final_url: str | None
     _validate_response(response=response)
     oidc_config = OIDCDiscoveryConfig(**response.json())
 
-    client = AsyncOAuth2Client(
-        client_id=provider.client_id,
-        client_secret=provider.client_secret,
-        scope=provider.scopes,
-    )
+    with trace.get_tracer(__name__).start_as_current_span("sso_oauth2_client_configuration") as span:
+        span.set_attribute("provider_name", provider_name)
+        span.set_attribute("scopes", provider.scopes)
+        span.set_attribute("discovery_url", provider.discovery_url)
+
+        client = AsyncOAuth2Client(
+            client_id=provider.client_id,
+            client_secret=provider.client_secret,
+            scope=provider.scopes,
+        )
 
     redirect_uri = _get_redirect_url(request=request, provider_name=provider_name)
     final_url = final_url or config.SETTINGS.main.public_url or str(request.base_url)
@@ -126,7 +135,10 @@ async def token(
 
     token_response = await service.http.post(str(oidc_config.token_endpoint), data=token_data)
     _validate_response(response=token_response)
-    payload = token_response.json()
+
+    with trace.get_tracer(__name__).start_as_current_span("sso_token_request") as span:
+        span.set_attribute("token_request_data", ujson.dumps(token_response.json()))
+        payload = token_response.json()
 
     headers = {"Authorization": f"{payload.get('token_type')} {payload.get('access_token')}"}
 
@@ -138,10 +150,14 @@ async def token(
     _validate_response(response=userinfo_response)
     user_info = userinfo_response.json()
     sso_groups = user_info.get("groups", [])
+
     if not sso_groups and config.SETTINGS.security.sso_user_default_group:
         sso_groups = [config.SETTINGS.security.sso_user_default_group]
 
-    user_token = await signin_sso_account(db=db, account_name=user_info["name"], sso_groups=sso_groups)
+    with trace.get_tracer(__name__).start_as_current_span("signin_sso_account") as span:
+        span.set_attribute("account_name", ujson.dumps(userinfo_response.json()))
+        span.set_attribute("sso_groups", sso_groups)
+        user_token = await signin_sso_account(db=db, account_name=user_info["name"], sso_groups=sso_groups)
 
     response.set_cookie(
         "access_token", user_token.access_token, httponly=True, max_age=config.SETTINGS.security.access_token_lifetime

infrahub/api/schema.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Optional, Union
+from typing import TYPE_CHECKING, Any
 
 from fastapi import APIRouter, Depends, Query, Request
 from pydantic import (
@@ -72,17 +72,17 @@ class APISchemaMixin:
 
 
 class APINodeSchema(NodeSchema, APISchemaMixin):
-    api_kind: Optional[str] = Field(default=None, alias="kind", validate_default=True)
+    api_kind: str | None = Field(default=None, alias="kind", validate_default=True)
     hash: str
 
 
 class APIGenericSchema(GenericSchema, APISchemaMixin):
-    api_kind: Optional[str] = Field(default=None, alias="kind", validate_default=True)
+    api_kind: str | None = Field(default=None, alias="kind", validate_default=True)
     hash: str
 
 
 class APIProfileSchema(ProfileSchema, APISchemaMixin):
-    api_kind: Optional[str] = Field(default=None, alias="kind", validate_default=True)
+    api_kind: str | None = Field(default=None, alias="kind", validate_default=True)
     hash: str
 
 
@@ -103,16 +103,16 @@ class SchemasLoadAPI(BaseModel):
 
 
 class JSONSchema(BaseModel):
-    title: Optional[str] = Field(None, description="Title of the schema")
-    description: Optional[str] = Field(None, description="Description of the schema")
+    title: str | None = Field(None, description="Title of the schema")
+    description: str | None = Field(None, description="Description of the schema")
     type: str = Field(..., description="Type of the schema element (e.g., 'object', 'array', 'string')")
-    properties: Optional[dict[str, Any]] = Field(None, description="Properties of the object if type is 'object'")
-    items: Optional[Union[dict[str, Any], list[dict[str, Any]]]] = Field(
+    properties: dict[str, Any] | None = Field(None, description="Properties of the object if type is 'object'")
+    items: dict[str, Any] | list[dict[str, Any]] | None = Field(
         None, description="Items of the array if type is 'array'"
     )
-    required: Optional[list[str]] = Field(None, description="List of required properties if type is 'object'")
-    schema_spec: Optional[str] = Field(None, alias="$schema", description="Schema version identifier")
-    additional_properties: Optional[Union[bool, dict[str, Any]]] = Field(
+    required: list[str] | None = Field(None, description="List of required properties if type is 'object'")
+    schema_spec: str | None = Field(None, alias="$schema", description="Schema version identifier")
+    additional_properties: bool | dict[str, Any] | None = Field(
         None, description="Specifies whether additional properties are allowed", alias="additionalProperties"
     )
 
@@ -152,7 +152,9 @@ def evaluate_candidate_schemas(
 
 @router.get("")
 async def get_schema(
-    branch: Branch = Depends(get_branch_dep), namespaces: Union[list[str], None] = Query(default=None)
+    branch: Branch = Depends(get_branch_dep),
+    namespaces: list[str] | None = Query(default=None),
+    _: AccountSession = Depends(get_current_user),
 ) -> SchemaReadAPI:
     log.debug("schema_request", branch=branch.name)
     schema_branch = registry.schema.get_schema_branch(name=branch.name)
@@ -180,7 +182,9 @@ async def get_schema(
 
 
 @router.get("/summary")
-async def get_schema_summary(branch: Branch = Depends(get_branch_dep)) -> SchemaBranchHash:
+async def get_schema_summary(
+    branch: Branch = Depends(get_branch_dep), _: AccountSession = Depends(get_current_user)
+) -> SchemaBranchHash:
     log.debug("schema_summary_request", branch=branch.name)
     schema_branch = registry.schema.get_schema_branch(name=branch.name)
     return schema_branch.get_hash_full()
@@ -188,13 +192,13 @@ async def get_schema_summary(branch: Branch = Depends(get_branch_dep)) -> Schema
 
 @router.get("/{schema_kind}")
 async def get_schema_by_kind(
-    schema_kind: str, branch: Branch = Depends(get_branch_dep)
-) -> Union[APIProfileSchema, APINodeSchema, APIGenericSchema]:
+    schema_kind: str, branch: Branch = Depends(get_branch_dep), _: AccountSession = Depends(get_current_user)
+) -> APIProfileSchema | APINodeSchema | APIGenericSchema:
     log.debug("schema_kind_request", branch=branch.name)
 
     schema = registry.schema.get(name=schema_kind, branch=branch, duplicate=False)
 
-    api_schema: dict[str, type[Union[APIProfileSchema, APINodeSchema, APIGenericSchema]]] = {
+    api_schema: dict[str, type[APIProfileSchema | APINodeSchema | APIGenericSchema]] = {
         "profile": APIProfileSchema,
         "node": APINodeSchema,
         "generic": APIGenericSchema,
@@ -212,7 +216,9 @@ async def get_schema_by_kind(
 
 
 @router.get("/json_schema/{schema_kind}")
-async def get_json_schema_by_kind(schema_kind: str, branch: Branch = Depends(get_branch_dep)) -> JSONSchema:
+async def get_json_schema_by_kind(
+    schema_kind: str, branch: Branch = Depends(get_branch_dep), _: AccountSession = Depends(get_current_user)
+) -> JSONSchema:
     log.debug("json_schema_kind_request", branch=branch.name)
 
     fields: dict[str, Any] = {}
@@ -368,7 +374,7 @@ async def check_schema(
     request: Request,
     schemas: SchemasLoadAPI,
     branch: Branch = Depends(get_branch_dep),
-    _: Any = Depends(get_current_user),
+    _: AccountSession = Depends(get_current_user),
 ) -> JSONResponse:
     service: InfrahubServices = request.app.state.service
     log.info("schema_check_request", branch=branch.name)

infrahub/api/storage.py

@@ -1,4 +1,7 @@
+from __future__ import annotations
+
 import hashlib
+from typing import TYPE_CHECKING
 
 from fastapi import APIRouter, Depends, File, Response, UploadFile
 from infrahub_sdk.uuidt import UUIDT
@@ -8,6 +11,9 @@ from infrahub.api.dependencies import get_current_user
 from infrahub.core import registry
 from infrahub.log import get_logger
 
+if TYPE_CHECKING:
+    from infrahub.auth import AccountSession
+
 log = get_logger()
 router = APIRouter(prefix="/storage")
 
@@ -22,10 +28,7 @@ class UploadContentPayload(BaseModel):
 
 
 @router.get("/object/{identifier:str}")
-def get_file(
-    identifier: str,
-    _: str = Depends(get_current_user),
-) -> Response:
+def get_file(identifier: str, _: AccountSession = Depends(get_current_user)) -> Response:
     content = registry.storage.retrieve(identifier=identifier)
     return Response(content=content)
 
@@ -48,10 +51,7 @@ def upload_content(
 
 
 @router.post("/upload/file")
-def upload_file(
-    file: UploadFile = File(...),
-    _: str = Depends(get_current_user),
-) -> UploadResponse:
+def upload_file(file: UploadFile = File(...), _: AccountSession = Depends(get_current_user)) -> UploadResponse:
     # TODO need to optimized how we read the content of the file, especially if the file is really large
     # Check this discussion for more details
     # https://stackoverflow.com/questions/63048825/how-to-upload-file-using-fastapi

infrahub/api/transformation.py

@@ -27,6 +27,7 @@ from infrahub.transformations.models import TransformJinjaTemplateData, Transfor
 from infrahub.workflows.catalogue import TRANSFORM_JINJA2_RENDER, TRANSFORM_PYTHON_RENDER
 
 if TYPE_CHECKING:
+    from infrahub.auth import AccountSession
     from infrahub.services import InfrahubServices
 
 router = APIRouter()
@@ -38,7 +39,7 @@ async def transform_python(
     transform_id: str,
     db: InfrahubDatabase = Depends(get_db),
     branch_params: BranchParams = Depends(get_branch_params),
-    _: str = Depends(get_current_user),
+    _: AccountSession = Depends(get_current_user),
 ) -> JSONResponse:
     params = {key: value for key, value in request.query_params.items() if key not in ["branch", "at"]}
 
@@ -97,7 +98,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),
-    _: str = Depends(get_current_user),
+    _: AccountSession = Depends(get_current_user),
 ) -> PlainTextResponse:
     params = {key: value for key, value in request.query_params.items() if key not in ["branch", "at"]}
 

infrahub/auth.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 import uuid
 from datetime import datetime, timedelta, timezone
 from enum import Enum
-from typing import TYPE_CHECKING, Callable
+from typing import TYPE_CHECKING
 
 import bcrypt
 import jwt
@@ -233,29 +233,6 @@ async def validate_api_key(db: InfrahubDatabase, token: str) -> AccountSession:
     return AccountSession(account_id=account_id, auth_type=AuthType.API)
 
 
-def _validate_update_account(account_session: AccountSession, node_id: str, fields: list[str]) -> None:
-    if account_session.account_id != node_id:
-        # A regular account is not allowed to modify another account
-        raise PermissionError("You are not allowed to modify this account")
-
-    allowed_fields = ["description", "label", "password"]
-    for field in fields:
-        if field not in allowed_fields:
-            raise PermissionError(f"You are not allowed to modify '{field}'")
-
-
-def validate_mutation_permissions_update_node(
-    operation: str, node_id: str, account_session: AccountSession, fields: list[str]
-) -> None:
-    validation_map: dict[str, Callable[[AccountSession, str, list[str]], None]] = {
-        f"{InfrahubKind.ACCOUNT}Update": _validate_update_account,
-        f"{InfrahubKind.ACCOUNT}Upsert": _validate_update_account,
-    }
-
-    if validator := validation_map.get(operation):
-        validator(account_session, node_id, fields)
-
-
 async def invalidate_refresh_token(db: InfrahubDatabase, token_id: str) -> None:
     refresh_token = await NodeManager.get_one(id=token_id, db=db)
     if refresh_token:

infrahub/cli/__init__.py

@@ -20,7 +20,7 @@ app = typer.Typer(name="Infrahub CLI", pretty_exceptions_enable=False)
 @app.callback()
 def common(ctx: typer.Context) -> None:
     """Infrahub CLI"""
-    ctx.obj = CliContext(database_class=InfrahubDatabase)
+    ctx.obj = CliContext()
 
 
 app.add_typer(server_app, name="server")

infrahub/cli/context.py

@@ -1,18 +1,15 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import TYPE_CHECKING
 
-from infrahub.database import get_db
-
-if TYPE_CHECKING:
-    from infrahub.database import InfrahubDatabase
+from infrahub.database import InfrahubDatabase, get_db
 
 
 @dataclass
 class CliContext:
-    database_class: type[InfrahubDatabase]
     application: str = "infrahub.server:app"
 
-    async def get_db(self, retry: int = 0) -> InfrahubDatabase:
-        return self.database_class(driver=await get_db(retry=retry))
+    # This method is inherited for Infrahub Enterprise.
+    @staticmethod
+    async def init_db(retry: int) -> InfrahubDatabase:
+        return InfrahubDatabase(driver=await get_db(retry=retry))

infrahub/cli/db.py

@@ -97,7 +97,7 @@ async def init(
     config.load_and_exit(config_file_name=config_file)
 
     context: CliContext = ctx.obj
-    dbdriver = await context.get_db(retry=1)
+    dbdriver = await context.init_db(retry=1)
     async with dbdriver.start_transaction() as db:
         log.info("Delete All Nodes")
         await delete_all_nodes(db=db)
@@ -120,7 +120,7 @@ async def load_test_data(
     config.load_and_exit(config_file_name=config_file)
 
     context: CliContext = ctx.obj
-    dbdriver = await context.get_db(retry=1)
+    dbdriver = await context.init_db(retry=1)
     async with dbdriver.start_session() as db:
         await initialization(db=db)
 
@@ -152,7 +152,7 @@ async def migrate(
     config.load_and_exit(config_file_name=config_file)
 
     context: CliContext = ctx.obj
-    dbdriver = await context.get_db(retry=1)
+    dbdriver = await context.init_db(retry=1)
     async with dbdriver.start_session() as db:
         rprint("Checking current state of the Database")
 
@@ -207,7 +207,7 @@ async def update_core_schema(  # pylint: disable=too-many-statements
     config.load_and_exit(config_file_name=config_file)
 
     context: CliContext = ctx.obj
-    dbdriver = await context.get_db(retry=1)
+    dbdriver = await context.init_db(retry=1)
 
     error_badge = "[bold red]ERROR[/bold red]"
 
@@ -332,7 +332,7 @@ async def constraint(
     config.load_and_exit(config_file_name=config_file)
 
     context: CliContext = ctx.obj
-    dbdriver = await context.get_db(retry=1)
+    dbdriver = await context.init_db(retry=1)
 
     manager: Optional[ConstraintManagerBase] = None
     if dbdriver.db_type == DatabaseType.NEO4J:
@@ -376,7 +376,7 @@ async def index(
     config.load_and_exit(config_file_name=config_file)
 
     context: CliContext = ctx.obj
-    dbdriver = await context.get_db(retry=1)
+    dbdriver = await context.init_db(retry=1)
     dbdriver.manager.index.init(nodes=node_indexes, rels=rel_indexes)
 
     if action == IndexAction.ADD: