fix(backend): performance issues when generating artifacts

backend/infrahub/artifacts/models.py

@@ -25,7 +25,8 @@ class CheckArtifactCreate(BaseModel):
     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: 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")
+    query: str = Field(..., description="The name of the query to use when collecting data")  # Deprecated
+    query_id: str = Field(..., description="The id 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")

backend/infrahub/computed_attribute/tasks.py

@@ -104,7 +104,7 @@ async def process_transform(
         )  # type: ignore[misc]
 
         data = await client.query_gql_query(
-            name=transform.query.peer.name.value,
+            name=transform.query.id,
             branch_name=branch_name,
             variables={"id": object_id},
             update_group=True,

backend/infrahub/core/schema/manager.py

@@ -2,6 +2,9 @@
 
 from typing import TYPE_CHECKING, Any
 
+from cachetools import LRUCache
+from infrahub_sdk.schema import BranchSchema as SDKBranchSchema
+
 from infrahub import lock
 from infrahub.core.manager import NodeManager
 from infrahub.core.models import (
@@ -40,6 +43,8 @@ class SchemaManager(NodeManager):
     def __init__(self) -> None:
         self._cache: dict[int, Any] = {}
         self._branches: dict[str, SchemaBranch] = {}
+        self._branch_hash_by_name: dict[str, str] = {}
+        self._sdk_branches: LRUCache[str, SDKBranchSchema] = LRUCache(maxsize=10)
 
     def _get_from_cache(self, key: int) -> Any:
         return self._cache[key]
@@ -140,12 +145,23 @@ def get_schema_branch(self, name: str) -> SchemaBranch:
         if name in self._branches:
             return self._branches[name]
 
-        self._branches[name] = SchemaBranch(cache=self._cache, name=name)
+        self.set_schema_branch(name, schema=SchemaBranch(cache=self._cache, name=name))
         return self._branches[name]
 
+    def get_sdk_schema_branch(self, name: str) -> SDKBranchSchema:
+        schema_hash = self._branch_hash_by_name[name]
+        branch_schema = self._sdk_branches.get(schema_hash)
+        if not branch_schema:
+            self._sdk_branches[schema_hash] = SDKBranchSchema.from_api_response(
+                data=self._branches[name].to_dict_api_schema_object()
+            )
+
+        return self._sdk_branches[schema_hash]
+
     def set_schema_branch(self, name: str, schema: SchemaBranch) -> None:
         schema.name = name
         self._branches[name] = schema
+        self._branch_hash_by_name[name] = schema.get_hash()
 
     def process_schema_branch(self, name: str) -> None:
         schema_branch = self.get_schema_branch(name=name)
@@ -764,6 +780,8 @@ def purge_inactive_branches(self, active_branches: list[str]) -> list[str]:
         for branch_name in list(self._branches.keys()):
             if branch_name not in active_branches:
                 del self._branches[branch_name]
+                if branch_name in self._branch_hash_by_name:
+                    del self._branch_hash_by_name[branch_name]
                 removed_branches.append(branch_name)
 
         for hash_key in list(self._cache.keys()):

backend/infrahub/core/schema/schema_branch.py

@@ -162,6 +162,14 @@ def to_dict_schema_object(self, duplicate: bool = False) -> dict[str, dict[str,
             "templates": {name: self.get(name, duplicate=duplicate) for name in self.templates},
         }
 
+    def to_dict_api_schema_object(self) -> dict[str, list[dict]]:
+        return {
+            "nodes": [self.get(name, duplicate=False).model_dump() for name in self.nodes],
+            "profiles": [self.get(name, duplicate=False).model_dump() for name in self.profiles],
+            "generics": [self.get(name, duplicate=False).model_dump() for name in self.generics],
+            "templates": [self.get(name, duplicate=False).model_dump() for name in self.templates],
+        }
+
     @classmethod
     def from_dict_schema_object(cls, data: dict) -> Self:
         type_mapping = {

backend/infrahub/generators/tasks.py

@@ -21,6 +21,7 @@
 )
 from infrahub.git.base import extract_repo_file_information
 from infrahub.git.repository import get_initialized_repo
+from infrahub.git.utils import fetch_proposed_change_generator_definition_targets
 from infrahub.workers.dependencies import get_client, get_workflow
 from infrahub.workflows.catalogue import REQUEST_GENERATOR_DEFINITION_RUN, REQUEST_GENERATOR_RUN
 from infrahub.workflows.utils import add_tags
@@ -177,14 +178,9 @@ async def request_generator_definition_run(
         branch=model.branch,
     )
 
-    group = await client.get(
-        kind=InfrahubKind.GENERICGROUP,
-        prefetch_relationships=True,
-        populate_store=True,
-        id=model.generator_definition.group_id,
-        branch=model.branch,
+    group = await fetch_proposed_change_generator_definition_targets(
+        client=client, branch=model.branch, definition=model.generator_definition
     )
-    await group.members.fetch()
 
     instance_by_member = {}
     for instance in existing_instances:

backend/infrahub/git/integrator.py

@@ -1363,7 +1363,7 @@ async def render_artifact(
         message: CheckArtifactCreate | RequestArtifactGenerate,
     ) -> ArtifactGenerateResult:
         response = await self.sdk.query_gql_query(
-            name=message.query,
+            name=message.query_id,
             variables=message.variables,
             update_group=True,
             subscribers=[artifact.id],

backend/infrahub/git/models.py

@@ -38,7 +38,8 @@ class RequestArtifactGenerate(BaseModel):
     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: 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")
+    query: str = Field(..., description="The name of the query to use when collecting data")  # Deprecated
+    query_id: str = Field(..., description="The id 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")
     context: InfrahubContext = Field(..., description="The context of the task")

backend/infrahub/git/repository.py

@@ -2,6 +2,9 @@
 
 from typing import TYPE_CHECKING, Any
 
+from cachetools import TTLCache
+from cachetools.keys import hashkey
+from cachetools_async import cached
 from git.exc import BadName, GitCommandError
 from infrahub_sdk.exceptions import GraphQLError
 from prefect import task
@@ -248,12 +251,13 @@ async def sync_from_remote(self, commit: str | None = None) -> None:
         await self.update_commit_value(branch_name=self.infrahub_branch_name, commit=commit)
 
 
-@task(
-    name="Fetch repository commit",
-    description="Retrieve a git repository at a given commit, if it does not already exist locally",
-    cache_policy=NONE,
+@cached(
+    TTLCache(maxsize=100, ttl=30),
+    key=lambda *_, **kwargs: hashkey(
+        kwargs.get("repository_id"), kwargs.get("name"), kwargs.get("repository_kind"), kwargs.get("commit")
+    ),
 )
-async def get_initialized_repo(
+async def _get_initialized_repo(
     client: InfrahubClient, repository_id: str, name: str, repository_kind: str, commit: str | None = None
 ) -> InfrahubReadOnlyRepository | InfrahubRepository:
     if repository_kind == InfrahubKind.REPOSITORY:
@@ -263,3 +267,16 @@ async def get_initialized_repo(
         return await InfrahubReadOnlyRepository.init(id=repository_id, name=name, commit=commit, client=client)
 
     raise NotImplementedError(f"The repository kind {repository_kind} has not been implemented")
+
+
+@task(
+    name="Fetch repository commit",
+    description="Retrieve a git repository at a given commit, if it does not already exist locally",
+    cache_policy=NONE,
+)
+async def get_initialized_repo(
+    client: InfrahubClient, repository_id: str, name: str, repository_kind: str, commit: str | None = None
+) -> InfrahubReadOnlyRepository | InfrahubRepository:
+    return await _get_initialized_repo(
+        client=client, repository_id=repository_id, name=name, repository_kind=repository_kind, commit=commit
+    )

backend/infrahub/git/tasks.py

@@ -53,6 +53,7 @@
     UserCheckDefinitionData,
 )
 from .repository import InfrahubReadOnlyRepository, InfrahubRepository, get_initialized_repo
+from .utils import fetch_artifact_definition_targets, fetch_check_definition_targets
 
 
 @flow(
@@ -323,9 +324,8 @@ async def generate_request_artifact_definition(
         kind=CoreArtifactDefinition, id=model.artifact_definition_id, branch=model.branch
     )
 
-    await artifact_definition.targets.fetch()
-    group = artifact_definition.targets.peer
-    await group.members.fetch()
+    group = await fetch_artifact_definition_targets(client=client, branch=model.branch, definition=artifact_definition)
+
     current_members = [member.id for member in group.members.peers]
 
     artifacts_by_member = {}
@@ -356,6 +356,7 @@ async def generate_request_artifact_definition(
         transform_location = f"{transform.file_path.value}::{transform.class_name.value}"
         convert_query_response = transform.convert_query_response.value
 
+    batch = await client.create_batch()
     for relationship in group.members.peers:
         member = relationship.peer
         artifact_id = artifacts_by_member.get(member.id)
@@ -376,6 +377,7 @@ async def generate_request_artifact_definition(
             repository_kind=repository.get_kind(),
             branch_name=model.branch,
             query=query.name.value,
+            query_id=query.id,
             variables=await member.extract(params=artifact_definition.parameters.value),
             target_id=member.id,
             target_name=member.display_label,
@@ -385,10 +387,16 @@ async def generate_request_artifact_definition(
             context=context,
         )
 
-        await get_workflow().submit_workflow(
-            workflow=REQUEST_ARTIFACT_GENERATE, context=context, parameters={"model": request_artifact_generate_model}
+        batch.add(
+            task=get_workflow().submit_workflow,
+            workflow=REQUEST_ARTIFACT_GENERATE,
+            context=context,
+            parameters={"model": request_artifact_generate_model},
         )
 
+    async for _, _ in batch.execute():
+        pass
+
 
 @flow(name="git-repository-pull-read-only", flow_run_name="Pull latest commit on {model.repository_name}")
 async def pull_read_only(model: GitRepositoryPullReadOnly) -> None:
@@ -569,9 +577,7 @@ async def trigger_repository_user_checks_definitions(model: UserCheckDefinitionD
 
     if definition.targets.id:
         # Check against a group of targets
-        await definition.targets.fetch()
-        group = definition.targets.peer
-        await group.members.fetch()
+        group = await fetch_check_definition_targets(client=client, branch=model.branch_name, definition=definition)
         check_models = []
         for relationship in group.members.peers:
             member = relationship.peer

backend/infrahub/git/utils.py

@@ -1,9 +1,16 @@
-from typing import TYPE_CHECKING
+from collections import defaultdict
+from typing import TYPE_CHECKING, Any
+
+from infrahub_sdk import InfrahubClient
+from infrahub_sdk.node import RelationshipManager
+from infrahub_sdk.protocols import CoreArtifactDefinition, CoreCheckDefinition, CoreGroup
+from infrahub_sdk.types import Order
 
 from infrahub.core import registry
 from infrahub.core.constants import InfrahubKind
 from infrahub.core.manager import NodeManager
 from infrahub.database import InfrahubDatabase
+from infrahub.generators.models import ProposedChangeGeneratorDefinition
 
 from .models import RepositoryBranchInfo, RepositoryData
 
@@ -46,3 +53,118 @@ async def get_repositories_commit_per_branch(
             )
 
     return repositories
+
+
+def _collect_parameter_first_segments(params: Any) -> set[str]:
+    segments: set[str] = set()
+
+    def _walk(value: Any) -> None:
+        if isinstance(value, str):
+            segment = value.split("__", 1)[0]
+            if segment:
+                segments.add(segment)
+        elif isinstance(value, dict):
+            for nested in value.values():
+                _walk(nested)
+        elif isinstance(value, (list, tuple, set)):
+            for nested in value:
+                _walk(nested)
+
+    _walk(params)
+    return segments
+
+
+async def _prefetch_group_member_nodes(
+    client: InfrahubClient,
+    members: RelationshipManager,
+    branch: str,
+    required_fields: set[str],
+) -> None:
+    ids_per_kind: dict[str, set[str]] = defaultdict(set)
+    for peer in members.peers:
+        if peer.id and peer.typename:
+            ids_per_kind[peer.typename].add(peer.id)
+
+    if not ids_per_kind:
+        return
+
+    batch = await client.create_batch()
+
+    for kind, ids in ids_per_kind.items():
+        schema = await client.schema.get(kind=kind, branch=branch)
+
+        # FIXME: https://github.com/opsmill/infrahub-sdk-python/pull/205
+        valid_fields = set(schema.attribute_names) | set(schema.relationship_names)
+        keep_relationships = set(schema.relationship_names) & required_fields
+        cleaned_fields = valid_fields - required_fields
+
+        kwargs: dict[str, Any] = {
+            "kind": kind,
+            "ids": list(ids),
+            "branch": branch,
+            "exclude": list(cleaned_fields),
+            "populate_store": True,
+            "order": Order(disable=True),
+        }
+
+        if keep_relationships:
+            kwargs["include"] = list(keep_relationships)
+
+        batch.add(task=client.filters, **kwargs)
+
+    async for _ in batch.execute():
+        pass
+
+
+async def _fetch_definition_targets(
+    client: InfrahubClient,
+    branch: str,
+    group_id: str,
+    parameters: Any,
+) -> CoreGroup:
+    group = await client.get(
+        kind=CoreGroup,
+        id=group_id,
+        branch=branch,
+        include=["members"],
+    )
+
+    parameter_fields = _collect_parameter_first_segments(parameters)
+    await _prefetch_group_member_nodes(
+        client=client,
+        members=group.members,
+        branch=branch,
+        required_fields=parameter_fields,
+    )
+
+    return group
+
+
+async def fetch_artifact_definition_targets(
+    client: InfrahubClient,
+    branch: str,
+    definition: CoreArtifactDefinition,
+) -> CoreGroup:
+    return await _fetch_definition_targets(
+        client=client, branch=branch, group_id=definition.targets.id, parameters=definition.parameters.value
+    )
+
+
+async def fetch_check_definition_targets(
+    client: InfrahubClient,
+    branch: str,
+    definition: CoreCheckDefinition,
+) -> CoreGroup:
+    return await _fetch_definition_targets(
+        client=client, branch=branch, group_id=definition.targets.id, parameters=definition.parameters.value
+    )
+
+
+async def fetch_proposed_change_generator_definition_targets(
+    client: InfrahubClient,
+    branch: str,
+    definition: ProposedChangeGeneratorDefinition,
+) -> CoreGroup:
+    return await _fetch_definition_targets(
+        client=client, branch=branch, group_id=definition.group_id, parameters=definition.parameters
+    )

backend/infrahub/message_bus/types.py

@@ -89,7 +89,8 @@ class ProposedChangeArtifactDefinition(BaseModel):
     definition_id: str
     definition_name: str
     artifact_name: str
-    query_name: str
+    query_name: str  # Deprecated
+    query_id: str
     query_models: list[str]
     repository_id: str
     transform_kind: str